Fix tests

UnitTestBot · Aug 8, 2023 · 7383744 · 7383744
1 parent b8099d8
commit 7383744
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/usvm-core/src/main/kotlin/org/usvm/ps/PathSelectorFactory.kt b/usvm-core/src/main/kotlin/org/usvm/ps/PathSelectorFactory.kt
@@ -29,7 +29,7 @@ fun <Method, Statement, State : UState<*, *, Method, Statement>> createPathSelec
             PathSelectionStrategy.BFS -> BfsPathSelector()
             PathSelectionStrategy.DFS -> DfsPathSelector()
             PathSelectionStrategy.RANDOM_PATH -> RandomTreePathSelector(
-                root = requireNotNull(initialState.pathLocation.cast()), // TODO message
+                root = requireNotNull(initialState.pathLocation.parent.cast()), // TODO message
                 randomNonNegativeInt = { random.nextInt(0, Int.MAX_VALUE) }
             )
 

diff --git a/usvm-core/src/main/kotlin/org/usvm/ps/RandomTreePathSelector.kt b/usvm-core/src/main/kotlin/org/usvm/ps/RandomTreePathSelector.kt
@@ -1,10 +1,11 @@
 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.RootNode
+import java.util.IdentityHashMap
 
 /**
  * [UPathSelector] implementation which selects the next state by descending
@@ -28,7 +29,13 @@ internal class RandomTreePathSelector<State : UState<*, *, *, Statement>, Statem
     private val ignoreToken: Long = 0,
 ) : UPathSelector<State> {
 
-    private val states = HashSet<State>()
+    private val states = hashSetOf<State>()
+    // TODO add comment
+    private val currentNodes = IdentityHashMap<PathsTreeNode<State, Statement>, Boolean>()
+
+    init {
+        currentNodes[root] = true
+    }
 
     override fun isEmpty(): Boolean {
         return states.isEmpty()
@@ -56,13 +63,17 @@ internal class RandomTreePathSelector<State : UState<*, *, *, Statement>, Statem
             val children = currentNode
                 .children
                 .values
-                .filterNot { it.ignoreTokens.contains(ignoreToken) || it.states.all { state -> state !in states } }
+                .filter { !it.ignoreTokens.contains(ignoreToken) && it in currentNodes }
 
             if (children.isEmpty()) {
                 currentNode.addIgnoreToken(ignoreToken)
-                val parent = currentNode.parent ?: break
+                states -= currentNode.states // TODO do we need to remove all the states?
 
-                if (parent is RootNode) break // TODO add a comment about it
+                if (currentNode.parent == null) {
+                    require(states.isEmpty())
+                }
+
+                val parent = currentNode.parent ?: break
 
                 currentNode = parent
                 // All children are excluded from search, exclude the parent from the
@@ -87,20 +98,26 @@ internal class RandomTreePathSelector<State : UState<*, *, *, Statement>, Statem
         // Peeked a state from tree, but it's not in our states. Exclude the node from further search
         // and try again
         currentNode.addIgnoreToken(ignoreToken)
+        states -= currentNode.states // TODO do we need to remove all the states?
+
         return peekRec()
     }
 
     override fun peek(): State {
         return peekRec()
     }
 
-    override fun update(state: State) { }
+    override fun update(state: State) {
+        currentNodes[state.pathLocation.cast()] = true
+    }
 
     override fun add(states: Collection<State>) {
         this.states.addAll(states)
+        states.forEach { currentNodes[it.pathLocation.cast()] = true }
     }
 
     override fun remove(state: State) {
         states.remove(state)
+        // TODO do we need to remove nodes?
     }
 }