Flow state collectors

usvm-core/src/main/kotlin/org/usvm/statistics/collectors/CoveredNewStatesCollector.kt

@@ -11,12 +11,10 @@ import org.usvm.statistics.UMachineObserver
  * @param isException if true, state is collected regardless of coverage.
  */
 class CoveredNewStatesCollector<State>(
+    private val statesCollector: StatesCollector<State>,
     private val coverageStatistics: CoverageStatistics<*, *, *>,
     private val isException: (State) -> Boolean
-) : StatesCollector<State> {
-    private val mutableCollectedStates = mutableListOf<State>()
-    override val collectedStates: List<State> = mutableCollectedStates
-
+) : UMachineObserver<State> {
     private var previousCoveredStatements = coverageStatistics.getTotalCoveredStatements()
 
     override fun onStateTerminated(state: State, stateReachable: Boolean) {
@@ -25,14 +23,14 @@ class CoveredNewStatesCollector<State>(
         }
 
         if (isException(state)) {
-            mutableCollectedStates.add(state)
+            statesCollector.addState(state)
             return
         }
 
         val currentCoveredStatements = coverageStatistics.getTotalCoveredStatements()
         if (currentCoveredStatements > previousCoveredStatements) {
             previousCoveredStatements = currentCoveredStatements
-            mutableCollectedStates += state
+            statesCollector.addState(state)
         }
     }
 }

usvm-core/src/main/kotlin/org/usvm/statistics/collectors/StatesCollector.kt

@@ -1,14 +1,44 @@
 package org.usvm.statistics.collectors
 
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.runBlocking
 import org.usvm.statistics.UMachineObserver
 
 /**
  * Interface for [UMachineObserver]s which are able to
  * collect states.
  */
-interface StatesCollector<State> : UMachineObserver<State> {
+interface StatesCollector<State> {
+    val count: Int
+    fun addState(state: State)
+}
+
+class FlowStatesCollector<State> : StatesCollector<State> {
+    private var statesCount: Int = 0
+    override val count: Int
+        get() = statesCount
+
+    private val flow = MutableSharedFlow<State>()
+    val collectedStatesFlow: Flow<State> = flow
+
+    override fun addState(state: State) = runBlocking {
+        statesCount++
+        flow.emit(state)
+    }
+}
+
+class ListStatesCollector<State> : StatesCollector<State> {
     /**
      * Current list of collected states.
      */
-    val collectedStates: List<State>
+    private val states = mutableListOf<State>()
+    val collectedStates: List<State> = states
+
+    override val count: Int
+        get() = states.size
+
+    override fun addState(state: State) {
+        states.add(state)
+    }
 }

usvm-core/src/main/kotlin/org/usvm/statistics/collectors/TargetsReachedStatesCollector.kt

@@ -1,20 +1,20 @@
 package org.usvm.statistics.collectors
 
 import org.usvm.UState
+import org.usvm.statistics.UMachineObserver
 
 /**
  * [StatesCollector] implementation collecting only those states which have reached
  * any terminal targets.
  */
-class TargetsReachedStatesCollector<State : UState<*, *, *, *, *, State>> : StatesCollector<State> {
-    private val mutableCollectedStates = mutableListOf<State>()
-    override val collectedStates: List<State> = mutableCollectedStates
-
+class TargetsReachedStatesCollector<State : UState<*, *, *, *, *, State>>(
+    val statesCollector: StatesCollector<State>
+) : UMachineObserver<State> {
     // TODO probably this should be called not only for terminated states
     //      Also, we should process more carefully clone operation for the states
     override fun onStateTerminated(state: State, stateReachable: Boolean) {
         if (state.targets.reachedTerminal.isNotEmpty()) {
-            mutableCollectedStates.add(state)
+            statesCollector.addState(state)
         }
     }
 }

usvm-jvm-instrumentation/src/main/kotlin/org/usvm/instrumentation/util/Classloader.kt

@@ -1,38 +1,34 @@
 @file:Suppress("JAVA_MODULE_DOES_NOT_EXPORT_PACKAGE")
 package org.usvm.instrumentation.util
-import jdk.internal.loader.URLClassPath
-import jdk.internal.loader.Resource as InternalResource
 
 import java.io.File
-import java.io.InputStream
 import java.net.URL
-import java.security.AccessController
+import java.net.URLClassLoader
 import java.security.CodeSigner
 
-class URLClassPathLoader(classPath: List<File>) {
+class URLClassPathLoader(private val classPath: List<File>) {
 
     interface Resource {
         fun getName(): String
         fun getURL(): URL
         fun getCodeSourceURL(): URL
-        fun getInputStream(): InputStream
-        fun getContentLength(): Int
         fun getBytes(): ByteArray
         fun getCodeSigners(): Array<CodeSigner>?
     }
 
-    private class InternalResourceWrapper(val resource: InternalResource): Resource {
-        override fun getName(): String = resource.name
-        override fun getURL(): URL = resource.url
-        override fun getCodeSourceURL(): URL = resource.codeSourceURL
-        override fun getInputStream(): InputStream = resource.inputStream
-        override fun getContentLength(): Int = resource.contentLength
-        override fun getBytes(): ByteArray = resource.bytes
-        override fun getCodeSigners(): Array<CodeSigner>? = resource.codeSigners
-    }
-
-    private val urlClassPath = URLClassPath(classPath.map { it.toURI().toURL() }.toTypedArray(), AccessController.getContext())
+    private val urlClassLoader = URLClassLoader(classPath.map { it.toURI().toURL() }.toTypedArray())
 
-    fun getResource(name: String): Resource = InternalResourceWrapper(urlClassPath.getResource(name, false))
+    fun getResource(name: String): Resource {
+        val resourceUrl = urlClassLoader.getResource(name) ?: error("Resource $name not found on classpath $classPath")
+        return object : Resource {
+            override fun getName(): String = name
+            override fun getURL(): URL = resourceUrl
+            // TODO usvm-sbft-merge: may be incorrect, especially for non-ASCII URLs
+            override fun getCodeSourceURL(): URL = resourceUrl
+            override fun getBytes(): ByteArray = resourceUrl.readBytes()
+            // TODO usvm-sbft-merge: figure out the way to get code signers
+            override fun getCodeSigners(): Array<CodeSigner>? = null
+        }
+    }
 
 }

usvm-jvm/src/main/kotlin/org/usvm/machine/JcMachine.kt

@@ -26,6 +26,8 @@ import org.usvm.statistics.TimeStatistics
 import org.usvm.statistics.TransitiveCoverageZoneObserver
 import org.usvm.statistics.UMachineObserver
 import org.usvm.statistics.collectors.CoveredNewStatesCollector
+import org.usvm.statistics.collectors.ListStatesCollector
+import org.usvm.statistics.collectors.StatesCollector
 import org.usvm.statistics.collectors.TargetsReachedStatesCollector
 import org.usvm.statistics.constraints.SoftConstraintsObserver
 import org.usvm.statistics.distances.CfgStatistics
@@ -54,7 +56,11 @@ class JcMachine(
 
     private val cfgStatistics = CfgStatisticsImpl(applicationGraph)
 
-    fun analyze(methods: List<JcMethod>, targets: List<JcTarget> = emptyList()): List<JcState> {
+    fun analyze(
+        methods: List<JcMethod>,
+        statesCollector: StatesCollector<JcState>,
+        targets: List<JcTarget> = emptyList()
+    ) {
         logger.debug("{}.analyze({})", this, methods)
         val initialStates = mutableMapOf<JcMethod, JcState>()
         methods.forEach {
@@ -103,13 +109,13 @@ class JcMachine(
             { callGraphStatistics }
         )
 
-        val statesCollector =
+        val statesCollectorObs =
             when (options.stateCollectionStrategy) {
-                StateCollectionStrategy.COVERED_NEW -> CoveredNewStatesCollector<JcState>(coverageStatistics) {
+                StateCollectionStrategy.COVERED_NEW -> CoveredNewStatesCollector(statesCollector, coverageStatistics) {
                     it.methodResult is JcMethodResult.JcException
                 }
 
-                StateCollectionStrategy.REACHED_TARGET -> TargetsReachedStatesCollector()
+                StateCollectionStrategy.REACHED_TARGET -> TargetsReachedStatesCollector(statesCollector)
             }
 
         val stepsStatistics = StepsStatistics<JcMethod, JcState>()
@@ -120,7 +126,7 @@ class JcMachine(
             timeStatisticsFactory = { timeStatistics },
             stepsStatisticsFactory = { stepsStatistics },
             coverageStatisticsFactory = { coverageStatistics },
-            getCollectedStatesCount = { statesCollector.collectedStates.size }
+            getCollectedStatesCount = { statesCollector.count }
         )
 
         val observers = mutableListOf<UMachineObserver<JcState>>(coverageStatistics)
@@ -148,7 +154,7 @@ class JcMachine(
                 )
             )
         }
-        observers.add(statesCollector)
+        observers.add(statesCollectorObs)
         // TODO: use the same calculator which is used for path selector
         if (targets.isNotEmpty()) {
             val distanceCalculator = MultiTargetDistanceCalculator<JcMethod, JcInst, InterprocDistance> { stmt ->
@@ -189,12 +195,16 @@ class JcMachine(
             isStateTerminated = ::isStateTerminated,
             stopStrategy = stopStrategy,
         )
-
-        return statesCollector.collectedStates
     }
 
-    fun analyze(method: JcMethod, targets: List<JcTarget> = emptyList()): List<JcState> =
-        analyze(listOf(method), targets)
+    fun analyze(method: JcMethod, collector: StatesCollector<JcState>, targets: List<JcTarget> = emptyList()) =
+        analyze(listOf(method), collector, targets)
+
+    fun analyze(method: JcMethod, targets: List<JcTarget> = emptyList()): List<JcState> {
+        val collector = ListStatesCollector<JcState>()
+        analyze(method, collector, targets)
+        return collector.collectedStates
+    }
 
     /**
      * Returns a wrapper for the [cfgStatistics] that ignores [JcTransparentInstruction]s.

usvm-sample-language/src/main/kotlin/org/usvm/machine/SampleMachine.kt

@@ -16,6 +16,8 @@ import org.usvm.statistics.StepsStatistics
 import org.usvm.statistics.TimeStatistics
 import org.usvm.statistics.UMachineObserver
 import org.usvm.statistics.collectors.CoveredNewStatesCollector
+import org.usvm.statistics.collectors.ListStatesCollector
+import org.usvm.statistics.collectors.StatesCollector
 import org.usvm.statistics.collectors.TargetsReachedStatesCollector
 import org.usvm.statistics.constraints.SoftConstraintsObserver
 import org.usvm.statistics.distances.CallGraphStatisticsImpl
@@ -44,8 +46,9 @@ class SampleMachine(
 
     fun analyze(
         methods: List<Method<*>>,
+        resultCollector: StatesCollector<ProgramExecutionResult>,
         targets: List<SampleTarget> = emptyList()
-    ): Collection<ProgramExecutionResult> {
+    ) {
         logger.debug("{}.analyze({})", this, methods)
         val initialStates = mutableMapOf<Method<*>, SampleState>()
         methods.forEach {
@@ -75,12 +78,21 @@ class SampleMachine(
             { callGraphStatistics }
         )
 
+        val stateCollector = object : StatesCollector<SampleState> {
+            override val count: Int
+                get() = resultCollector.count
+
+            override fun addState(state: SampleState) {
+                resultCollector.addState(resultModelConverter.convert(state, state.entrypoint))
+            }
+        }
+
         val statesCollector =
             when (options.stateCollectionStrategy) {
-                StateCollectionStrategy.COVERED_NEW -> CoveredNewStatesCollector<SampleState>(coverageStatistics) {
+                StateCollectionStrategy.COVERED_NEW -> CoveredNewStatesCollector(stateCollector, coverageStatistics) {
                     it.exceptionRegister != null
                 }
-                StateCollectionStrategy.REACHED_TARGET -> TargetsReachedStatesCollector()
+                StateCollectionStrategy.REACHED_TARGET -> TargetsReachedStatesCollector(stateCollector)
             }
 
         val stepsStatistics = StepsStatistics<Method<*>, SampleState>()
@@ -91,7 +103,7 @@ class SampleMachine(
             { timeStatistics },
             { stepsStatistics },
             { coverageStatistics },
-            { statesCollector.collectedStates.size }
+            { stateCollector.count }
         )
 
         val observers = mutableListOf<UMachineObserver<SampleState>>(coverageStatistics)
@@ -110,12 +122,19 @@ class SampleMachine(
             isStateTerminated = ::isStateTerminated,
             stopStrategy = stopStrategy,
         )
-
-        return statesCollector.collectedStates.map { resultModelConverter.convert(it, it.entrypoint) }
     }
 
-    fun analyze(method: Method<*>, targets: List<SampleTarget> = emptyList()): Collection<ProgramExecutionResult> =
-        analyze(listOf(method), targets)
+    fun analyze(
+        method: Method<*>,
+        resultCollector: StatesCollector<ProgramExecutionResult>,
+        targets: List<SampleTarget> = emptyList()
+    ) = analyze(listOf(method), resultCollector, targets)
+
+    fun analyze(method: Method<*>, targets: List<SampleTarget> = emptyList()): List<ProgramExecutionResult> {
+        val collector = ListStatesCollector<ProgramExecutionResult>()
+        analyze(method, collector, targets)
+        return collector.collectedStates
+    }
 
     private fun getInitialState(
         method: Method<*>,