Create ApplicationBlockGraph interface, jvm impl

Author: emnigma

.gitignore

@@ -8,3 +8,6 @@ buildSrc/.gradle
 
 # Ignore Idea directory
 .idea
+
+# Ignore MacOS-specific files
+/**/.DS_Store

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

@@ -89,6 +89,10 @@ class CoverageStatistics<Method, Statement, State : UState<*, Method, Statement,
         return uncoveredStatements.values.flatten()
     }
 
+    fun isCovered(statement: Statement) = statement in coveredStatements.values.flatten()
+
+    fun inCoverageZone(method: Method) = coveredStatements.containsKey(method) || uncoveredStatements.containsKey(method)
+
     /**
      * Adds a listener triggered when a new statement is covered.
      */

usvm-path-selection/build.gradle

@@ -7,4 +7,5 @@ dependencies {
     implementation(project(":usvm-jvm"))
 
     implementation("org.jacodb:jacodb-analysis:${Versions.jcdb}")
-}
+    implementation("org.jacodb:jacodb-approximations:${Versions.jcdb}")
+}

usvm-path-selection/src/main/kotlin/org/usvm/ApplicationBlockGraph.kt

@@ -3,7 +3,7 @@ package org.usvm
 import org.usvm.statistics.ApplicationGraph
 
 interface ApplicationBlockGraph<Method, BasicBlock, Statement> : ApplicationGraph<Method, BasicBlock> {
-    fun block(stmt: Statement): BasicBlock
+    fun blockOf(stmt: Statement): BasicBlock
     fun instructions(block: BasicBlock): Sequence<Statement>
     fun blocks(): Sequence<BasicBlock>
 }

usvm-path-selection/src/main/kotlin/org/usvm/MLMachineOptions.kt

@@ -0,0 +1,25 @@
+package org.usvm
+
+enum class MLPathSelectionStrategy {
+    /**
+     * Collects features according to states selected by any other path selector.
+     */
+    FEATURES_LOGGING,
+
+    /**
+     * Collects features and feeds them to the ML model to select states.
+     * Extends FEATURE_LOGGING path selector.
+     */
+    MACHINE_LEARNING,
+
+    /**
+     * Selects states with best Graph Neural Network state score
+     */
+    GNN,
+}
+
+data class MLMachineOptions(
+    val basicOptions: UMachineOptions,
+    val pathSelectionStrategy: MLPathSelectionStrategy,
+    val heteroGNNModelPath: String = ""
+)

usvm-path-selection/src/main/kotlin/org/usvm/MLPathSelectorStrategy.kt

@@ -0,0 +1,53 @@
+package org.usvm
+
+import org.usvm.ps.ExceptionPropagationPathSelector
+import org.usvm.ps.GNNPathSelector
+import org.usvm.statistics.CoverageStatistics
+import org.usvm.statistics.StateVisitsStatistics
+
+fun <Method, Statement, BasicBlock, State : UState<*, Method, Statement, *, *, State>> createPathSelector(
+    initialState: State,
+    options: MLMachineOptions,
+    applicationGraph: ApplicationBlockGraph<Method, BasicBlock, Statement>,
+    stateVisitsStatistics: StateVisitsStatistics<Method, Statement, State>,
+    coverageStatistics: CoverageStatistics<Method, Statement, State>,
+): UPathSelector<State> {
+    val selector = when (options.pathSelectionStrategy) {
+        MLPathSelectionStrategy.GNN -> createGNNPathSelector(
+            stateVisitsStatistics,
+            coverageStatistics, applicationGraph, options.heteroGNNModelPath
+        )
+
+        else -> {
+            throw NotImplementedError()
+        }
+    }
+
+    val propagateExceptions = options.basicOptions.exceptionsPropagation
+
+    val resultSelector = selector.wrapIfRequired(propagateExceptions)
+    resultSelector.add(listOf(initialState))
+
+    return selector
+}
+
+private fun <State : UState<*, *, *, *, *, State>> UPathSelector<State>.wrapIfRequired(propagateExceptions: Boolean) =
+    if (propagateExceptions && this !is ExceptionPropagationPathSelector<State>) {
+        ExceptionPropagationPathSelector(this)
+    } else {
+        this
+    }
+
+private fun <Method, Statement, BasicBlock, State : UState<*, Method, Statement, *, *, State>> createGNNPathSelector(
+    stateVisitsStatistics: StateVisitsStatistics<Method, Statement, State>,
+    coverageStatistics: CoverageStatistics<Method, Statement, State>,
+    applicationGraph: ApplicationBlockGraph<Method, BasicBlock, Statement>,
+    heteroGNNModelPath: String,
+): UPathSelector<State> {
+    return GNNPathSelector(
+        coverageStatistics,
+        stateVisitsStatistics,
+        applicationGraph,
+        heteroGNNModelPath
+    )
+}

usvm-path-selection/src/main/kotlin/org/usvm/jvm/JcApplicationBlockGraph.kt

@@ -1,24 +1,42 @@
 package org.usvm.jvm
 
+import org.jacodb.api.JcClasspath
 import org.jacodb.api.JcMethod
+import org.jacodb.api.JcTypedMethod
 import org.jacodb.api.cfg.JcBasicBlock
 import org.jacodb.api.cfg.JcBlockGraph
-import org.jacodb.api.cfg.JcGraph
 import org.jacodb.api.cfg.JcInst
-import org.jacodb.impl.cfg.JcBlockGraphImpl
+import org.jacodb.api.ext.toType
 import org.usvm.ApplicationBlockGraph
 import org.usvm.machine.JcApplicationGraph
+import java.util.concurrent.ConcurrentHashMap
 
-class JcApplicationBlockGraph(jcGraph: JcGraph) :
+class JcApplicationBlockGraph(cp: JcClasspath) :
     ApplicationBlockGraph<JcMethod, JcBasicBlock, JcInst> {
-    private val jcBlockGraphImpl: JcBlockGraph = JcBlockGraphImpl(jcGraph)
-    private val jcApplicationGraph: JcApplicationGraph = JcApplicationGraph(jcGraph.classpath)
-    override fun predecessors(node: JcBasicBlock): Sequence<JcBasicBlock> =
-        jcBlockGraphImpl.predecessors(node).asSequence()
+    val jcApplicationGraph: JcApplicationGraph = JcApplicationGraph(cp)
+    var initialStatement: JcInst? = null
+    private fun getInitialStatement(): JcInst {
+        if (initialStatement == null) {
+            throw RuntimeException("initial statement not set")
+        }
+        return initialStatement!!
+    }
+
+    private fun getBlockGraph() = getInitialStatement().location.method.flowGraph().blockGraph()
 
-    override fun successors(node: JcBasicBlock): Sequence<JcBasicBlock> = jcBlockGraphImpl.successors(node).asSequence()
+    override fun predecessors(node: JcBasicBlock): Sequence<JcBasicBlock> {
+        val jcBlockGraphImpl: JcBlockGraph = getBlockGraph()
+        return jcBlockGraphImpl.predecessors(node).asSequence()
+    }
+
+    override fun successors(node: JcBasicBlock): Sequence<JcBasicBlock> {
+        val jcBlockGraphImpl: JcBlockGraph = getBlockGraph()
+        return jcBlockGraphImpl.successors(node).asSequence()
+    }
 
     override fun callees(node: JcBasicBlock): Sequence<JcMethod> {
+        val jcBlockGraphImpl: JcBlockGraph = getBlockGraph()
+
         return jcBlockGraphImpl.instructions(node)
             .map { jcApplicationGraph.callees(it) }
             .reduce { acc, sequence -> acc + sequence }
@@ -29,26 +47,53 @@ class JcApplicationBlockGraph(jcGraph: JcGraph) :
     override fun callers(method: JcMethod): Sequence<JcBasicBlock> {
         return jcApplicationGraph
             .callers(method)
-            .map { stmt -> block(stmt) }
+            .map { stmt -> blockOf(stmt) }
             .toSet()
             .asSequence()
     }
 
-    override fun entryPoints(method: JcMethod): Sequence<JcBasicBlock> = jcBlockGraphImpl.entries.asSequence()
+    override fun entryPoints(method: JcMethod): Sequence<JcBasicBlock> =
+        method.flowGraph().blockGraph().entries.asSequence()
+
+    override fun exitPoints(method: JcMethod): Sequence<JcBasicBlock> =
+        method.flowGraph().blockGraph().exits.asSequence()
 
-    override fun exitPoints(method: JcMethod): Sequence<JcBasicBlock> = jcBlockGraphImpl.exits.asSequence()
+    override fun methodOf(node: JcBasicBlock): JcMethod {
+        val firstInstruction = getBlockGraph().instructions(node).first()
+        return jcApplicationGraph.methodOf(firstInstruction)
+    }
 
-    override fun methodOf(node: JcBasicBlock): JcMethod = jcBlockGraphImpl.instructions(node).first().location.method
+    override fun instructions(block: JcBasicBlock): Sequence<JcInst> {
+        return getBlockGraph().instructions(block).asSequence()
+    }
 
     override fun statementsOf(method: JcMethod): Sequence<JcBasicBlock> {
         return jcApplicationGraph
             .statementsOf(method)
-            .map { stmt -> block(stmt) }
+            .map { stmt -> blockOf(stmt) }
             .toSet()
             .asSequence()
     }
 
-    override fun block(stmt: JcInst): JcBasicBlock = jcBlockGraphImpl.block(stmt)
-    override fun instructions(block: JcBasicBlock): Sequence<JcInst> = jcBlockGraphImpl.instructions(block).asSequence()
-    override fun blocks(): Sequence<JcBasicBlock> = jcBlockGraphImpl.asSequence()
+    override fun blockOf(stmt: JcInst): JcBasicBlock {
+        val jcBlockGraphImpl: JcBlockGraph = stmt.location.method.flowGraph().blockGraph()
+        val blocks = blocks()
+        for (block in blocks) {
+            if (stmt in jcBlockGraphImpl.instructions(block)) {
+                return block
+            }
+        }
+        throw IllegalStateException("block not found for $stmt in ${jcBlockGraphImpl.jcGraph.method}")
+    }
+
+    override fun blocks(): Sequence<JcBasicBlock> {
+        return getInitialStatement().location.method.flowGraph().blockGraph().asSequence()
+    }
+
+    private val typedMethodsCache = ConcurrentHashMap<JcMethod, JcTypedMethod>()
+
+    val JcMethod.typed
+        get() = typedMethodsCache.getOrPut(this) {
+            enclosingClass.toType().declaredMethods.first { it.method == this }
+        }
 }