Track last peeked state

PySymGym · Oct 12, 2024 · 5d18860 · 5d18860
1 parent 2bb956c
commit 5d18860
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/usvm-core/src/main/kotlin/org/usvm/ps/AIPathSelector.kt b/usvm-core/src/main/kotlin/org/usvm/ps/AIPathSelector.kt
@@ -20,7 +20,8 @@ class AIPathSelector<Statement, State, Block>(
 State : UState<*, *, Statement, *, *, State>,
 Block : BasicBlock {
     private val statesMap = mutableMapOf<State, StateWrapper<Statement, State, Block>>()
-    private var lastPeekedState: State? = null
+    private val lastPeekedState: State?
+        get() = stepsStatistics.lastPeekedState
     private val totalSteps
         get() = stepsStatistics.totalSteps.toInt()
 
@@ -29,6 +30,8 @@ Block : BasicBlock {
     private var touchedBlocks = mutableSetOf<Block>()
     private var newBlocks = listOf<Block>()
 
+    private var firstSend = true
+
     private fun predict(): State {
         val wrappers = statesMap.values
         val vertices = blockGraph.blocks
@@ -46,9 +49,10 @@ Block : BasicBlock {
         val game = when (predictor) {
             is OnnxModel<Game<Block>> -> Game(vertices, wrappers, blockGraph)
             is Oracle<Game<Block>> -> {
-                // if we played with default searcher before
+                // if we played with default searcher before,
                 // client has no information about the game
-                if (lastPeekedState == null) {
+                if (firstSend) {
+                    firstSend = false
                     Game(vertices, wrappers, blockGraph)
                 } else {
                     if (blockGraph.newBlocks != newBlocks) {
@@ -71,12 +75,11 @@ Block : BasicBlock {
 
     override fun peek(): State {
         if (totalSteps == 0) {
-            return statesMap.keys.single().also { lastPeekedState = it }
+            return statesMap.keys.single()
         }
 
         val predictedState = predict()
-        lastPeekedState = predictedState
-        val wrapper = checkNotNull(statesMap[lastPeekedState])
+        val wrapper = checkNotNull(statesMap[predictedState])
         touchedBlocks.add(wrapper.currentBlock)
         touchedStates.add(wrapper)
 

diff --git a/usvm-core/src/main/kotlin/org/usvm/statistics/StepsStatistics.kt b/usvm-core/src/main/kotlin/org/usvm/statistics/StepsStatistics.kt
@@ -13,13 +13,19 @@ class StepsStatistics<Method, State : UState<*, Method, *, *, *, State>> : UMach
     var totalSteps = 0UL
         private set
 
+    var lastPeekedState: State? = null
+
     private val methodSteps = mutableMapOf<Method, ULong>()
 
     /**
      * Returns number of steps machine made during [method] exploration.
      */
     fun getMethodSteps(method: Method) = methodSteps.getOrDefault(method, 0UL)
 
+    override fun onStatePeeked(state: State) {
+        lastPeekedState = state
+    }
+
     override fun onState(parent: State, forks: Sequence<State>) {
         totalSteps++
         parent.callStack.forEach { (m, _) ->