Types prioritization during test generation (#128)

UnitTestBot · Nov 22, 2023 · 14f383a · 14f383a
1 parent 3a8cd7c
commit 14f383a
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 26 deletions.
diff --git a/usvm-jvm/src/main/kotlin/org/usvm/api/util/JcTestInterpreter.kt b/usvm-jvm/src/main/kotlin/org/usvm/api/util/JcTestInterpreter.kt
@@ -51,10 +51,11 @@ import org.usvm.memory.URegisterStackLValue
 import org.usvm.collection.array.UArrayIndexLValue
 import org.usvm.collection.array.length.UArrayLengthLValue
 import org.usvm.collection.field.UFieldLValue
+import org.usvm.machine.interpreter.JcFixedInheritorsNumberTypeSelector
+import org.usvm.machine.interpreter.JcTypeStreamPrioritization
 import org.usvm.model.UModelBase
 import org.usvm.sizeSort
 import org.usvm.types.first
-import org.usvm.types.firstOrNull
 
 /**
  * A class, responsible for resolving a single [JcTest] for a specific method from a symbolic state.
@@ -125,6 +126,9 @@ class JcTestInterpreter(
         private val classLoader: ClassLoader = JcClassLoader,
     ) {
         private val resolvedCache = mutableMapOf<UConcreteHeapAddress, Any?>()
+        private val typeSelector = JcTypeStreamPrioritization(
+            typesToScore = JcFixedInheritorsNumberTypeSelector.DEFAULT_INHERITORS_NUMBER_TO_SCORE
+        )
 
         fun resolveState(): JcParametersState {
             // TODO: now we need to explicitly evaluate indices of registers, because we don't have specific ULValues
@@ -192,7 +196,7 @@ class JcTestInterpreter(
             // and array types right now.
             // In such cases, we need to resolve this element to null.
 
-            val evaluatedType = typeStream.firstOrNull() ?: return null
+            val evaluatedType = typeSelector.firstOrNull(typeStream, type.jcClass) ?: return null
 
             // We check for the type stream emptiness firsly and only then for the resolved cache,
             // because even if the object is already resolved, it could be incompatible with the [type], if it

diff --git a/usvm-jvm/src/main/kotlin/org/usvm/machine/interpreter/JcFixedInheritorsNumberTypeSelector.kt b/usvm-jvm/src/main/kotlin/org/usvm/machine/interpreter/JcFixedInheritorsNumberTypeSelector.kt
@@ -1,9 +1,11 @@
 package org.usvm.machine.interpreter
 
 import org.jacodb.api.JcArrayType
+import org.jacodb.api.JcClassOrInterface
 import org.jacodb.api.JcClassType
 import org.jacodb.api.JcMethod
 import org.jacodb.api.JcType
+import org.jacodb.api.ext.constructors
 import org.usvm.machine.logger
 import org.usvm.types.TypesResult
 import org.usvm.types.UTypeStream
@@ -14,35 +16,63 @@ interface JcTypeSelector {
 
 class JcFixedInheritorsNumberTypeSelector(
     private val inheritorsNumberToChoose: Int = DEFAULT_INHERITORS_NUMBER_TO_CHOOSE,
-    private val inheritorsNumberToSelectFrom: Int = DEFAULT_INHERITORS_NUMBER_TO_SCORE,
+    inheritorsNumberToSelectFrom: Int = DEFAULT_INHERITORS_NUMBER_TO_SCORE,
 ) : JcTypeSelector {
+    private val typesPriorities = JcTypeStreamPrioritization(inheritorsNumberToSelectFrom)
 
-    override fun choose(method: JcMethod, typeStream: UTypeStream<out JcType>): Collection<JcType> {
-        return typeStream
-            .take(inheritorsNumberToSelectFrom)
+    override fun choose(method: JcMethod, typeStream: UTypeStream<out JcType>): Collection<JcType> =
+        typesPriorities.take(typeStream, method.enclosingClass, inheritorsNumberToChoose)
+            .also {
+                logger.debug { "Select types for ${method.enclosingClass.name} : ${it.map { it.typeName }}" }
+            }
+
+    companion object {
+        const val DEFAULT_INHERITORS_NUMBER_TO_CHOOSE: Int = 4
+        // TODO: elaborate on better constant choosing
+        const val DEFAULT_INHERITORS_NUMBER_TO_SCORE: Int = 100
+    }
+}
+
+class JcTypeStreamPrioritization(private val typesToScore: Int) {
+    fun take(
+        typeStream: UTypeStream<out JcType>,
+        referenceClass: JcClassOrInterface,
+        limit: Int
+    ): Collection<JcType> = fetchTypes(typeStream)
+        .sortedByDescending { type -> typeScore(referenceClass, type) }
+        .take(limit)
+
+    fun firstOrNull(
+        typeStream: UTypeStream<out JcType>,
+        referenceClass: JcClassOrInterface,
+    ): JcType? = fetchTypes(typeStream)
+        .maxByOrNull { type -> typeScore(referenceClass, type) }
+
+    private fun fetchTypes(typeStream: UTypeStream<out JcType>): Collection<JcType> =
+        typeStream
+            .take(typesToScore)
             .let {
                 when (it) {
                     TypesResult.EmptyTypesResult -> emptyList()
                     is TypesResult.SuccessfulTypesResult -> it
                     is TypesResult.TypesResultWithExpiredTimeout -> it.collectedTypes
                 }
             }
-            .sortedByDescending { type -> typeScore(method, type) }
-            .take(inheritorsNumberToChoose)
-            .also {
-                logger.debug { "Select types for ${method.enclosingClass.name} : ${it.map { it.typeName }}" }
-            }
-    }
 
-    private fun typeScore(method: JcMethod, type: JcType): Double {
+    private fun typeScore(referenceClass: JcClassOrInterface, type: JcType): Double {
         var score = 0.0
 
         if (type is JcClassType) {
             // prefer class types over arrays
             score += 1
 
             if (type.isPublic) {
-                score += 1
+                score += 3
+            }
+
+            // Prefer easy instantiable classes
+            if (type.constructors.any { it.isPublic }) {
+                score += 3
             }
 
             if (type.isFinal) {
@@ -54,7 +84,7 @@ class JcFixedInheritorsNumberTypeSelector(
             }
 
             val typePkg = type.jcClass.name.split(".")
-            val methodPkg = method.enclosingClass.name.split(".")
+            val methodPkg = referenceClass.name.split(".")
 
             for ((typePkgPart, methodPkgPart) in typePkg.zip(methodPkg)) {
                 if (typePkgPart != methodPkgPart) break
@@ -63,16 +93,10 @@ class JcFixedInheritorsNumberTypeSelector(
         }
 
         if (type is JcArrayType) {
-            val elementScore = typeScore(method, type.elementType)
+            val elementScore = typeScore(referenceClass, type.elementType)
             score += elementScore / 10
         }
 
         return score
     }
-
-    companion object {
-        const val DEFAULT_INHERITORS_NUMBER_TO_CHOOSE: Int = 4
-        // TODO: elaborate on better constant choosing
-        const val DEFAULT_INHERITORS_NUMBER_TO_SCORE: Int = 100
-    }
-}
+}
diff --git a/usvm-jvm/src/test/kotlin/org/usvm/util/JcTestExecutor.kt b/usvm-jvm/src/test/kotlin/org/usvm/util/JcTestExecutor.kt
@@ -20,14 +20,15 @@ import org.usvm.instrumentation.testcase.UTest
 import org.usvm.instrumentation.testcase.api.*
 import org.usvm.instrumentation.testcase.descriptor.Descriptor2ValueConverter
 import org.usvm.machine.*
+import org.usvm.machine.interpreter.JcFixedInheritorsNumberTypeSelector
+import org.usvm.machine.interpreter.JcTypeStreamPrioritization
 import org.usvm.machine.state.JcState
 import org.usvm.machine.state.localIdx
 import org.usvm.memory.ULValue
 import org.usvm.memory.UReadOnlyMemory
 import org.usvm.memory.URegisterStackLValue
 import org.usvm.model.UModelBase
 import org.usvm.types.first
-import org.usvm.types.firstOrNull
 
 /**
  * A class, responsible for resolving a single [JcTest] for a specific method from a symbolic state.
@@ -171,7 +172,9 @@ class JcTestExecutor(
         private val memory: UReadOnlyMemory<JcType>,
         private val method: JcTypedMethod,
     ) {
-
+        private val typeSelector = JcTypeStreamPrioritization(
+            typesToScore = JcFixedInheritorsNumberTypeSelector.DEFAULT_INHERITORS_NUMBER_TO_SCORE
+        )
         private val resolvedCache = mutableMapOf<UConcreteHeapAddress, Pair<UTestExpression, List<UTestInst>>>()
 
         fun createUTest(): UTest {
@@ -272,7 +275,8 @@ class JcTestExecutor(
             // and array types right now.
             // In such cases, we need to resolve this element to null.
 
-            val evaluatedType = typeStream.firstOrNull() ?: return UTestNullExpression(type) to listOf()
+            val evaluatedType = typeSelector.firstOrNull(typeStream, type.jcClass)
+                ?: return UTestNullExpression(type) to listOf()
 
             // We check for the type stream emptiness firsly and only then for the resolved cache,
             // because even if the object is already resolved, it could be incompatible with the [type], if it