Improved get-or-create-variable logic

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/Domain.kt

@@ -12,7 +12,6 @@ import org.utbot.framework.plugin.api.CodeGenerationSettingItem
 import org.utbot.framework.plugin.api.MethodId
 import org.utbot.framework.plugin.api.TypeParameters
 import org.utbot.framework.plugin.api.UtModel
-import org.utbot.framework.plugin.api.idOrNull
 import org.utbot.framework.plugin.api.isolateCommandLineArgumentsToArgumentFile
 import org.utbot.framework.plugin.api.util.booleanArrayClassId
 import org.utbot.framework.plugin.api.util.booleanClassId
@@ -730,18 +729,13 @@ object SpringBoot : DependencyInjectionFramework(
 )
 
 /**
- * Extended id of [UtModel], unique for whole test set.
+ * Extended [UtModel] model with testSet id and execution id.
  *
- * Allows distinguishing models from different executions and test sets,
- * even if they have the same value of `UtModel.id` that is allowed.
+ * Used as a key in [valueByUtModelWrapper].
+ * Was introduced primarily for shared among all test methods global variables.
  */
-data class ModelId private constructor(
-    private val id: Int?,
-    private val executionId: Int,
-    private val testSetId: Int,
-) {
-    companion object {
-        fun create(model: UtModel, executionId: Int = -1, testSetId: Int = -1) = ModelId(model.idOrNull(), executionId, testSetId)
-    }
-}
-
+data class UtModelWrapper(
+    val testSetId: Int,
+    val executionId: Int,
+    val model: UtModel
+)

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/context/CgContext.kt

@@ -15,14 +15,13 @@ import org.utbot.framework.codegen.domain.models.CgTestMethod
 import org.utbot.framework.codegen.domain.models.CgThisInstance
 import org.utbot.framework.codegen.domain.models.CgValue
 import org.utbot.framework.codegen.domain.models.CgVariable
-import java.util.IdentityHashMap
 import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.PersistentSet
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.persistentMapOf
 import kotlinx.collections.immutable.persistentSetOf
-import org.utbot.framework.codegen.domain.ModelId
+import org.utbot.framework.codegen.domain.UtModelWrapper
 import org.utbot.framework.codegen.domain.ProjectType
 import org.utbot.framework.codegen.domain.models.CgMethodTestSet
 import org.utbot.framework.codegen.domain.builtin.TestClassUtilMethodProvider
@@ -195,10 +194,30 @@ interface CgContextOwner {
 
     val shouldOptimizeImports: Boolean
 
-    var valueByModel: IdentityHashMap<UtModel, CgValue>
+    /**
+     * Used for differentiating models in codegen
+     * because comparisons by [UtModel] are not enough,
+     * especially for Spring related variable processing.
+     *
+     * Default value is -1, meaning that we are not processing any test set.
+     *
+     * @see [CgContext.withTestSetIdScope]
+     */
+    var currentTestSetId: Int
+
+    /**
+     * Used for differentiating models in codegen
+     * because comparisons by [UtModel] are not enough,
+     * especially for Spring related variable processing.
+     *
+     * Default value is -1, meaning that we are not processing any execution.
+     *
+     * @see [CgContext.withExecutionIdScope]
+     */
+    var currentExecutionId: Int
 
-    // use it to compare stateBefore and result variables - in case of equality do not create new variable
-    var valueByModelId: MutableMap<ModelId, CgValue>
+    // used for comparing stateBefore and result variables -- in case of equality do not create new variable
+    var valueByUtModelWrapper: MutableMap<UtModelWrapper, CgValue>
 
     // parameters of the method currently being generated
     val currentMethodParameters: MutableMap<CgParameterKind, CgVariable>
@@ -227,12 +246,6 @@ interface CgContextOwner {
      */
     var successfulExecutionsModels: List<UtModel>
 
-    /**
-     * Gives a unique identifier to model in test set.
-     * Determines which execution current model belongs to.
-     */
-    var modelIds: MutableMap<UtModel, ModelId>
-
     fun block(init: () -> Unit): Block {
         val prevBlock = currentBlock
         return try {
@@ -260,6 +273,16 @@ interface CgContextOwner {
         currentExecutable = executableId
     }
 
+    fun <R> withTestSetIdScope(testSetId: Int, block: () -> R): R {
+        currentTestSetId = testSetId
+        return block()
+    }
+
+    fun <R> withExecutionIdScope(executionId: Int, block: () -> R): R {
+        currentExecutionId = executionId
+        return block()
+    }
+
     fun addExceptionIfNeeded(exception: ClassId) {
         if (exception !is BuiltinClassId) {
             require(exception isSubtypeOf Throwable::class.id) {
@@ -316,11 +339,10 @@ interface CgContextOwner {
 
     fun updateVariableScope(variable: CgVariable, model: UtModel? = null) {
         model?.let {
-            valueByModel[it] = variable
+            valueByUtModelWrapper[it.wrap()] = variable
             (model as UtReferenceModel).let { refModel ->
                 refModel.id.let {
-                    val modelId = getIdByModel(model)
-                    valueByModelId[modelId] = variable
+                    valueByUtModelWrapper[model.wrap()] = variable
                 }
             }
         }
@@ -331,17 +353,15 @@ interface CgContextOwner {
         val prevDeclaredClassRefs = declaredClassRefs
         val prevDeclaredExecutableRefs = declaredExecutableRefs
         val prevDeclaredFieldRefs = declaredFieldRefs
-        val prevValueByModel = IdentityHashMap(valueByModel)
-        val prevValueByModelId = valueByModelId.toMutableMap()
+        val prevValueByCgModel = valueByUtModelWrapper.toMutableMap()
         return try {
             block()
         } finally {
             existingVariableNames = prevVariableNames
             declaredClassRefs = prevDeclaredClassRefs
             declaredExecutableRefs = prevDeclaredExecutableRefs
             declaredFieldRefs = prevDeclaredFieldRefs
-            valueByModel = prevValueByModel
-            valueByModelId = prevValueByModelId
+            valueByUtModelWrapper = prevValueByCgModel
         }
     }
 
@@ -439,7 +459,7 @@ interface CgContextOwner {
     val getLambdaMethod: MethodId
         get() = utilMethodProvider.getLambdaMethodMethodId
 
-    fun getIdByModel(model: UtModel): ModelId
+    fun UtModel.wrap(): UtModelWrapper
 }
 
 /**
@@ -491,8 +511,6 @@ class CgContext(
     override lateinit var actual: CgVariable
     override lateinit var successfulExecutionsModels: List<UtModel>
 
-    override var modelIds: MutableMap<UtModel, ModelId> = mutableMapOf()
-
     /**
      * This property cannot be accessed outside of test class file scope
      * (i.e. outside of [CgContextOwner.withTestClassFileScope]).
@@ -570,7 +588,12 @@ class CgContext(
         }
     }
 
-    override fun getIdByModel(model: UtModel): ModelId = modelIds.getOrPut(model) { ModelId.create(model) }
+    override fun UtModel.wrap(): UtModelWrapper =
+        UtModelWrapper(
+            testSetId = currentTestSetId,
+            executionId = currentExecutionId,
+            model = this
+        )
 
     private fun createClassIdForNestedClass(testClassModel: SimpleTestClassModel): ClassId {
         val simpleName = "${testClassModel.classUnderTest.simpleName}Test"
@@ -588,15 +611,15 @@ class CgContext(
         importedClasses.clear()
         testMethods.clear()
         requiredUtilMethods.clear()
-        valueByModel.clear()
-        valueByModelId.clear()
-        modelIds.clear()
+        valueByUtModelWrapper.clear()
         mockFrameworkUsed = false
     }
 
-    override var valueByModel: IdentityHashMap<UtModel, CgValue> = IdentityHashMap()
+    override var currentTestSetId: Int = -1
+
+    override var currentExecutionId: Int = -1
 
-    override var valueByModelId: MutableMap<ModelId, CgValue> = mutableMapOf()
+    override var valueByUtModelWrapper: MutableMap<UtModelWrapper, CgValue> = mutableMapOf()
 
     override val currentMethodParameters: MutableMap<CgParameterKind, CgVariable> = mutableMapOf()
 

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/models/TestClassModel.kt

@@ -1,9 +1,7 @@
 package org.utbot.framework.codegen.domain.models
 
+import org.utbot.framework.codegen.domain.UtModelWrapper
 import org.utbot.framework.plugin.api.ClassId
-import org.utbot.framework.plugin.api.UtModel
-
-typealias ClassModels = Map<ClassId, Set<UtModel>>
 
 /**
  * Stores method test sets in a structure that replicates structure of their methods in [classUnderTest].
@@ -32,7 +30,7 @@ class SpringTestClassModel(
     classUnderTest: ClassId,
     methodTestSets: List<CgMethodTestSet>,
     nestedClasses: List<SimpleTestClassModel>,
-    val injectedMockModels: ClassModels = mapOf(),
-    val mockedModels: ClassModels = mapOf(),
+    val injectedMockModels: Map<ClassId, Set<UtModelWrapper>> = mapOf(),
+    val mockedModels: Map<ClassId, Set<UtModelWrapper>> = mapOf(),
 ): TestClassModel(classUnderTest, methodTestSets, nestedClasses)
 

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/domain/models/builders/SpringTestClassModelBuilder.kt

@@ -1,9 +1,8 @@
 package org.utbot.framework.codegen.domain.models.builders
 
-import org.utbot.framework.codegen.domain.ModelId
+import org.utbot.framework.codegen.domain.UtModelWrapper
 import org.utbot.framework.codegen.domain.context.CgContext
 import org.utbot.framework.codegen.domain.models.CgMethodTestSet
-import org.utbot.framework.codegen.domain.models.ClassModels
 import org.utbot.framework.codegen.domain.models.SpringTestClassModel
 import org.utbot.framework.plugin.api.ClassId
 import org.utbot.framework.plugin.api.UtArrayModel
@@ -27,61 +26,64 @@ class SpringTestClassModelBuilder(val context: CgContext): TestClassModelBuilder
         val (injectedModels, mockedModels) = collectInjectedAndMockedModels(testSets)
 
         return SpringTestClassModel(
-            baseModel.classUnderTest,
-            baseModel.methodTestSets,
-            baseModel.nestedClasses,
-            injectedModels,
-            mockedModels,
+            classUnderTest = baseModel.classUnderTest,
+            methodTestSets = baseModel.methodTestSets,
+            nestedClasses = baseModel.nestedClasses,
+            injectedMockModels = injectedModels,
+            mockedModels = mockedModels
         )
     }
 
-    private fun collectInjectedAndMockedModels(testSets: List<CgMethodTestSet>): Pair<ClassModels, ClassModels> {
-        val thisInstances = mutableSetOf<UtModel>()
-        val thisInstancesDependentModels = mutableSetOf<UtModel>()
-
-        for ((testSetIndex, testSet) in testSets.withIndex()) {
-            for ((executionIndex, execution) in testSet.executions.withIndex()) {
-
-                setOf(execution.stateBefore.thisInstance, execution.stateAfter.thisInstance)
-                    .filterNotNull()
-                    .forEach { model ->
-                        thisInstances += model
-                        thisInstancesDependentModels += collectByThisInstanceModel(model, executionIndex, testSetIndex)
+    private fun collectInjectedAndMockedModels(testSets: List<CgMethodTestSet>): Pair<Map<ClassId, Set<UtModelWrapper>>, Map<ClassId, Set<UtModelWrapper>>> {
+        val thisInstances = mutableSetOf<UtModelWrapper>()
+        val thisInstancesDependentModels = mutableSetOf<UtModelWrapper>()
+
+        with(context) {
+            for ((testSetIndex, testSet) in testSets.withIndex()) {
+                withTestSetIdScope(testSetIndex) {
+                    for ((executionIndex, execution) in testSet.executions.withIndex()) {
+                        withExecutionIdScope(executionIndex) {
+                            setOf(execution.stateBefore.thisInstance, execution.stateAfter.thisInstance)
+                                .filterNotNull()
+                                .forEach { model ->
+                                    thisInstances += model.wrap()
+                                    thisInstancesDependentModels += collectByThisInstanceModel(model)
+                                }
+                        }
                     }
+                }
             }
         }
 
         val dependentMockModels =
-            thisInstancesDependentModels.filterTo(mutableSetOf()) { it.isMockModel() && it !in thisInstances }
+            thisInstancesDependentModels
+                .filterTo(mutableSetOf()) { cgModel ->
+                    cgModel.model.isMockModel() && cgModel !in thisInstances
+                }
 
         return thisInstances.groupByClassId() to dependentMockModels.groupByClassId()
     }
 
-    private fun collectByThisInstanceModel(model: UtModel, executionIndex: Int, testSetIndex: Int): Set<UtModel> {
-        context.modelIds[model] = ModelId.create(model, executionIndex, testSetIndex)
-
-        val dependentModels = mutableSetOf<UtModel>()
+    private fun collectByThisInstanceModel(model: UtModel): Set<UtModelWrapper> {
+        val dependentModels = mutableSetOf<UtModelWrapper>()
         collectRecursively(model, dependentModels)
 
-        dependentModels.forEach { model ->
-            context.modelIds[model] = ModelId.create(model, executionIndex, testSetIndex)
-        }
-
         return dependentModels
     }
 
-    private fun Set<UtModel>.groupByClassId(): ClassModels {
-        val classModels = mutableMapOf<ClassId, Set<UtModel>>()
+    private fun Set<UtModelWrapper>.groupByClassId(): Map<ClassId, Set<UtModelWrapper>> {
+        val classModels = mutableMapOf<ClassId, Set<UtModelWrapper>>()
 
-        for (modelGroup in this.groupBy { it.classId }) {
+        for (modelGroup in this.groupBy { it.model.classId }) {
             classModels[modelGroup.key] = modelGroup.value.toSet()
         }
 
         return classModels
     }
 
-    private fun collectRecursively(currentModel: UtModel, allModels: MutableSet<UtModel>) {
-        if (!allModels.add(currentModel)) {
+    private fun collectRecursively(currentModel: UtModel, allModels: MutableSet<UtModelWrapper>) {
+        val cgModel = with(context) { currentModel.wrap() }
+        if (!allModels.add(cgModel)) {
             return
         }
 

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgAbstractTestClassConstructor.kt

@@ -98,7 +98,9 @@ abstract class CgAbstractTestClassConstructor<T : TestClassModel>(val context: C
             }
 
             for (i in checkedRange) {
-                currentTestCaseTestMethods += methodConstructor.createTestMethod(methodUnderTest, testSet.executions[i])
+                withExecutionIdScope(i) {
+                    currentTestCaseTestMethods += methodConstructor.createTestMethod(methodUnderTest, testSet.executions[i])
+                }
             }
 
             val comments = listOf("Actual number of generated tests (${executionIndices.last - executionIndices.first}) exceeds per-method limit (${UtSettings.maxTestsPerMethodInRegion})",

utbot-framework/src/main/kotlin/org/utbot/framework/codegen/tree/CgSimpleTestClassConstructor.kt

@@ -49,14 +49,16 @@ open class CgSimpleTestClassConstructor(context: CgContext): CgAbstractTestClass
                 }
             }
 
-            for (testSet in notYetConstructedTestSets) {
+            for ((testSetIndex, testSet) in notYetConstructedTestSets.withIndex()) {
                 updateCurrentExecutable(testSet.executableId)
-                val currentMethodUnderTestRegions = constructTestSet(testSet) ?: continue
-                val executableUnderTestCluster = CgMethodsCluster(
-                    "Test suites for executable $currentExecutable",
-                    currentMethodUnderTestRegions
-                )
-                methodRegions += executableUnderTestCluster
+                withTestSetIdScope(testSetIndex) {
+                    val currentMethodUnderTestRegions = constructTestSet(testSet) ?: return@withTestSetIdScope
+                    val executableUnderTestCluster = CgMethodsCluster(
+                        "Test suites for executable $currentExecutable",
+                        currentMethodUnderTestRegions
+                    )
+                    methodRegions += executableUnderTestCluster
+                }
             }
 
             val currentTestClassDataProviderMethods = currentTestClassContext.cgDataProviderMethods
@@ -159,8 +161,11 @@ open class CgSimpleTestClassConstructor(context: CgContext): CgAbstractTestClass
 
         testSet.executions
             .filterIsInstance<UtFuzzedExecution>()
-            .forEach { execution ->
-                testMethods += methodConstructor.createTestMethod(testSet.executableId, execution)
+            .withIndex()
+            .forEach { (index, execution) ->
+                withExecutionIdScope(index) {
+                    testMethods += methodConstructor.createTestMethod(testSet.executableId, execution)
+                }
             }
 
         return testMethods
@@ -176,8 +181,11 @@ open class CgSimpleTestClassConstructor(context: CgContext): CgAbstractTestClass
         testSet.executions
             .filterIsInstance<UtSymbolicExecution>()
             .filter { it.containsMocking }
-            .forEach { execution ->
-                testMethods += methodConstructor.createTestMethod(testSet.executableId, execution)
+            .withIndex()
+            .forEach { (index, execution) ->
+                withExecutionIdScope(index) {
+                    testMethods += methodConstructor.createTestMethod(testSet.executableId, execution)
+                }
             }
 
         return testMethods