Common ml base

.gitignore

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

buildSrc/src/main/kotlin/Versions.kt

@@ -8,6 +8,7 @@ object Versions {
     const val mockk = "1.13.4"
     const val junitParams = "5.9.3"
     const val logback = "1.4.8"
+    const val onnxruntime = "1.15.1"
 
     // versions for jvm samples
     const val samplesLombok = "1.18.20"

buildSrc/src/main/kotlin/usvm.kotlin-conventions.gradle.kts

@@ -19,6 +19,7 @@ dependencies {
 
     implementation(kotlin("stdlib-jdk8"))
     implementation(kotlin("reflect"))
+    implementation("com.microsoft.onnxruntime", "onnxruntime", Versions.onnxruntime)
     implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:${Versions.coroutines}")
 
     testImplementation(kotlin("test"))

settings.gradle.kts

@@ -4,3 +4,4 @@ include("usvm-core")
 include("usvm-jvm")
 include("usvm-util")
 include("usvm-sample-language")
+include("usvm-ml-path-selection")

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

@@ -34,6 +34,11 @@ sealed class PathsTrieNode<State : UState<*, *, Statement, *, *, State>, Stateme
      */
     abstract val depth: Int
 
+    /**
+     * States that were forked from this node
+     */
+    abstract val accumulatedForks: MutableCollection<State>
+
     /**
      * Adds a new label to [labels] collection.
      */
@@ -80,6 +85,7 @@ class PathsTrieNodeImpl<State : UState<*, *, Statement, *, *, State>, Statement>
         statement = statement
     ) {
         parentNode.children[statement] = this
+        parentNode.accumulatedForks.addAll(this.states)
     }
 
     internal constructor(parentNode: PathsTrieNodeImpl<State, Statement>, statement: Statement, state: State) : this(
@@ -89,11 +95,14 @@ class PathsTrieNodeImpl<State : UState<*, *, Statement, *, *, State>, Statement>
         statement = statement
     ) {
         parentNode.children[statement] = this
+        parentNode.accumulatedForks.addAll(this.states)
         parentNode.states -= state
     }
 
     override val labels: MutableSet<Any> = hashSetOf()
 
+    override val accumulatedForks: MutableCollection<State> = mutableSetOf()
+
     override fun addLabel(label: Any) {
         labels.add(label)
     }
@@ -115,6 +124,8 @@ class RootNode<State : UState<*, *, Statement, *, *, State>, Statement> : PathsT
 
     override val labels: MutableSet<Any> = hashSetOf()
 
+    override val accumulatedForks: MutableCollection<State> = mutableSetOf()
+
     override val depth: Int = 0
 
     override fun addLabel(label: Any) {

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-ml-path-selection/build.gradle

@@ -0,0 +1,11 @@
+plugins {
+    id("usvm.kotlin-conventions")
+}
+
+dependencies {
+    implementation(project(":usvm-core"))
+    implementation(project(":usvm-jvm"))
+
+    implementation("org.jacodb:jacodb-analysis:${Versions.jcdb}")
+    implementation("org.jacodb:jacodb-approximations:${Versions.jcdb}")
+}

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

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

usvm-ml-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-ml-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-ml-path-selection/src/main/kotlin/org/usvm/jvm/JcApplicationBlockGraph.kt

@@ -0,0 +1,100 @@
+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.JcInst
+import org.jacodb.api.ext.toType
+import org.usvm.ApplicationBlockGraph
+import org.usvm.machine.JcApplicationGraph
+import java.util.concurrent.ConcurrentHashMap
+
+class JcApplicationBlockGraph(cp: JcClasspath) :
+    ApplicationBlockGraph<JcMethod, JcBasicBlock, JcInst> {
+    val jcApplicationGraph: JcApplicationGraph = JcApplicationGraph(cp)
+    var initialStatement: JcInst? = null
+
+    private fun initialStatement(): JcInst {
+        if (initialStatement == null) {
+            throw RuntimeException("initial statement not set")
+        }
+        return initialStatement!!
+    }
+
+    private fun getBlockGraph() = initialStatement().location.method.flowGraph().blockGraph()
+
+    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() + jcBlockGraphImpl.throwers(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 }
+            .toSet()
+            .asSequence()
+    }
+
+    override fun callers(method: JcMethod): Sequence<JcBasicBlock> {
+        return jcApplicationGraph
+            .callers(method)
+            .map { stmt -> blockOf(stmt) }
+            .toSet()
+            .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 methodOf(node: JcBasicBlock): JcMethod {
+        val firstInstruction = getBlockGraph().instructions(node).first()
+        return jcApplicationGraph.methodOf(firstInstruction)
+    }
+
+    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 -> blockOf(stmt) }
+            .toSet()
+            .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 initialStatement().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 }
+        }
+}