ML path selection

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-ml-path-selection/README.md

@@ -0,0 +1,44 @@
+## Machine Learning Path Selector
+
+### Entry point
+
+To run tests with this path selector use `jarRunner.kt`. You can pass a path to a configuration json as the first argument. Gathered statistics will be put in a folder according to your configuration.
+
+### Config
+
+A config object is declared inside `MLConfig.kt`. A detailed description of all the options is listed below:
+
+- `gameEnvPath` - a path to a folder that contains trained models (`rnn_cell.onnx`, `gnn_model.onnx`, `actor_model.onnx`) and a blacklist of tests to be skipped (`blacklist.txt`), also some logs are saved to this folder
+- `dataPath` - a path to a folder to save all statistics into
+- `defaultAlgorithm` - an algorithm to use if a trained model is not found, must be one of: `BFS`, `ForkDepthRandom`
+- `postprocessing` - how actor model's outputs should be processed, must be one of: `Argmax` (choose an id of the maximum value), `Softmax` (sample from a distribution derived from the outputs via the softmax), `None` (sample from the outputs — only when they form a distribution)
+- `mode` - a mode for `jarRunner.kt`, must be one of: `Calculation` (to calculate statistics used to train models), `Aggregation` (to aggregate statistics for different tests into one file), `Both` (to both calculate statistics and aggregate them), `Test` (to test this path selector with different time limits and compare it to other path selectors)
+- `logFeatures` - whether to save statistics used to train models
+- `shuffleTests` - whether to shuffle tests before running (affects the tests being run if the `dataConsumption` option is less than 100)
+- `discounts` - time discounts used when testing path selectors
+- `inputShape` - an input shape of an actor model
+- `maxAttentionLength` - a maximum attention length of a PPO actor model
+- `useGnn` - whether to use a GNN model
+- `dataConsumption` - a percentage of tests to run
+- `hardTimeLimit` - a time limit for one test
+- `solverTimeLimit` - a time limit for one solver call
+- `maxConcurrency` - a maximum number of threads running different tests concurrently
+- `graphUpdate` - when to update block graph data, must be one of: `Once` (at the beginning of a test), `TestGeneration` (every time a new test is generated)
+- `logGraphFeatuers` - whether to save graph statistics used to train a GNN model to a dataset file
+- `gnnFeaturesCount` - a number of features that a GNN model returns
+- `useRnn` - whether to use an RNN model
+- `rnnStateShape` - a shape of an RNN state
+- `rnnFeaturesCount` - a number of features that an RNN model returns
+- `inputJars` - jars and their packages to run tests on
+
+### How to modify the metric
+
+To modify the metric you may change values of the `reward` property of the `ActionData` objects. They are written inside the property `path` of the `FeaturesLoggingPathSelector`. Currently, the metric is calculated in the `remove` method of the `FeaturesLoggingPathSelector`.
+
+### Training environment
+
+The training environment and its description are inside `environment.zip`.
+
+### "Modified" files
+
+Source files which names start with "Modified" are modified copies of files from other modules. They were modified to support this path selector.

usvm-ml-path-selection/build.gradle.kts

@@ -0,0 +1,22 @@
+object MLVersions {
+    const val serialization = "1.5.1"
+    const val onnxruntime = "1.15.1"
+    const val dotlin = "1.0.2"
+}
+
+plugins {
+    id("usvm.kotlin-conventions")
+    kotlin("plugin.serialization") version "1.8.21"
+}
+
+dependencies {
+    implementation(project(":usvm-jvm"))
+    implementation(project(":usvm-core"))
+
+    implementation("org.jacodb:jacodb-analysis:${Versions.jcdb}")
+    implementation("ch.qos.logback:logback-classic:${Versions.logback}")
+
+    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:${MLVersions.serialization}")
+    implementation("io.github.rchowell:dotlin:${MLVersions.dotlin}")
+    implementation("com.microsoft.onnxruntime:onnxruntime:${MLVersions.onnxruntime}")
+}

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

@@ -0,0 +1,74 @@
+package org.usvm
+
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.buildJsonObject
+import kotlinx.serialization.json.put
+import kotlinx.serialization.json.putJsonObject
+import java.util.concurrent.ConcurrentHashMap
+
+class CoverageCounter {
+    private val testCoverages = ConcurrentHashMap<String, List<Float>>()
+    private val testStatementsCounts = ConcurrentHashMap<String, Float>()
+    private val testDiscounts = ConcurrentHashMap<String, List<Float>>()
+    private val testFinished = ConcurrentHashMap<String, Boolean>()
+
+    fun addTest(testName: String, statementsCount: Float) {
+        testCoverages[testName] = List(MLConfig.discounts.size) { 0.0f }
+        testStatementsCounts[testName] = statementsCount
+        testDiscounts[testName] = List(MLConfig.discounts.size) { 1.0f }
+        testFinished[testName] = false
+    }
+
+    fun updateDiscounts(testName: String) {
+        testDiscounts[testName] = testDiscounts.getValue(testName)
+            .mapIndexed { id, currentDiscount -> MLConfig.discounts[id] * currentDiscount }
+    }
+
+    fun updateResults(testName: String, newCoverage: Float) {
+        val currentDiscounts = testDiscounts.getValue(testName)
+        testCoverages[testName] = testCoverages.getValue(testName)
+            .mapIndexed { id, currentCoverage -> currentCoverage + currentDiscounts[id] * newCoverage }
+    }
+
+    fun finishTest(testName: String) {
+        testFinished[testName] = true
+    }
+
+    fun reset() {
+        testCoverages.clear()
+        testStatementsCounts.clear()
+        testDiscounts.clear()
+        testFinished.clear()
+    }
+
+    private fun getTotalCoverages(): List<Float> {
+        return testCoverages.values.reduce { acc, floats ->
+            acc.zip(floats).map { (total, value) -> total + value }
+        }
+    }
+
+    fun getStatistics(): JsonObject {
+        return buildJsonObject {
+            putJsonObject("Tests") {
+                testCoverages.forEach { (test, coverages) ->
+                    putJsonObject(test) {
+                        putJsonObject("discounts") {
+                            MLConfig.discounts.zip(coverages).forEach { (discount, coverage) ->
+                                put(discount.toString(), coverage)
+                            }
+                        }
+                        put("statementsCount", testStatementsCounts[test])
+                        put("finished", testFinished[test])
+                    }
+                }
+            }
+            putJsonObject("totalDiscounts") {
+                MLConfig.discounts.zip(getTotalCoverages()).forEach { (discount, coverage) ->
+                    put(discount.toString(), coverage)
+                }
+            }
+            put("totalStatementsCount", testStatementsCounts.values.sum())
+            put("finishedTestsCount", testFinished.values.sumOf { if (it) 1.0 else 0.0 })
+        }
+    }
+}

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

@@ -0,0 +1,51 @@
+package org.usvm
+
+enum class Postprocessing {
+    Argmax,
+    Softmax,
+    None,
+}
+
+enum class Mode {
+    Calculation,
+    Aggregation,
+    Both,
+    Test,
+}
+
+enum class Algorithm {
+    BFS,
+    ForkDepthRandom,
+}
+
+enum class GraphUpdate {
+    Once,
+    TestGeneration,
+}
+
+object MLConfig {
+    var gameEnvPath = "../Game_env"
+    var dataPath = "../Data"
+    var defaultAlgorithm = Algorithm.BFS
+    var postprocessing = Postprocessing.Argmax
+    var mode = Mode.Both
+    var logFeatures = true
+    var shuffleTests = true
+    var discounts = listOf(1.0f, 0.998f, 0.99f)
+    var inputShape = listOf<Long>(1, -1, 77)
+    var maxAttentionLength = -1
+    var useGnn = true
+    var dataConsumption = 100.0f
+    var hardTimeLimit = 30000 // in ms
+    var solverTimeLimit = 10000 // in ms
+    var maxConcurrency = 64
+    var graphUpdate = GraphUpdate.Once
+    var logGraphFeatures = false
+    var gnnFeaturesCount = 8
+    var useRnn = true
+    var rnnStateShape = listOf<Long>(4, 1, 512)
+    var rnnFeaturesCount = 33
+    var inputJars = mapOf(
+        Pair("../Game_env/jars/usvm-jvm-new.jar", listOf("org.usvm.samples", "com.thealgorithms"))
+    ) // path to jar file -> list of package names
+}

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

@@ -0,0 +1,67 @@
+package org.usvm
+
+enum class ModifiedPathSelectionStrategy {
+    /**
+     * 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,
+}
+
+data class ModifiedUMachineOptions(
+    /**
+     * State selection heuristics.
+     * If multiple heuristics are specified, they are combined according to [pathSelectorCombinationStrategy].
+     *
+     * @see PathSelectionStrategy
+     */
+    val pathSelectionStrategies: List<ModifiedPathSelectionStrategy> = listOf(ModifiedPathSelectionStrategy.MACHINE_LEARNING),
+    /**
+     * Strategy to combine multiple [pathSelectionStrategies].
+     *
+     * @see PathSelectorCombinationStrategy
+     */
+    val pathSelectorCombinationStrategy: PathSelectorCombinationStrategy = PathSelectorCombinationStrategy.INTERLEAVED,
+    /**
+     * Seed used for random operations.
+     */
+    val randomSeed: Long = 0,
+    /**
+     * Code coverage percent to stop execution on. Considered only if in range [1..100].
+     */
+    val stopOnCoverage: Int = 100,
+    /**
+     * Optional limit of symbolic execution steps to stop execution on.
+     */
+    val stepLimit: ULong? = 1500u,
+    /**
+     * Optional limit of collected states to stop execution on.
+     */
+    val collectedStatesLimit: Int? = null,
+    /**
+     * Optional timeout in milliseconds to stop execution on.
+     */
+    val timeoutMs: Long? = 20_000,
+    /**
+     * A number of steps from the last terminated state.
+     */
+    val stepsFromLastCovered: Long? = null,
+    /**
+     * Scope of methods which coverage is considered.
+     *
+     * @see CoverageZone
+     */
+    val coverageZone: CoverageZone = CoverageZone.METHOD,
+    /**
+     * Whether we should prefer exceptional state in the queue to the regular ones.
+     */
+    val exceptionsPropagation: Boolean = true,
+    /**
+     * SMT solver type used for path constraint solving.
+     */
+    val solverType: SolverType = SolverType.Z3
+)

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

@@ -0,0 +1,130 @@
+package org.usvm.machine
+
+import mu.KLogging
+import org.jacodb.api.JcClasspath
+import org.jacodb.api.JcMethod
+import org.jacodb.api.cfg.JcInst
+import org.jacodb.api.ext.methods
+import org.usvm.*
+import org.usvm.machine.interpreter.JcInterpreter
+import org.usvm.machine.state.JcMethodResult
+import org.usvm.machine.state.JcState
+import org.usvm.machine.state.lastStmt
+import org.usvm.ps.FeaturesLoggingPathSelector
+import org.usvm.ps.modifiedCreatePathSelector
+import org.usvm.statistics.*
+import org.usvm.stopstrategies.createStopStrategy
+
+val logger = object : KLogging() {}.logger
+
+class ModifiedJcMachine(
+    cp: JcClasspath,
+    private val options: ModifiedUMachineOptions
+) : UMachine<JcState>() {
+        private val applicationGraph = JcApplicationGraph(cp)
+
+    private val typeSystem = JcTypeSystem(cp)
+    private val components = JcComponents(typeSystem, options.solverType)
+    private val ctx = JcContext(cp, components)
+
+    private val interpreter = JcInterpreter(ctx, applicationGraph)
+
+    private val distanceStatistics = DistanceStatistics(applicationGraph)
+
+    fun analyze(
+        method: JcMethod,
+        coverageCounter: CoverageCounter? = null
+    ): List<JcState> {
+        logger.debug("$this.analyze($method)")
+        val initialState = interpreter.getInitialState(method)
+
+        val methodsToTrackCoverage =
+            when (options.coverageZone) {
+                CoverageZone.METHOD -> setOf(method)
+                CoverageZone.TRANSITIVE -> setOf(method)
+                // TODO: more adequate method filtering. !it.isConstructor is used to exclude default constructor which is often not covered
+                CoverageZone.CLASS -> method.enclosingClass.methods.filter {
+                    it.enclosingClass == method.enclosingClass && !it.isConstructor
+                }.toSet()
+            }
+
+        val coverageStatistics: CoverageStatistics<JcMethod, JcInst, JcState> = CoverageStatistics(
+            methodsToTrackCoverage,
+            applicationGraph
+        )
+
+        val pathSelector = modifiedCreatePathSelector(
+            initialState,
+            options,
+            { coverageStatistics },
+            { distanceStatistics },
+            { applicationGraph }
+        )
+
+        val statesCollector = CoveredNewStatesCollector<JcState>(coverageStatistics) {
+            it.methodResult is JcMethodResult.JcException
+        }
+        val stopStrategy = createStopStrategy(
+            UMachineOptions(
+                pathSelectorCombinationStrategy = options.pathSelectorCombinationStrategy,
+                randomSeed = options.randomSeed,
+                stopOnCoverage = options.stopOnCoverage,
+                stepLimit = options.stepLimit,
+                collectedStatesLimit = options.collectedStatesLimit,
+                timeoutMs = options.timeoutMs,
+                stepsFromLastCovered = options.stepsFromLastCovered,
+                coverageZone = options.coverageZone,
+                exceptionsPropagation = options.exceptionsPropagation,
+                solverType = options.solverType
+            ),
+            coverageStatistics = { coverageStatistics },
+            getCollectedStatesCount = { statesCollector.collectedStates.size }
+        )
+
+        val observers = mutableListOf<UMachineObserver<JcState>>(coverageStatistics)
+
+        observers.add(TerminatedStateRemover())
+
+        if (options.coverageZone == CoverageZone.TRANSITIVE) {
+            observers.add(
+                TransitiveCoverageZoneObserver(
+                    initialMethod = method,
+                    methodExtractor = { state -> state.lastStmt.location.method },
+                    addCoverageZone = { coverageStatistics.addCoverageZone(it) },
+                    ignoreMethod = { false } // TODO replace with a configurable setting
+                )
+            )
+        }
+        observers.add(statesCollector)
+
+        val methodName = method.toString().dropWhile { it != ')' }.drop(1)
+        if (coverageCounter != null) {
+            observers.add(CoverageCounterStatistics(coverageStatistics, coverageCounter, methodName))
+        }
+
+        run(
+            interpreter,
+            pathSelector,
+            observer = CompositeUMachineObserver(observers),
+            isStateTerminated = ::isStateTerminated,
+            stopStrategy = stopStrategy,
+        )
+
+        coverageCounter?.finishTest(methodName)
+        if (pathSelector is FeaturesLoggingPathSelector<*, *, *>) {
+            if (MLConfig.logFeatures && MLConfig.mode != Mode.Test) {
+                pathSelector.savePath()
+            }
+        }
+
+        return statesCollector.collectedStates
+    }
+
+    private fun isStateTerminated(state: JcState): Boolean {
+        return state.callStack.isEmpty()
+    }
+
+    override fun close() {
+        components.close()
+    }
+}