Supported switch instruction in JaCoDB interpreter

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

@@ -18,7 +18,7 @@ abstract class UState<Type, Field, Method, Statement>(
     open val pathConstraints: UPathConstraints<Type>,
     open val memory: UMemoryBase<Field, Type, Method>,
     open var models: List<UModelBase<Field, Type>>,
-    open var path: PersistentList<Statement>
+    open var path: PersistentList<Statement>,
 ) {
     /**
      * Deterministic state id.
@@ -40,67 +40,71 @@ abstract class UState<Type, Field, Method, Statement>(
         get() = path.lastOrNull()
 }
 
-class ForkResult<T>(
+data class ForkResult<T>(
     val positiveState: T?,
     val negativeState: T?,
-) {
-    operator fun component1(): T? = positiveState
-    operator fun component2(): T? = negativeState
-}
+)
+
+private typealias StateToCheck = Boolean
+
+private const val ForkedState = true
+private const val OriginalState = false
 
 /**
- * Checks if [conditionToCheck] is satisfiable within path constraints of [state].
- * If it does, clones [state] and returns it with enriched constraints:
- * - if [forkToSatisfied], then adds constraint [satisfiedCondition];
- * - if ![forkToSatisfied], then adds constraint [conditionToCheck].
- * Otherwise, returns null.
- * If [conditionToCheck] is not unsatisfiable (i.e., solver returns sat or unknown),
- * mutates [state] by adding new path constraint c:
- * - if [forkToSatisfied], then c = [conditionToCheck]
- * - if ![forkToSatisfied], then c = [satisfiedCondition]
+ * Checks [newConstraintToOriginalState] or [newConstraintToForkedState], depending on the value of [stateToCheck].
+ * Depending on the result of checking this condition, do the following:
+ * - On [UUnsatResult] - returns `null`;
+ * - On [UUnknownResult] - adds [newConstraintToOriginalState] to the path constraints of the [state],
+ * iff [addConstraintOnUnknown] is `true`, and returns null;
+ * - On [USatResult] - clones the original state and adds the [newConstraintToForkedState] to it, adds [newConstraintToOriginalState]
+ * to the original state, sets the satisfiable model to the corresponding state depending on the [stateToCheck], and returns the
+ * forked state.
+ *
  */
 private fun <T : UState<Type, Field, *, *>, Type, Field> forkIfSat(
     state: T,
-    satisfiedCondition: UBoolExpr,
-    conditionToCheck: UBoolExpr,
-    forkToSatisfied: Boolean,
+    newConstraintToOriginalState: UBoolExpr,
+    newConstraintToForkedState: UBoolExpr,
+    stateToCheck: StateToCheck,
+    addConstraintOnUnknown: Boolean = true,
 ): T? {
-    val pathConstraints = state.pathConstraints.clone()
-    pathConstraints += conditionToCheck
+    val constraintsToCheck = state.pathConstraints.clone()
 
-    if (pathConstraints.isFalse) {
-        return null
+    constraintsToCheck += if (stateToCheck) {
+        newConstraintToForkedState
+    } else {
+        newConstraintToOriginalState
     }
-
-    val solver = satisfiedCondition.uctx.solver<Field, Type, Any?>()
-    val satResult = solver.check(pathConstraints, useSoftConstraints = true)
+    val solver = newConstraintToForkedState.uctx.solver<Field, Type, Any?>()
+    val satResult = solver.check(constraintsToCheck, useSoftConstraints = true)
 
     return when (satResult) {
+        is UUnsatResult -> null
+
         is USatResult -> {
-            if (forkToSatisfied) {
+            // Note that we cannot extract common code here due to
+            // heavy plusAssign operator in path constraints.
+            // Therefore, it is better to reuse already constructed [constraintToCheck].
+            if (stateToCheck) {
                 @Suppress("UNCHECKED_CAST")
-                val forkedState = state.clone() as T
-                // TODO: implement path condition setter (don't forget to reset UMemoryBase:types!)
-                state.pathConstraints += conditionToCheck
-                state.models = listOf(satResult.model)
-                forkedState.pathConstraints += satisfiedCondition
+                val forkedState = state.clone(constraintsToCheck) as T
+                state.pathConstraints += newConstraintToOriginalState
+                forkedState.models = listOf(satResult.model)
                 forkedState
             } else {
                 @Suppress("UNCHECKED_CAST")
-                val forkedState = state.clone(pathConstraints) as T
-                state.pathConstraints += satisfiedCondition
-                forkedState.models = listOf(satResult.model)
+                val forkedState = state.clone() as T
+                state.pathConstraints += newConstraintToOriginalState
+                state.models = listOf(satResult.model)
+                // TODO: implement path condition setter (don't forget to reset UMemoryBase:types!)
+                forkedState.pathConstraints += newConstraintToForkedState
                 forkedState
             }
         }
 
-        is UUnsatResult -> null
-
         is UUnknownResult -> {
-            if (forkToSatisfied) {
-                state.pathConstraints += conditionToCheck
-            } else {
-                state.pathConstraints += satisfiedCondition
+            if (addConstraintOnUnknown) {
+                state.pathConstraints += newConstraintToOriginalState
             }
             null
         }
@@ -148,13 +152,18 @@ fun <T : UState<Type, Field, *, *>, Type, Field> fork(
 
         trueModels.isNotEmpty() -> state to forkIfSat(
             state,
-            condition,
-            notCondition,
-            forkToSatisfied = false
+            newConstraintToOriginalState = condition,
+            newConstraintToForkedState = notCondition,
+            stateToCheck = ForkedState
         )
 
         falseModels.isNotEmpty() -> {
-            val forkedState = forkIfSat(state, notCondition, condition, forkToSatisfied = true)
+            val forkedState = forkIfSat(
+                state,
+                newConstraintToOriginalState = condition,
+                newConstraintToForkedState = notCondition,
+                stateToCheck = OriginalState
+            )
 
             if (forkedState != null) {
                 state to forkedState
@@ -169,3 +178,55 @@ fun <T : UState<Type, Field, *, *>, Type, Field> fork(
     @Suppress("UNCHECKED_CAST")
     return ForkResult(posState, negState as T?)
 }
+
+/**
+ * Implements symbolic branching on few disjoint conditions.
+ *
+ * @return a list of states for each condition - `null` state
+ * means [UUnknownResult] or [UUnsatResult] of checking condition.
+ */
+fun <T : UState<Type, Field, *, *>, Type, Field> forkMulti(
+    state: T,
+    conditions: Iterable<UBoolExpr>,
+): List<T?> {
+    var curState = state
+    val result = mutableListOf<T?>()
+    for (condition in conditions) {
+        val (trueModels, _) = curState.models.partition { model ->
+            val holdsInModel = model.eval(condition)
+            check(holdsInModel is KInterpretedValue<UBoolSort>) {
+                "Evaluation in $model on condition $condition: expected true or false, but got $holdsInModel"
+            }
+            holdsInModel.isTrue
+        }
+
+        val nextRoot = if (trueModels.isNotEmpty()) {
+            @Suppress("UNCHECKED_CAST")
+            val root = curState.clone() as T
+
+            curState.models = trueModels
+            curState.pathConstraints += condition
+
+            root
+        } else {
+            val root = forkIfSat(
+                curState,
+                newConstraintToOriginalState = condition,
+                newConstraintToForkedState = condition.ctx.trueExpr,
+                stateToCheck = OriginalState,
+                addConstraintOnUnknown = false
+            )
+
+            root
+        }
+
+        if (nextRoot != null) {
+            result += curState
+            curState = nextRoot
+        } else {
+            result += null
+        }
+    }
+
+    return result
+}

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

@@ -1,24 +1,36 @@
 package org.usvm
 
+import org.usvm.StepScope.StepScopeState.CANNOT_BE_PROCESSED
+import org.usvm.StepScope.StepScopeState.CAN_BE_PROCESSED
+import org.usvm.StepScope.StepScopeState.DEAD
+
 /**
  * An auxiliary class, which carefully maintains forks and asserts via [fork] and [assert].
  * It should be created on every step in an interpreter.
  * You can think about an instance of [StepScope] as a monad `ExceptT null (State [T])`.
  *
- * An underlying state is `null`, iff one of the `condition`s passed to the [fork] was unsatisfiable.
+ * This scope is considered as [DEAD], iff the condition in [assert] was unsatisfiable or unknown.
+ * The underlying state cannot be processed further (see [CANNOT_BE_PROCESSED]),
+ * if the first passed to [fork] or [forkMulti] condition was unsatisfiable or unknown.
  *
  * To execute some function on a state, you should use [doWithState] or [calcOnState]. `null` is returned, when
- * the current state is `null`.
+ * this scope cannot be processed on the current step - see [CANNOT_BE_PROCESSED].
  *
  * @param originalState an initial state.
  */
 class StepScope<T : UState<Type, Field, *, *>, Type, Field>(
-    originalState: T,
+    private val originalState: T,
 ) {
     private val forkedStates = mutableListOf<T>()
-    private var curState: T? = originalState
-    var alive: Boolean = true
-        private set
+
+    private val alive: Boolean get() = stepScopeState != DEAD
+    private val canProcessFurtherOnCurrentStep: Boolean get() = stepScopeState == CAN_BE_PROCESSED
+
+    /**
+     * Determines whether we interact this scope on the current step.
+     * @see [StepScopeState].
+     */
+    private var stepScopeState: StepScopeState = CAN_BE_PROCESSED
 
     /**
      * @return forked states and the status of initial state.
@@ -30,45 +42,54 @@ class StepScope<T : UState<Type, Field, *, *>, Type, Field>(
      *
      * @return `null` if the underlying state is `null`.
      */
-    fun doWithState(block: T.() -> Unit): Unit? {
-        val state = curState ?: return null
-        state.block()
-        return Unit
-    }
+    fun doWithState(block: T.() -> Unit): Unit? =
+        if (canProcessFurtherOnCurrentStep) {
+            originalState.block()
+        } else {
+            null
+        }
 
     /**
      * Executes [block] on a state.
      *
      * @return `null` if the underlying state is `null`, otherwise returns result of calling [block].
      */
-    fun <R> calcOnState(block: T.() -> R): R? {
-        val state = curState ?: return null
-        return state.block()
-    }
+    fun <R> calcOnState(block: T.() -> R): R? =
+        if (canProcessFurtherOnCurrentStep) {
+            originalState.block()
+        } else {
+            null
+        }
 
     /**
      * Forks on a [condition], performing [blockOnTrueState] on a state satisfying [condition] and
      * [blockOnFalseState] on a state satisfying [condition].not().
      *
-     * If the [condition] is unsatisfiable, sets underlying state to `null`.
+     * If the [condition] is unsatisfiable or unknown, sets the scope state to the [CANNOT_BE_PROCESSED].
      *
-     * @return `null` if the [condition] is unsatisfiable.
+     * @return `null` if the [condition] is unsatisfiable or unknown.
      */
     fun fork(
         condition: UBoolExpr,
         blockOnTrueState: T.() -> Unit = {},
         blockOnFalseState: T.() -> Unit = {},
     ): Unit? {
-        val state = curState ?: return null
+        check(canProcessFurtherOnCurrentStep)
 
-        val (posState, negState) = fork(state, condition)
+        val (posState, negState) = fork(originalState, condition)
 
         posState?.blockOnTrueState()
-        curState = posState
+
+        if (posState == null) {
+            stepScopeState = CANNOT_BE_PROCESSED
+            check(negState === originalState)
+        } else {
+            check(posState === originalState)
+        }
 
         if (negState != null) {
             negState.blockOnFalseState()
-            if (negState !== state) {
+            if (negState !== originalState) {
                 forkedStates += negState
             }
         }
@@ -77,23 +98,75 @@ class StepScope<T : UState<Type, Field, *, *>, Type, Field>(
         return posState?.let { }
     }
 
+    /**
+     * Forks on a few disjoint conditions using `forkMulti` in `State.kt`
+     * and executes the corresponding block on each not-null state.
+     *
+     * NOTE: always sets the [stepScopeState] to the [CANNOT_BE_PROCESSED] value.
+     */
+    fun forkMulti(conditionsWithBlockOnStates: List<Pair<UBoolExpr, T.() -> Unit>>) {
+        check(canProcessFurtherOnCurrentStep)
+
+        val conditions = conditionsWithBlockOnStates.map { it.first }
+
+        val conditionStates = forkMulti(originalState, conditions)
+
+        val forkedStates = conditionStates.mapIndexedNotNull { idx, positiveState ->
+            val block = conditionsWithBlockOnStates[idx].second
+
+            positiveState?.apply(block)
+        }
+
+        stepScopeState = CANNOT_BE_PROCESSED
+        if (forkedStates.isEmpty()) {
+            stepScopeState = DEAD
+            return
+        }
+
+        val firstForkedState = forkedStates.first()
+        require(firstForkedState == originalState) {
+            "The original state $originalState was expected to become the first of forked states but $firstForkedState found"
+        }
+
+        // Interpret the first state as original and others as forked
+        this.forkedStates += forkedStates.subList(1, forkedStates.size)
+    }
+
     fun assert(
         constraint: UBoolExpr,
         block: T.() -> Unit = {},
     ): Unit? {
-        val state = curState ?: return null
+        check(canProcessFurtherOnCurrentStep)
 
-        val (posState, _) = fork(state, constraint)
+        val (posState, _) = fork(originalState, constraint)
 
         posState?.block()
-        curState = posState
 
         if (posState == null) {
-            alive = false
+            stepScopeState = DEAD
         }
 
         return posState?.let { }
     }
+
+    /**
+     * Represents the current state of this [StepScope].
+     */
+    private enum class StepScopeState {
+        /**
+         * Cannot be processed further with any actions.
+         */
+        DEAD,
+        /**
+         * Cannot be forked or asserted using [fork], [forkMulti] or [assert],
+         * but is considered as alive from the Machine's point of view.
+         */
+        CANNOT_BE_PROCESSED,
+        /**
+         * Can be forked using [fork] or [forkMulti] and asserted using [assert].
+         */
+        CAN_BE_PROCESSED;
+    }
 }
 
 /**