Improvements

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

@@ -1,10 +1,9 @@
 package org.usvm.ps
 
-import io.ksmt.utils.cast
 import org.usvm.UPathSelector
 import org.usvm.UState
 import org.usvm.statistics.PathsTreeNode
-import org.usvm.statistics.PathsTreeStatistics
+import org.usvm.statistics.TerminatedStateRemover
 import java.util.IdentityHashMap
 
 /**
@@ -21,7 +20,7 @@ import java.util.IdentityHashMap
  * @param pathsTreeStatistics symbolic execution tree statistics used to peek states from.
  * @param randomNonNegativeInt function returning non negative random integer used to select the next child in tree.
  * @param ignoreToken token to visit only the subtree of not removed states. Should be different for
- * different instances of [PathsTreeStatistics] consumers.
+ * different instances of [TerminatedStateRemover] consumers.
  */
 internal class RandomTreePathSelector<State : UState<*, *, *, Statement, *, State>, Statement>(
     private val root: PathsTreeNode<State, Statement>, // TODO fix docs
@@ -108,12 +107,16 @@ internal class RandomTreePathSelector<State : UState<*, *, *, Statement, *, Stat
     }
 
     override fun update(state: State) {
-        currentNodes[state.pathLocation.cast()] = true
+        registerLocation(state.pathLocation)
+    }
+
+    internal fun registerLocation(pathLocation: PathsTreeNode<State, Statement>) {
+        currentNodes[pathLocation] = true
     }
 
     override fun add(states: Collection<State>) {
         this.states.addAll(states)
-        states.forEach { currentNodes[it.pathLocation.cast()] = true }
+        states.forEach { currentNodes[it.pathLocation] = true }
     }
 
     override fun remove(state: State) {

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

@@ -1,19 +1,17 @@
 package org.usvm.statistics
 
-import io.ksmt.utils.cast
-import org.usvm.UContext
 import org.usvm.UState
 
 /**
  * Symbolic execution tree node.
  *
- * @see [PathsTreeStatistics]
+ * @see [TerminatedStateRemover]
  */
 sealed class PathsTreeNode<State : UState<*, *, *, Statement, *, State>, Statement> {
     /**
      * Forked states' nodes.
      */
-    abstract val children: MutableMap<Statement, PathsTreeNode<State, Statement>>
+    abstract val children: MutableMap<Statement, PathsTreeNodeImpl<State, Statement>>
 
     /**
      * State which is located in this node. Null for non-leaf nodes.
@@ -43,32 +41,47 @@ sealed class PathsTreeNode<State : UState<*, *, *, Statement, *, State>, Stateme
      */
     abstract fun addIgnoreToken(token: Long)
 
-    fun getChildOrCreate(statement: Statement, childCreator: () -> PathsTreeNode<State, Statement>) =
-        children.getOrPut(statement, childCreator)
+    /**
+     * TODO add docs
+     */
+    fun propagateState(statement: Statement, state: State): PathsTreeNodeImpl<State, Statement> {
+        val child = children[statement]
+
+        if (child != null) {
+            child.states += state
+            states -= state
+            return child
+        }
+
+        val node = when (this) {
+            is RootNode -> PathsTreeNodeImpl(parentNode = this, statement, state)
+            is PathsTreeNodeImpl -> PathsTreeNodeImpl(parentNode = this, statement, state)
+        }
+
+        return node
+    }
+
 }
 
 // TODO move it from here
-class PathsTreeNodeImpl<State : UState<*, *, *, Statement, *, State>, Statement : Any> private constructor(
+class PathsTreeNodeImpl<State : UState<*, *, *, Statement, *, State>, Statement> private constructor(
     override val depth: Int,
-    private val childrenImpl: MutableMap<Statement, PathsTreeNode<State, Statement>>,
     override var states: MutableSet<State>,
+    override val children: MutableMap<Statement, PathsTreeNodeImpl<State, Statement>> = mutableMapOf(), // Note: order is important for tests
     override val parent: PathsTreeNode<State, Statement>?,
-    override val statement: Statement
+    override val statement: Statement,
 ) : PathsTreeNode<State, Statement>() {
-
-    constructor(parentNode: RootNode<State, Statement>, statement: Statement, state: State) : this(
+    internal constructor(parentNode: RootNode<State, Statement>, statement: Statement, state: State) : this(
         depth = parentNode.depth + 1,
-        childrenImpl = hashMapOf(), // TODO is order important?
         parent = parentNode,
         states = hashSetOf(state),
         statement = statement
     ) {
         parentNode.children[statement] = this
     }
 
-    constructor(parentNode: PathsTreeNodeImpl<State, Statement>, statement: Statement, state: State) : this(
+    internal constructor(parentNode: PathsTreeNodeImpl<State, Statement>, statement: Statement, state: State) : this(
         depth = parentNode.depth + 1,
-        childrenImpl = hashMapOf(), // TODO is order important?
         parent = parentNode,
         states = mutableSetOf(state),
         statement = statement
@@ -77,23 +90,18 @@ class PathsTreeNodeImpl<State : UState<*, *, *, Statement, *, State>, Statement
         parentNode.states -= state
     }
 
-    val locker = Any()
     override val ignoreTokens: MutableSet<Long> = hashSetOf()
 
-    override val children: MutableMap<Statement, PathsTreeNode<State, Statement>>
-        get() = synchronized(locker) { childrenImpl }
-
     override fun addIgnoreToken(token: Long) {
-        synchronized(locker) {
-            ignoreTokens.add(token)
-        }
+        ignoreTokens.add(token)
     }
 
+
     override fun toString(): String = "Depth: $depth, statement: $statement"
 }
 
 class RootNode<State : UState<*, *, *, Statement, *, State>, Statement> : PathsTreeNode<State, Statement>() {
-    override val children: MutableMap<Statement, PathsTreeNode<State, Statement>> = mutableMapOf()
+    override val children: MutableMap<Statement, PathsTreeNodeImpl<State, Statement>> = mutableMapOf()
 
     override var states: MutableSet<State> = hashSetOf()
 
@@ -122,21 +130,10 @@ class RootNode<State : UState<*, *, *, Statement, *, State>, Statement> : PathsT
  *
  * @param initialState root state.
  */
-class PathsTreeStatistics<Statement, Context : UContext, State : UState<*, *, *, Statement, Context, State>>(
-    initialState: State,
+// TODO rename it
+class TerminatedStateRemover<State : UState<*, *, *, *, *, State>>(
 ) : UMachineObserver<State> {
-
-    val root: PathsTreeNode<State, Statement> = // TODO that's not cool to have `parent` here
-        initialState.pathLocation.parent.cast() // TODO either we add one generic to state, or cast it here
-
-    override fun onState(
-        parent: State,
-        forks: Sequence<State>,
-    ) {
-        // TODO log forks???
-    }
-
     override fun onStateTerminated(state: State) {
-        state.pathLocation.states.remove(state) // TODO do we need it? State is contained only in itself
+        state.pathLocation.states.remove(state)
     }
 }

usvm-core/src/test/kotlin/org/usvm/PathTests.kt

@@ -0,0 +1,52 @@
+package org.usvm
+
+import io.mockk.mockk
+import org.junit.jupiter.api.Assertions.assertSame
+import org.junit.jupiter.api.Test
+import org.usvm.statistics.RootNode
+import kotlin.test.assertTrue
+
+class PathTests {
+    @Test
+    fun smokeTest() {
+        val initialState = mockk<TestState>()
+        val fork = mockk<TestState>()
+
+        val root = RootNode<TestState, Int>()
+
+        val firstRealNode = root.propagateState(1, initialState)
+
+        val updatedInitialState = firstRealNode.propagateState(statement = 2, initialState)
+        val forkedState = firstRealNode.propagateState(statement = 3, fork)
+
+        assertTrue { root.children.size == 1 }
+        assertTrue { root.children.values.single() == firstRealNode }
+
+        assertTrue { firstRealNode.states.isEmpty() }
+        assertTrue { firstRealNode.children.size == 2 }
+
+        assertTrue { updatedInitialState.states.single() == initialState }
+        assertTrue { forkedState.states.single() == fork }
+
+        assertTrue { updatedInitialState.children.isEmpty() }
+        assertTrue { forkedState.children.isEmpty() }
+    }
+
+    @Test
+    fun severalStatesInOneNode() {
+        val firstState = mockk<TestState>()
+        val secondState = mockk<TestState>()
+
+        val root = RootNode<TestState, Int>()
+
+        val firstRealNode = root.propagateState(1, firstState)
+
+        val updatedFirstState = firstRealNode.propagateState(statement = 2, firstState)
+        val updatedSecondState = firstRealNode.propagateState(statement = 2, secondState)
+
+        assertSame(updatedFirstState, updatedSecondState)
+
+        assertTrue { updatedFirstState.children.isEmpty() }
+        assertTrue { updatedFirstState.states.size == 2 }
+    }
+}

usvm-core/src/test/kotlin/org/usvm/TestUtil.kt

@@ -1,5 +1,10 @@
 package org.usvm
 
+import org.usvm.constraints.UPathConstraints
+import org.usvm.memory.UMemoryBase
+import org.usvm.model.UModelBase
+import org.usvm.statistics.PathsTreeNode
+
 typealias Field = java.lang.reflect.Field
 typealias Type = kotlin.reflect.KClass<*>
 typealias Method = kotlin.reflect.KFunction<*>
@@ -19,3 +24,14 @@ internal fun pseudoRandom(i: Int): Int {
     res = (res shr 16) xor res
     return res
 }
+
+internal class TestState(
+    ctx: UContext,
+    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, UContext, TestState>(ctx, callStack, pathConstraints, memory, models, pathLocation) {
+    override fun clone(newConstraints: UPathConstraints<Any, UContext>?): TestState = this
+
+    override val isExceptional = false
+}