Refactoring (optimization) (#5)

* Add steps statistics to path selector

* Remove redundant `runBlocking`

* Add observing for coverage zone additions

* Update StateWrapper manually

* Change search of `blockOf` to binary search

* Fix not adding initial methods to coverage zone

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

@@ -3,28 +3,27 @@ import org.usvm.UPathSelector
 import org.usvm.UState
 import org.usvm.statistics.BasicBlock
 import org.usvm.statistics.BlockGraph
+import org.usvm.statistics.StepsStatistics
+import org.usvm.util.Predictor
 import org.usvm.utils.Game
 import org.usvm.utils.StateWrapper
-import org.usvm.util.Predictor
 
 
 class AIPathSelector<Statement, State, Block>(
-    private val isInCoverageZone: (Block) -> Boolean,
     private val blockGraph: BlockGraph<*, Block, Statement>,
+    private val stepsStatistics: StepsStatistics<*, State>,
     private val predictor: Predictor<Game<Block>>,
 ) : UPathSelector<State> where
 State : UState<*, *, Statement , *, *, State>,
 Block : BasicBlock {
     private val statesMap = mutableMapOf<State, StateWrapper<Statement, State, Block>>()
     private var lastPeekedState: State? = null
-    private var totalSteps = 0
+    private val totalSteps
+        get() = stepsStatistics.totalSteps.toInt()
 
     private fun predict(): State {
         val wrappers = statesMap.values
         val vertices = blockGraph.blocks
-        vertices.forEach{
-            it.inCoverageZone = isInCoverageZone(it)
-        }
         val predictedId = predictor.predictState(Game(vertices, wrappers, blockGraph))
         val predictedState = statesMap.keys.find { it.id == predictedId }
 
@@ -34,8 +33,6 @@ Block : BasicBlock {
     override fun isEmpty() = statesMap.isEmpty()
 
     override fun peek(): State {
-        totalSteps++
-
         if (statesMap.size == 1) {
             return statesMap.keys.single().also { lastPeekedState = it }
         }
@@ -53,7 +50,7 @@ Block : BasicBlock {
             parent?.children?.remove(wrapper)
 
             state.pathNode += state.currentStatement
-            wrapper.updateBlock(totalSteps)
+            wrapper.update(totalSteps)
             wrapper.children.clear()
         }
     }
@@ -68,8 +65,7 @@ Block : BasicBlock {
                 state,
                 parentPathConditionSize,
                 parentHistory,
-                blockGraph,
-                totalSteps
+                blockGraph
             )
             //
             statesMap[state] = wrapper

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

@@ -15,6 +15,7 @@ import org.usvm.statistics.ApplicationGraph
 import org.usvm.statistics.BasicBlock
 import org.usvm.statistics.BlockGraph
 import org.usvm.statistics.CoverageStatistics
+import org.usvm.statistics.StepsStatistics
 import org.usvm.statistics.TimeStatistics
 import org.usvm.statistics.distances.CallGraphStatistics
 import org.usvm.statistics.distances.CallStackDistanceCalculator
@@ -40,6 +41,7 @@ private fun <Method, Statement, Target, State, Block> createPathSelector(
     coverageStatisticsFactory: () -> CoverageStatistics<Method, Statement, State>? = { null },
     cfgStatisticsFactory: () -> CfgStatistics<Method, Statement>? = { null },
     callGraphStatisticsFactory: () -> CallGraphStatistics<Method>? = { null },
+    stepsStatisticsFactory: () -> StepsStatistics<Method, State>? = { null },
     loopStatisticFactory: () -> StateLoopTracker<*, Statement, State>? = { null },
 ): UPathSelector<State>
     where Target : UTarget<Statement, Target>,
@@ -111,6 +113,7 @@ private fun <Method, Statement, Target, State, Block> createPathSelector(
 
             PathSelectionStrategy.AI -> createAIPathSelector(
                 requireNotNull(coverageStatisticsFactory()) { "Coverage statistics is required for GNN path selector" },
+                requireNotNull(stepsStatisticsFactory()) { "Steps statistics is required for GNN path selector" },
                 blockGraph,
                 options
             )
@@ -154,7 +157,7 @@ private fun <Method, Statement, Target, State, Block> createPathSelector(
         }
 
         PathSelectorCombinationStrategy.SEQUENTIAL -> {
-            val selector = SequentialPathSelector(selectors, options.stepsToStart)
+            val selector = SequentialPathSelector(selectors, options.stepsToStart, requireNotNull(stepsStatisticsFactory()))
 
             val mergingSelector = createMergingPathSelector(initialStates, selector, options, cfgStatisticsFactory)
             val resultSelector = mergingSelector.wrapIfRequired(options, loopStatisticFactory)
@@ -171,20 +174,34 @@ private fun <Method, Statement, Target, State, Block> createPathSelector(
 @Suppress("UNCHECKED_CAST")
 fun <Method, Statement, State, Block> createAIPathSelector(
     coverageStatistics: CoverageStatistics<Method, Statement, State>,
+    stepsStatistics: StepsStatistics<Method, State>,
     blockGraph: BlockGraph<Method, Block, Statement>,
     options: UMachineOptions
 ): UPathSelector<State> where State : UState<*, Method, Statement, *, *, State>,
                               Block : BasicBlock {
-    coverageStatistics.addOnCoveredObserver { _, _, statement ->
-        blockGraph.blockOf(statement).coveredByTest = true
+    with(coverageStatistics) {
+        addOnCoveredObserver { _, _, statement ->
+            blockGraph.blockOf(statement).coveredByTest = true
+        }
+        fun addBlocksToCoverageZone(method: Method) {
+            val blocksOfMethod = blockGraph.blocks.filter { blockGraph.methodOf(it) == method }
+            blocksOfMethod.forEach { it.inCoverageZone = true }
+        }
+        // initial methods
+        coverageZone.forEach(::addBlocksToCoverageZone)
+        addOnMethodAddedObserver { method ->
+            addBlocksToCoverageZone(method)
+        }
     }
 
-    val isInCoverageZone = { block: Block -> blockGraph.methodOf(block) in coverageStatistics.coverageZone }
-
     val predictor = options.oracle as? Predictor<Game<Block>>
     checkNotNull(predictor)
 
-    return AIPathSelector(isInCoverageZone, blockGraph, predictor)
+    return AIPathSelector(
+        blockGraph,
+        stepsStatistics,
+        predictor
+    )
 }
 
 fun <Method, Statement, Target, State, Block> createPathSelector(
@@ -195,6 +212,7 @@ fun <Method, Statement, Target, State, Block> createPathSelector(
     coverageStatisticsFactory: () -> CoverageStatistics<Method, Statement, State>? = { null },
     cfgStatisticsFactory: () -> CfgStatistics<Method, Statement>? = { null },
     callGraphStatisticsFactory: () -> CallGraphStatistics<Method>? = { null },
+    stepsStatisticsFactory: () -> StepsStatistics<Method, State>? = { null },
     loopStatisticFactory: () -> StateLoopTracker<*, Statement, State>? = { null },
 ): UPathSelector<State> where Target : UTarget<Statement, Target>, State : UState<*, Method, Statement, *, Target, State>,
                               Block : BasicBlock =
@@ -206,6 +224,7 @@ fun <Method, Statement, Target, State, Block> createPathSelector(
         coverageStatisticsFactory,
         cfgStatisticsFactory,
         callGraphStatisticsFactory,
+        stepsStatisticsFactory,
         loopStatisticFactory
     )
 
@@ -218,6 +237,7 @@ fun <Method, Statement, Target, State, Block> createPathSelector(
     coverageStatisticsFactory: () -> CoverageStatistics<Method, Statement, State>? = { null },
     cfgStatisticsFactory: () -> CfgStatistics<Method, Statement>? = { null },
     callGraphStatisticsFactory: () -> CallGraphStatistics<Method>? = { null },
+    stepsStatisticsFactory: () -> StepsStatistics<Method, State>? = { null },
     loopStatisticFactory: () -> StateLoopTracker<*, Statement, State>? = { null },
 ): UPathSelector<State> where Target : UTarget<Statement, Target>, State : UState<*, Method, Statement, *, Target, State>,
                               Block : BasicBlock {
@@ -230,6 +250,7 @@ fun <Method, Statement, Target, State, Block> createPathSelector(
             coverageStatisticsFactory,
             cfgStatisticsFactory,
             callGraphStatisticsFactory,
+            stepsStatisticsFactory,
             loopStatisticFactory
         )
     }
@@ -248,6 +269,7 @@ fun <Method, Statement, Target, State, Block> createPathSelector(
             coverageStatisticsFactory,
             cfgStatisticsFactory,
             callGraphStatisticsFactory,
+            stepsStatisticsFactory,
             loopStatisticFactory
         )
 

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

@@ -1,17 +1,20 @@
 package org.usvm.ps
 
 import org.usvm.UPathSelector
+import org.usvm.statistics.StepsStatistics
 
 class SequentialPathSelector<State>(
     private val selectors: List<UPathSelector<State>>,
-    private val stepsToSwitch: UInt
+    private val stepsToSwitch: UInt,
+    private val stepsStatistics: StepsStatistics<*, *>
 ): UPathSelector<State> {
     init {
         require(selectors.size >= 2) { "Cannot create sequential path selector from less than 2 selectors" }
     }
 
     private var currentSelector = selectors.first()
-    private var totalSteps = 0u
+    private val totalSteps
+        get() = stepsStatistics.totalSteps.toUInt()
 
     override fun isEmpty() = currentSelector.isEmpty() && selectors.all { it.isEmpty() }
 
@@ -20,7 +23,6 @@ class SequentialPathSelector<State>(
             selectors.removeFirst()
             currentSelector = selectors.first()
         }
-        totalSteps++
         return currentSelector.peek()
     }
 

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

@@ -1,7 +1,7 @@
 package org.usvm.statistics
 
 interface BlockGraph<Method, Block, Statement> {
-    val blocks: Set<Block>
+    val blocks: List<Block>
 
     fun predecessors(node: Block): Sequence<Block>
     fun successors(node: Block): Sequence<Block>

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

@@ -19,6 +19,7 @@ class CoverageStatistics<Method, Statement, State : UState<*, Method, Statement,
 ) : UMachineObserver<State> {
 
     private val onStatementCoveredObservers: MutableSet<(State, Method, Statement) -> Unit> = ConcurrentHashMap.newKeySet()
+    private val onMethodAddedToCoverageZoneObservers: MutableSet<(Method) -> Unit> = ConcurrentHashMap.newKeySet()
 
     // Set is actually concurrent
     private val uncoveredStatements = HashMap<Method, MutableSet<Statement>>()
@@ -56,6 +57,7 @@ class CoverageStatistics<Method, Statement, State : UState<*, Method, Statement,
         uncoveredStatements[method] = methodStatements
         totalUncoveredStatements += methodStatements.size
         coveredStatements[method] = hashSetOf()
+        onMethodAddedToCoverageZoneObservers.forEach { it(method) }
     }
 
     private fun computeCoverage(covered: Int, uncovered: Int): Float {
@@ -119,6 +121,10 @@ class CoverageStatistics<Method, Statement, State : UState<*, Method, Statement,
         onStatementCoveredObservers.add(observer)
     }
 
+    fun addOnMethodAddedObserver(observer: (Method) -> Unit) {
+        onMethodAddedToCoverageZoneObservers.add(observer)
+    }
+
     // TODO: don't consider coverage of runtime exceptions states
     override fun onStateTerminated(state: State, stateReachable: Boolean) {
         if (!stateReachable) return

usvm-core/src/main/kotlin/org/usvm/utils/StateWrapper.kt

@@ -11,45 +11,39 @@ data class StateHistoryElement(
 )
 
 class StateWrapper<Statement, State, Block>(
-    // TODO: too many parameters
     private val state: State,
     private val parentPathConditionSize: Int,
     private val parentHistory: MutableMap<Block, StateHistoryElement>,
     private val blockGraph: BlockGraph<*, Block, Statement>,
-    steps: Int
 ) where State : UState<*, *, Statement, *, *, State>, Block : BasicBlock {
     val children = mutableSetOf<StateWrapper<Statement, State, Block>>()
-    val history = parentHistory
+    val history = parentHistory.toMutableMap()
     val id = state.id
 
-    val position: Int
-        get() = currentBlock.id
-
-    val pathConditionSize: Int
-        get() = parentPathConditionSize + state.forkPoints.depth
-
-    val visitedAgainVertices: Int
-        get() = history.values.count { it.numOfVisits > 1 }
-
     val visitedNotCoveredVerticesInZone: Int
         get() = history.keys.count { !it.coveredByTest && it.inCoverageZone }
 
     val visitedNotCoveredVerticesOutOfZone: Int
         get() = history.keys.count { !it.coveredByTest && !it.inCoverageZone }
 
-    val visitedStatement: Statement
-        get() = checkNotNull(state.pathNode.parent?.statement)
-
-    var stepWhenMovedLastTime = steps
-
-    val instructionsVisitedInCurrentBlock: Int
-        get() = blockGraph.statementsOf(currentBlock).indexOf(visitedStatement) + 1
-
-    val currentBlock: Block
-        get() = blockGraph.blockOf(visitedStatement)
+    var visitedStatement: Statement? = null
+    lateinit var currentBlock: Block
+    var position: Int = 0
+    var pathConditionSize: Int = 0
+    var visitedAgainVertices: Int = 0
+    var instructionsVisitedInCurrentBlock: Int = 0
+    var stepWhenMovedLastTime: Int = 0
 
     fun update(steps: Int) {
+        visitedStatement = checkNotNull(state.pathNode.parent?.statement)
+        currentBlock = blockGraph.blockOf(visitedStatement!!)
+
+        position = currentBlock.id
+        pathConditionSize = parentPathConditionSize + state.forkPoints.depth
+        visitedAgainVertices = history.values.count { it.numOfVisits > 1 }
+        instructionsVisitedInCurrentBlock = blockGraph.statementsOf(currentBlock).indexOf(visitedStatement) + 1
         stepWhenMovedLastTime = steps
+
         updateBlock(stepWhenMovedLastTime)
     }
 
@@ -67,6 +61,8 @@ class StateWrapper<Statement, State, Block>(
         }
     }
 
-    fun addChildren(wrappers: Collection<StateWrapper<Statement, State, Block>>) =
+    fun addChildren(wrappers: Collection<StateWrapper<Statement, State, Block>>) {
+        wrappers.forEach { it.update(stepWhenMovedLastTime) }
         children.addAll(wrappers)
+    }
 }