Drain path selector queues without scheduling or merging

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

@@ -22,4 +22,17 @@ interface UPathSelector<State> {
      * Updates the internal priority of the [state].
      */
     fun update(state: State)
-}
+
+    /**
+     * Removes and yields every queued state after execution has stopped.
+     * Wrappers with scheduling or merging side effects must drain their children directly.
+     * Each state is removed before it is yielded; the sequence must be consumed once.
+     */
+    fun drain(): Sequence<State> = sequence {
+        while (!isEmpty()) {
+            val state = peek()
+            remove(state)
+            yield(state)
+        }
+    }
+}

usvm-core/src/main/kotlin/org/usvm/merging/MergingPathSelector.kt

@@ -75,6 +75,13 @@ class MergingPathSelector<State : UState<*, *, *, *, *, State>>(
         closeStatesSearcher.add(states)
     }
 
+    override fun drain(): Sequence<State> = sequence {
+        for (state in underlyingPathSelector.drain()) {
+            closeStatesSearcher.remove(state)
+            yield(state)
+        }
+    }
+
     override fun remove(state: State) {
         underlyingPathSelector.remove(state)
         closeStatesSearcher.remove(state)

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

@@ -2,6 +2,7 @@ package org.usvm.ps
 
 import org.usvm.UPathSelector
 import org.usvm.UState
+import java.util.Collections
 import java.util.IdentityHashMap
 
 /**
@@ -92,6 +93,26 @@ class ExceptionPropagationPathSelector<State : UState<*, *, *, *, *, State>>(
         selector.add(statesToAdd)
     }
 
+    override fun drain(): Sequence<State> = sequence {
+        val emittedStates = Collections.newSetFromMap(IdentityHashMap<State, Boolean>())
+        for (state in selector.drain()) {
+            statesInSelector.remove(state)
+            exceptionalStates.remove(state)
+            if (emittedStates.add(state)) {
+                yield(state)
+            }
+        }
+
+        while (exceptionalStates.isNotEmpty()) {
+            val state = exceptionalStates.keys.first()
+            exceptionalStates.remove(state)
+            statesInSelector.remove(state)
+            if (emittedStates.add(state)) {
+                yield(state)
+            }
+        }
+    }
+
     override fun remove(state: State) {
         if (state in exceptionalStates) {
             exceptionalStates -= state

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

@@ -1,6 +1,8 @@
 package org.usvm.ps
 
 import org.usvm.UPathSelector
+import java.util.Collections
+import java.util.IdentityHashMap
 
 /**
  * An interleaved path selector.
@@ -40,6 +42,17 @@ class InterleavedPathSelector<State>(
         ptr = (ptr + 1) % selectors.size
     }
 
+    override fun drain(): Sequence<State> = sequence {
+        val emittedStates = Collections.newSetFromMap(IdentityHashMap<State, Boolean>())
+        for (selector in selectors) {
+            for (state in selector.drain()) {
+                if (emittedStates.add(state)) {
+                    yield(state)
+                }
+            }
+        }
+    }
+
     override fun remove(state: State) {
         selectors.forEach { it.remove(state) }
     }

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

@@ -80,6 +80,24 @@ class IterativeDeepeningPs<Stmt, Method, Loop : Any, State : UState<*, Method, S
         underlyingPs.remove(state)
     }
 
+    override fun drain(): Sequence<State> = sequence {
+        yieldAll(underlyingPs.drain())
+
+        for (level in pendingStates.keys.sorted()) {
+            val states = pendingStates.getValue(level)
+            while (states.isNotEmpty()) {
+                yield(states.removeAt(states.lastIndex))
+            }
+            pendingStates.remove(level)
+        }
+
+        stateStats.clear()
+        lastPeekedState = null
+        lastPeekedStateStats = StateLoopStatistic.rootStats()
+        currentLevel = 0
+        maxLevel = 0
+    }
+
     private fun pushStates() {
         while (underlyingPs.isEmpty() && currentLevel <= maxLevel) {
             currentLevel++

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

@@ -25,6 +25,12 @@ abstract class KeyedPathSelector<State, Key>(private val getKey: (State) -> Key)
         pathSelectors[key]?.remove(state) ?: throw IllegalStateException("Trying to remove state with unknown key")
     }
 
+    override fun drain(): Sequence<State> = sequence {
+        for (selector in pathSelectors.values) {
+            yieldAll(selector.drain())
+        }
+    }
+
     fun removeKey(key: Key) {
         pathSelectors.remove(key)
     }

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

@@ -60,6 +60,14 @@ class LoopLimiterPs<Stmt, Method, Loop : Any, State : UState<*, Method, Stmt, *,
         addSingleState(state)
     }
 
+    override fun drain(): Sequence<State> = sequence {
+        yieldAll(underlyingPs.drain())
+
+        stateStats.clear()
+        lastPeekedState = null
+        lastPeekedStateStats = StateLoopStatistic.rootStats()
+    }
+
     override fun remove(state: State) {
         check(state === lastPeekedState) { "Try to remove not peeked state" }
 

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

@@ -1,6 +1,8 @@
 package org.usvm.ps
 
 import org.usvm.UPathSelector
+import java.util.Collections
+import java.util.IdentityHashMap
 
 /**
  * A parallel path selector.
@@ -40,6 +42,17 @@ class ParallelPathSelector<State>(
         ptr = (ptr + 1) % selectors.size
     }
 
+    override fun drain(): Sequence<State> = sequence {
+        val emittedStates = Collections.newSetFromMap(IdentityHashMap<State, Boolean>())
+        for (selector in selectors) {
+            for (state in selector.drain()) {
+                if (emittedStates.add(state)) {
+                    yield(state)
+                }
+            }
+        }
+    }
+
     override fun remove(state: State) {
         selectors[ptr].remove(state)
     }