Add generics

usvm-core/src/main/kotlin/org/usvm/Context.kt

@@ -43,8 +43,8 @@ open class UContext(
     }
 
     @Suppress("UNCHECKED_CAST")
-    fun <Field, Type, Method> solver(): USolverBase<Field, Type, Method> =
-        this.solver as USolverBase<Field, Type, Method>
+    fun <Field, Type, Method, Context : UContext> solver(): USolverBase<Field, Type, Method, Context> =
+        this.solver as USolverBase<Field, Type, Method, Context>
 
     @Suppress("UNCHECKED_CAST")
     fun <Type> typeSystem(): UTypeSystem<Type> =
@@ -194,7 +194,7 @@ open class UContext(
 
     private val initialLocation = RootNode<Nothing, Nothing>() // TODO think about it later
 
-    fun <State : UState<*, *, *, Statement, State>, Statement> mkInitialLocation()
+    fun <State : UState<*, *, *, Statement, *, State>, Statement> mkInitialLocation()
         : PathsTreeNode<State, Statement> = initialLocation.uncheckedCast()
 
     fun mkUValueSampler(): KSortVisitor<KExpr<*>> {

usvm-core/src/main/kotlin/org/usvm/Merging.kt

@@ -14,7 +14,7 @@ class URegionHeapMerger<FieldDescr, ArrayType> : UMerger<URegionHeap<FieldDescr,
     override fun merge(left: URegionHeap<FieldDescr, ArrayType>, right: URegionHeap<FieldDescr, ArrayType>) = null
 }
 
-open class UStateMerger<State : UState<*, *, *, *, State>> : UMerger<State> {
+open class UStateMerger<State : UState<*, *, *, *, *, State>> : UMerger<State> {
     // Never merge for now
     override fun merge(left: State, right: State) = null
 }

usvm-core/src/main/kotlin/org/usvm/State.kt

@@ -12,11 +12,11 @@ import org.usvm.statistics.RootNode
 
 typealias StateId = UInt
 
-abstract class UState<Type, Field, Method, Statement, State : UState<Type, Field, Method, Statement, State>>(
+abstract class UState<Type, Field, Method, Statement, Context : UContext, State : UState<Type, Field, Method, Statement, Context, State>>(
     // TODO: add interpreter-specific information
     ctx: UContext,
     open val callStack: UCallStack<Method, Statement>,
-    open val pathConstraints: UPathConstraints<Type>,
+    open val pathConstraints: UPathConstraints<Type, Context>,
     open val memory: UMemoryBase<Field, Type, Method>,
     open var models: List<UModelBase<Field, Type>>,
     open var pathLocation: PathsTreeNode<State, Statement>,
@@ -50,7 +50,7 @@ abstract class UState<Type, Field, Method, Statement, State : UState<Type, Field
      * Creates new state structurally identical to this.
      * If [newConstraints] is null, clones [pathConstraints]. Otherwise, uses [newConstraints] in cloned state.
      */
-    abstract fun clone(newConstraints: UPathConstraints<Type>? = null): State
+    abstract fun clone(newConstraints: UPathConstraints<Type, Context>? = null): State
 
     val lastEnteredMethod: Method
         get() = callStack.lastMethod()
@@ -85,7 +85,7 @@ private const val OriginalState = false
  * forked state.
  *
  */
-private fun <T : UState<Type, Field, *, *, T>, Type, Field> forkIfSat(
+private fun <T : UState<Type, Field, *, *, Context, T>, Type, Field, Context : UContext> forkIfSat(
     state: T,
     newConstraintToOriginalState: UBoolExpr,
     newConstraintToForkedState: UBoolExpr,
@@ -99,7 +99,7 @@ private fun <T : UState<Type, Field, *, *, T>, Type, Field> forkIfSat(
     } else {
         newConstraintToOriginalState
     }
-    val solver = newConstraintToForkedState.uctx.solver<Field, Type, Any?>()
+    val solver = newConstraintToForkedState.uctx.solver<Field, Type, Any?, Context>()
     val satResult = solver.check(constraintsToCheck, useSoftConstraints = true)
 
     return when (satResult) {
@@ -145,7 +145,7 @@ private fun <T : UState<Type, Field, *, *, T>, Type, Field> forkIfSat(
  * 2. makes not more than one query to USolver;
  * 3. if both [condition] and ![condition] are satisfiable, then [ForkResult.positiveState] === [state].
  */
-fun <T : UState<Type, Field, *, *, T>, Type, Field> fork(
+fun <T : UState<Type, Field, *, *, Context, T>, Type, Field, Context : UContext> fork(
     state: T,
     condition: UBoolExpr,
 ): ForkResult<T> {
@@ -206,7 +206,7 @@ fun <T : UState<Type, Field, *, *, T>, Type, Field> fork(
  * @return a list of states for each condition - `null` state
  * means [UUnknownResult] or [UUnsatResult] of checking condition.
  */
-fun <T : UState<Type, Field, *, *, T>, Type, Field> forkMulti(
+fun <T : UState<Type, Field, *, *, Context, T>, Type, Field, Context : UContext> forkMulti(
     state: T,
     conditions: Iterable<UBoolExpr>,
 ): List<T?> {

usvm-core/src/main/kotlin/org/usvm/StepScope.kt

@@ -18,7 +18,7 @@ import org.usvm.StepScope.StepScopeState.DEAD
  *
  * @param originalState an initial state.
  */
-class StepScope<T : UState<Type, Field, *, *, T>, Type, Field>(
+class StepScope<T : UState<Type, Field, *, *, Context, T>, Type, Field, Context : UContext>(
     private val originalState: T,
 ) {
     private val forkedStates = mutableListOf<T>()

usvm-core/src/main/kotlin/org/usvm/UComponents.kt

@@ -8,6 +8,6 @@ import org.usvm.types.UTypeSystem
  * Instatiated once per [UContext].
  */
 interface UComponents<Field, Type, Method> {
-    fun mkSolver(ctx: UContext): USolverBase<Field, Type, Method>
+    fun <Context : UContext> mkSolver(ctx: Context): USolverBase<Field, Type, Method, Context>
     fun mkTypeSystem(ctx: UContext): UTypeSystem<Type>
 }

usvm-core/src/main/kotlin/org/usvm/constraints/PathConstraints.kt

@@ -18,8 +18,8 @@ import org.usvm.uctx
 /**
  * Mutable representation of path constraints.
  */
-open class UPathConstraints<Type> private constructor(
-    val ctx: UContext,
+open class UPathConstraints<Type, Context : UContext> private constructor(
+    val ctx: Context,
     logicalConstraints: PersistentSet<UBoolExpr> = persistentSetOf(),
     /**
      * Specially represented equalities and disequalities between objects, used in various part of constraints management.
@@ -39,7 +39,7 @@ open class UPathConstraints<Type> private constructor(
     var logicalConstraints: PersistentSet<UBoolExpr> = logicalConstraints
         private set
 
-    constructor(ctx: UContext) : this(ctx, persistentSetOf())
+    constructor(ctx: Context) : this(ctx, persistentSetOf())
 
     open val isFalse: Boolean
         get() = equalityConstraints.isContradicting ||
@@ -105,7 +105,7 @@ open class UPathConstraints<Type> private constructor(
             }
         }
 
-    open fun clone(): UPathConstraints<Type> {
+    open fun clone(): UPathConstraints<Type, Context> {
         val clonedEqualityConstraints = equalityConstraints.clone()
         val clonedTypeConstraints = typeConstraints.clone(clonedEqualityConstraints)
         return UPathConstraints(ctx, logicalConstraints, clonedEqualityConstraints, clonedTypeConstraints)

usvm-core/src/main/kotlin/org/usvm/ps/ExceptionPropagationPathSelector.kt

@@ -7,7 +7,7 @@ import java.util.IdentityHashMap
 /**
  * A class designed to give the highest priority to the states containing exceptions.
  */
-class ExceptionPropagationPathSelector<State : UState<*, *, *, *, State>>(
+class ExceptionPropagationPathSelector<State : UState<*, *, *, *, *, State>>(
     private val selector: UPathSelector<State>,
 ) : UPathSelector<State> {
     // An internal queue for states containing exceptions.

usvm-core/src/main/kotlin/org/usvm/ps/PathSelectorFactory.kt

@@ -13,7 +13,7 @@ import org.usvm.util.RandomizedPriorityCollection
 import kotlin.math.max
 import kotlin.random.Random
 
-fun <Method, Statement, State : UState<*, *, Method, Statement, State>> createPathSelector(
+fun <Method, Statement, State : UState<*, *, Method, Statement, *, State>> createPathSelector(
     initialState: State,
     options: UMachineOptions,
     coverageStatistics: () -> CoverageStatistics<Method, Statement, State>? = { null },
@@ -87,16 +87,16 @@ fun <Method, Statement, State : UState<*, *, Method, Statement, State>> createPa
 /**
  * Wraps the selector into an [ExceptionPropagationPathSelector] if [propagateExceptions] is true.
  */
-private fun <State : UState<*, *, *, *, State>> UPathSelector<State>.wrapIfRequired(propagateExceptions: Boolean) =
+private fun <State : UState<*, *, *, *, *, State>> UPathSelector<State>.wrapIfRequired(propagateExceptions: Boolean) =
     if (propagateExceptions && this !is ExceptionPropagationPathSelector<State>) {
         ExceptionPropagationPathSelector(this)
     } else {
         this
     }
 
-private fun <State : UState<*, *, *, *, State>> compareById(): Comparator<State> = compareBy { it.id }
+private fun <State : UState<*, *, *, *, *, State>> compareById(): Comparator<State> = compareBy { it.id }
 
-private fun <State : UState<*, *, *, *, State>> createDepthPathSelector(random: Random? = null): UPathSelector<State> {
+private fun <State : UState<*, *, *, *, *, State>> createDepthPathSelector(random: Random? = null): UPathSelector<State> {
     if (random == null) {
         return WeightedPathSelector(
             priorityCollectionFactory = { DeterministicPriorityCollection(Comparator.naturalOrder()) },
@@ -111,7 +111,7 @@ private fun <State : UState<*, *, *, *, State>> createDepthPathSelector(random:
     )
 }
 
-private fun <Method, Statement, State : UState<*, *, Method, Statement, State>> createClosestToUncoveredPathSelector(
+private fun <Method, Statement, State : UState<*, *, Method, Statement, *, State>> createClosestToUncoveredPathSelector(
     coverageStatistics: CoverageStatistics<Method, Statement, State>,
     distanceStatistics: DistanceStatistics<Method, Statement>,
     random: Random? = null,
@@ -137,7 +137,7 @@ private fun <Method, Statement, State : UState<*, *, Method, Statement, State>>
     )
 }
 
-private fun <Method, Statement, State : UState<*, *, Method, Statement, State>> createForkDepthPathSelector(
+private fun <Method, Statement, State : UState<*, *, Method, Statement, *, State>> createForkDepthPathSelector(
     random: Random? = null,
 ): UPathSelector<State> {
     if (random == null) {

usvm-core/src/main/kotlin/org/usvm/ps/RandomTreePathSelector.kt

@@ -23,7 +23,7 @@ import java.util.IdentityHashMap
  * @param ignoreToken token to visit only the subtree of not removed states. Should be different for
  * different instances of [PathsTreeStatistics] consumers.
  */
-internal class RandomTreePathSelector<State : UState<*, *, *, Statement, State>, Statement>(
+internal class RandomTreePathSelector<State : UState<*, *, *, Statement, *, State>, Statement>(
     private val root: PathsTreeNode<State, Statement>, // TODO fix docs
     private val randomNonNegativeInt: () -> Int,
     private val ignoreToken: Long = 0,

usvm-core/src/main/kotlin/org/usvm/ps/ShortestDistanceToTargetsStateWeighter.kt

@@ -16,7 +16,7 @@ import kotlin.math.min
  * @param getCfgDistanceToExitPoint function with the following signature:
  * (method, stmt) -> shortest CFG distance from stmt to any of method's exit points.
  */
-class ShortestDistanceToTargetsStateWeighter<Method, Statement, State : UState<*, *, Method, Statement, State>>(
+class ShortestDistanceToTargetsStateWeighter<Method, Statement, State : UState<*, *, Method, Statement, *, State>>(
     targets: Collection<Pair<Method, Statement>>,
     private val getCfgDistance: (Method, Statement, Statement) -> UInt,
     private val getCfgDistanceToExitPoint: (Method, Statement) -> UInt

usvm-core/src/main/kotlin/org/usvm/solver/Solver.kt

@@ -27,13 +27,13 @@ abstract class USolver<in PathCondition, out Model> : AutoCloseable {
     abstract fun check(pc: PathCondition, useSoftConstraints: Boolean): USolverResult<Model>
 }
 
-open class USolverBase<Field, Type, Method>(
-    protected val ctx: UContext,
+open class USolverBase<Field, Type, Method, Context : UContext>(
+    protected val ctx: Context,
     protected val smtSolver: KSolver<*>,
     protected val translator: UExprTranslator<Field, Type>,
     protected val decoder: UModelDecoder<UMemoryBase<Field, Type, Method>, UModelBase<Field, Type>>,
     protected val softConstraintsProvider: USoftConstraintsProvider<Field, Type>,
-) : USolver<UPathConstraints<Type>, UModelBase<Field, Type>>() {
+) : USolver<UPathConstraints<Type, Context>, UModelBase<Field, Type>>() {
 
     protected fun translateLogicalConstraints(constraints: Iterable<UBoolExpr>) {
         for (constraint in constraints) {
@@ -90,17 +90,17 @@ open class USolverBase<Field, Type, Method>(
         }
     }
 
-    protected fun translateToSmt(pc: UPathConstraints<Type>) {
+    protected fun translateToSmt(pc: UPathConstraints<Type, Context>) {
         translateEqualityConstraints(pc.equalityConstraints)
         translateLogicalConstraints(pc.logicalConstraints)
     }
 
     internal fun checkWithSoftConstraints(
-        pc: UPathConstraints<Type>,
+        pc: UPathConstraints<Type, Context>,
     ) = check(pc, useSoftConstraints = true)
 
     override fun check(
-        pc: UPathConstraints<Type>,
+        pc: UPathConstraints<Type, Context>,
         useSoftConstraints: Boolean,
     ): USolverResult<UModelBase<Field, Type>> {
         if (pc.isFalse) {

usvm-core/src/main/kotlin/org/usvm/statistics/CoverageStatistics.kt

@@ -13,7 +13,7 @@ import java.util.concurrent.ConcurrentHashMap
  * @param methods methods to track coverage of.
  * @param applicationGraph [ApplicationGraph] used to retrieve statements by method.
  */
-class CoverageStatistics<Method, Statement, State : UState<*, *, Method, Statement, State>>(
+class CoverageStatistics<Method, Statement, State : UState<*, *, Method, Statement, *, State>>(
     methods: Set<Method>,
     private val applicationGraph: ApplicationGraph<Method, Statement>
 ) : UMachineObserver<State> {

usvm-core/src/main/kotlin/org/usvm/statistics/PathsTreeStatistics.kt

@@ -1,14 +1,15 @@
 package org.usvm.statistics
 
 import io.ksmt.utils.cast
+import org.usvm.UContext
 import org.usvm.UState
 
 /**
  * Symbolic execution tree node.
  *
  * @see [PathsTreeStatistics]
  */
-sealed class PathsTreeNode<State : UState<*, *, *, Statement, State>, Statement> {
+sealed class PathsTreeNode<State : UState<*, *, *, Statement, *, State>, Statement> {
     /**
      * Forked states' nodes.
      */
@@ -47,7 +48,7 @@ sealed class PathsTreeNode<State : UState<*, *, *, Statement, State>, Statement>
 }
 
 // TODO move it from here
-class PathsTreeNodeImpl<State : UState<*, *, *, Statement, State>, Statement : Any> private constructor(
+class PathsTreeNodeImpl<State : UState<*, *, *, Statement, *, State>, Statement : Any> private constructor(
     override val depth: Int,
     private val childrenImpl: MutableMap<Statement, PathsTreeNode<State, Statement>>,
     override var states: MutableSet<State>,
@@ -91,7 +92,7 @@ class PathsTreeNodeImpl<State : UState<*, *, *, Statement, State>, Statement : A
     override fun toString(): String = "Depth: $depth, statement: $statement"
 }
 
-class RootNode<State : UState<*, *, *, Statement, State>, Statement> : PathsTreeNode<State, Statement>() {
+class RootNode<State : UState<*, *, *, Statement, *, State>, Statement> : PathsTreeNode<State, Statement>() {
     override val children: MutableMap<Statement, PathsTreeNode<State, Statement>> = mutableMapOf()
 
     override var states: MutableSet<State> = hashSetOf()
@@ -121,7 +122,7 @@ class RootNode<State : UState<*, *, *, Statement, State>, Statement> : PathsTree
  *
  * @param initialState root state.
  */
-class PathsTreeStatistics<Statement, State : UState<*, *, *, Statement, State>>(
+class PathsTreeStatistics<Statement, Context : UContext, State : UState<*, *, *, Statement, Context, State>>(
     initialState: State,
 ) : UMachineObserver<State> {
 

usvm-core/src/test/kotlin/org/usvm/model/ModelDecodingTest.kt

@@ -23,15 +23,14 @@ import org.usvm.solver.USatResult
 import org.usvm.solver.USoftConstraintsProvider
 import org.usvm.solver.USolverBase
 import org.usvm.solver.UUnsatResult
-import org.usvm.types.USingleTypeStream
 import org.usvm.types.single.SingleTypeSystem
 import kotlin.test.assertIs
 
 class ModelDecodingTest {
     private lateinit var ctx: UContext
-    private lateinit var solver: USolverBase<Field, Type, Method>
+    private lateinit var solver: USolverBase<Field, Type, Method, UContext>
 
-    private lateinit var pc: UPathConstraints<Type>
+    private lateinit var pc: UPathConstraints<Type, UContext>
     private lateinit var stack: URegistersStack
     private lateinit var heap: URegionHeap<Field, Type>
     private lateinit var mocker: UIndexedMocker<Method>

usvm-core/src/test/kotlin/org/usvm/ps/ShortestDistanceToTargetsStateWeighterTests.kt

@@ -137,11 +137,11 @@ internal class ShortestDistanceToTargetsStateWeighterTests {
 
     private class TestState(
         ctx: UContext,
-        callStack: UCallStack<String, Int>, pathConstraints: UPathConstraints<Any>,
+        callStack: UCallStack<String, Int>, pathConstraints: UPathConstraints<Any, UContext>,
         memory: UMemoryBase<Any, Any, String>, models: List<UModelBase<Any, Any>>,
         pathLocation: PathsTreeNode<TestState, Int>,
-    ) : UState<Any, Any, String, Int, TestState>(ctx, callStack, pathConstraints, memory, models, pathLocation) {
-        override fun clone(newConstraints: UPathConstraints<Any>?): TestState = this
+    ) : UState<Any, Any, String, Int, UContext, TestState>(ctx, callStack, pathConstraints, memory, models, pathLocation) {
+        override fun clone(newConstraints: UPathConstraints<Any, UContext>?): TestState = this
 
         override val isExceptional = false
     }

usvm-core/src/test/kotlin/org/usvm/solver/SoftConstraintsTest.kt

@@ -14,7 +14,6 @@ import org.usvm.constraints.UPathConstraints
 import org.usvm.memory.emptyInputArrayLengthRegion
 import org.usvm.model.ULazyModelDecoder
 import org.usvm.model.buildTranslatorAndLazyDecoder
-import org.usvm.types.UTypeSystem
 import org.usvm.types.single.SingleTypeSystem
 import kotlin.test.assertSame
 
@@ -23,7 +22,7 @@ open class SoftConstraintsTest<Field, Type, Method> {
     private lateinit var softConstraintsProvider: USoftConstraintsProvider<Field, Type>
     private lateinit var translator: UExprTranslator<Field, Type>
     private lateinit var decoder: ULazyModelDecoder<Field, Type, Method>
-    private lateinit var solver: USolverBase<Field, Type, Method>
+    private lateinit var solver: USolverBase<Field, Type, Method, UContext>
 
     @BeforeEach
     fun initialize() {
@@ -46,7 +45,7 @@ open class SoftConstraintsTest<Field, Type, Method> {
         val sndRegister = mkRegisterReading(idx = 1, bv32Sort)
         val expr = mkBvSignedLessOrEqualExpr(fstRegister, sndRegister)
 
-        val pc = UPathConstraints<Type>(ctx)
+        val pc = UPathConstraints<Type, UContext>(ctx)
         pc += expr
 
         val result = solver.checkWithSoftConstraints(pc) as USatResult
@@ -73,7 +72,7 @@ open class SoftConstraintsTest<Field, Type, Method> {
 
         every { softConstraintsProvider.provide(any()) } answers { callOriginal() }
 
-        val pc = UPathConstraints<Type>(ctx)
+        val pc = UPathConstraints<Type, UContext>(ctx)
         pc += fstExpr
         pc += sndExpr
         pc += sameAsFirstExpr
@@ -118,7 +117,7 @@ open class SoftConstraintsTest<Field, Type, Method> {
 
         val reading = region.read(secondInputRef)
 
-        val pc = UPathConstraints<Type>(ctx)
+        val pc = UPathConstraints<Type, UContext>(ctx)
         pc += reading eq size.toBv()
         pc += inputRef eq secondInputRef
         pc += (inputRef eq nullRef).not()
@@ -138,7 +137,7 @@ open class SoftConstraintsTest<Field, Type, Method> {
         val region = emptyInputArrayLengthRegion(arrayType, sizeSort)
             .write(inputRef, mkRegisterReading(3, sizeSort), guard = trueExpr)
 
-        val pc = UPathConstraints<Type>(ctx)
+        val pc = UPathConstraints<Type, UContext>(ctx)
         pc += (inputRef eq nullRef).not()
         val result = (solver.checkWithSoftConstraints(pc)) as USatResult
 
@@ -154,7 +153,7 @@ open class SoftConstraintsTest<Field, Type, Method> {
         val bvValue = 0.toBv()
         val expression = mkBvSignedLessOrEqualExpr(bvValue, inputRef).not()
 
-        val pc = UPathConstraints<Type>(ctx)
+        val pc = UPathConstraints<Type, UContext>(ctx)
         pc += expression
         val result = (solver.checkWithSoftConstraints(pc)) as USatResult