Add model validation

Author: Parzival-05

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

@@ -1,13 +1,13 @@
 package org.usvm.ps
+
 import org.usvm.UPathSelector
 import org.usvm.UState
 import org.usvm.statistics.BasicBlock
 import org.usvm.statistics.BlockGraph
 import org.usvm.statistics.StepsStatistics
-import org.usvm.util.OnnxModel
-import org.usvm.util.Oracle
 import org.usvm.util.Predictor
 import org.usvm.utils.Game
+import org.usvm.utils.OnnxModelImpl
 import org.usvm.utils.StateWrapper
 import org.usvm.utils.isSat
 
@@ -17,8 +17,7 @@ class AIPathSelector<Statement, State, Block>(
     private val stepsStatistics: StepsStatistics<*, State>,
     private val predictor: Predictor<Game<Block>>,
 ) : UPathSelector<State> where
-State : UState<*, *, Statement, *, *, State>,
-Block : BasicBlock {
+State : UState<*, *, Statement, *, *, State>, Block : BasicBlock {
     private val statesMap = mutableMapOf<State, StateWrapper<Statement, State, Block>>()
     private val lastPeekedState: State?
         get() = stepsStatistics.lastPeekedState
@@ -43,12 +42,14 @@ Block : BasicBlock {
     }
 
     private fun buildGame(
-        vertices: List<Block>,
-        wrappers: MutableCollection<StateWrapper<Statement, State, Block>>
+        vertices: List<Block>, wrappers: MutableCollection<StateWrapper<Statement, State, Block>>
     ): Game<Block> {
-        val game = when (predictor) {
-            is OnnxModel<Game<Block>> -> Game(vertices, wrappers, blockGraph)
-            is Oracle<Game<Block>> -> {
+        val game = when {
+            predictor is OnnxModelImpl<*> && !predictor.isTrainMode -> {
+                Game(vertices, wrappers, blockGraph)
+            }
+
+            else -> {
                 // if we played with default searcher before,
                 // client has no information about the game
                 if (firstSend) {
@@ -110,8 +111,7 @@ Block : BasicBlock {
              * If state [isSat], some blocks may have been covered by
              * [org.usvm.statistics.CoverageStatistics.onStateTerminated]
              */
-            if (state.isSat())
-                touchedBlocks.addAll(wrapper.history.keys)
+            if (state.isSat()) touchedBlocks.addAll(wrapper.history.keys)
 
             touchedStates.remove(wrapper)
         }
@@ -120,15 +120,9 @@ Block : BasicBlock {
     override fun add(states: Collection<State>) {
         // is null iff we are adding initial state
         val lastPeekedStateWrapper = statesMap[lastPeekedState]
-        val parentPathConditionSize = lastPeekedStateWrapper?.pathConditionSize ?: 0
         val parentHistory = lastPeekedStateWrapper?.history ?: mutableMapOf()
         val wrappers = states.map { state ->
-            val wrapper = StateWrapper(
-                state,
-                parentPathConditionSize,
-                parentHistory,
-                blockGraph
-            )
+            val wrapper = StateWrapper(state, parentHistory, blockGraph)
             statesMap[state] = wrapper
             wrapper
         }

usvm-core/src/main/kotlin/org/usvm/utils/OnnxModelImpl.kt

@@ -1,6 +1,7 @@
 package org.usvm.utils
 
 import ai.onnxruntime.OnnxTensor
+import ai.onnxruntime.OnnxTensor.createTensor
 import ai.onnxruntime.OrtEnvironment
 import ai.onnxruntime.OrtSession
 import org.usvm.StateId
@@ -10,16 +11,38 @@ import org.usvm.statistics.BlockGraph
 import org.usvm.util.OnnxModel
 import java.nio.FloatBuffer
 import java.nio.LongBuffer
+import kotlin.collections.toLongArray
 
 enum class Mode {
-    CPU,
-    GPU
+    CPU, GPU
 }
 
-class OnnxModelImpl<Block: BasicBlock>(
-    pathToONNX: String,
-    mode: Mode
-): OnnxModel<Game<Block>> {
+data class AIGameStep<T : BasicBlock>(
+    val Game: Game<T>, val NNOutput: Array<FloatArray>
+) {
+    // autogenerated
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as AIGameStep<*>
+
+        if (Game != other.Game) return false
+        if (!NNOutput.contentDeepEquals(other.NNOutput)) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = Game.hashCode()
+        result = 31 * result + NNOutput.contentDeepHashCode()
+        return result
+    }
+}
+
+class OnnxModelImpl<Block : BasicBlock>(
+    pathToONNX: String, mode: Mode, val isTrainMode: Boolean, val saveStep: (AIGameStep<Block>) -> Unit = {}
+) : OnnxModel<Game<Block>> {
     private val env: OrtEnvironment = OrtEnvironment.getEnvironment()
     private val sessionOptions: OrtSession.SessionOptions = OrtSession.SessionOptions().apply {
         if (mode == Mode.GPU) {
@@ -34,48 +57,47 @@ class OnnxModelImpl<Block: BasicBlock>(
         val stateIds = mutableMapOf<StateId, Int>()
         val input = generateInput(game, stateIds)
         val output = session.run(input)
-
-        val predictedStatesRanks =
-            (output["out"].get().value as Array<*>).map { (it as FloatArray).toList() }
-
+        val predictedStates = output["out"].get().value as Array<*>
+        val predictedStatesRanks = predictedStates.map { (it as FloatArray).toList() }
+        saveStep(AIGameStep(game, predictedStates.map { it as FloatArray }.toTypedArray()))
         return checkNotNull(getPredictedState(predictedStatesRanks, stateIds))
     }
 
     private fun generateInput(
-        game: Game<Block>,
-        stateIds: MutableMap<StateId, Int>
+        game: Game<Block>, stateIds: MutableMap<StateId, Int>
     ): Map<String, OnnxTensor> {
         val (vertices, stateWrappers, blockGraph) = game
         val vertexIds = mutableMapOf<Int, Int>()
         val gameVertices = tensorFromBasicBlocks(vertices, vertexIds)
         val states = tensorFromStates(stateWrappers, stateIds)
         val (vertexToVertexEdgesIndex, vertexToVertexEdgesAttributes) = tensorFromVertexEdges(
-            vertices,
-            blockGraph,
-            vertexIds
+            vertices, blockGraph, vertexIds
         )
         val (parentOfEdges, historyEdgesIndexVertexToState, historyEdgesAttributes) = tensorFromStateEdges(
-            stateWrappers,
-            stateIds,
-            vertexIds
+            stateWrappers, stateIds, vertexIds
         )
         val vertexToState = tensorFromStatePositions(stateWrappers, stateIds, vertexIds)
-
+        val mockPC1 = createTensor(env, FloatBuffer.wrap(FloatArray(0)), longArrayOf(0, 49))
+        val mockPC2 = createTensor(env, LongBuffer.wrap(LongArray(0)), longArrayOf(2, 0))
+        val mockPC3 =
+            createTensor(env, LongBuffer.wrap(LongArray(0)), longArrayOf(2, 0))
         return mapOf(
             "game_vertex" to gameVertices,
             "state_vertex" to states,
+            "path_condition_vertex" to mockPC1,
             "gamevertex_to_gamevertex_index" to vertexToVertexEdgesIndex,
             "gamevertex_to_gamevertex_type" to vertexToVertexEdgesAttributes,
             "gamevertex_history_statevertex_index" to historyEdgesIndexVertexToState,
             "gamevertex_history_statevertex_attrs" to historyEdgesAttributes,
             "gamevertex_in_statevertex" to vertexToState,
-            "statevertex_parentof_statevertex" to parentOfEdges
+            "statevertex_parentof_statevertex" to parentOfEdges,
+            "pathconditionvertex_to_pathconditionvertex" to mockPC2,
+            "pathconditionvertex_to_statevertex" to mockPC3,
         )
     }
 
     private fun getPredictedState(stateRank: List<List<Float>>, stateIds: Map<StateId, Int>): StateId? {
-        return stateRank
-            .mapIndexed { index, ranks -> stateIds.entries.find { it.value == index }?.key to ranks.sum() }
+        return stateRank.mapIndexed { index, ranks -> stateIds.entries.find { it.value == index }?.key to ranks.sum() }
             .maxBy { it.second }.first
     }
 
@@ -97,13 +119,14 @@ class OnnxModelImpl<Block: BasicBlock>(
         return createFloatTensor(verticesArray, vertices.size, numOfVertexAttributes)
     }
 
-    private fun tensorFromStates(states: Collection<StateWrapper<*, *, *>>, stateIds: MutableMap<StateId, Int>): OnnxTensor {
-        val numOfStateAttributes = 7
+    private fun tensorFromStates(
+        states: Collection<StateWrapper<*, *, *>>, stateIds: MutableMap<StateId, Int>
+    ): OnnxTensor {
+        val numOfStateAttributes = 6
         val statesArray = states.flatMapIndexed { i, state ->
             stateIds[state.id] = i
             listOf(
                 state.position,
-                state.pathConditionSize,
                 state.visitedAgainVertices,
                 state.visitedNotCoveredVerticesInZone,
                 state.visitedNotCoveredVerticesOutOfZone,
@@ -116,9 +139,7 @@ class OnnxModelImpl<Block: BasicBlock>(
     }
 
     private fun tensorFromVertexEdges(
-        blocks: Collection<Block>,
-        blockGraph: BlockGraph<*, Block, *>,
-        vertexIds: MutableMap<Int, Int>
+        blocks: Collection<Block>, blockGraph: BlockGraph<*, Block, *>, vertexIds: MutableMap<Int, Int>
     ): Pair<OnnxTensor, OnnxTensor> {
         val vertexFrom = mutableListOf<Long>()
         val vertexTo = mutableListOf<Long>()
@@ -136,15 +157,12 @@ class OnnxModelImpl<Block: BasicBlock>(
         val indexList = (vertexFrom + vertexTo).toLongArray()
 
         return createLongTensor(indexList, 2, vertexFrom.size) to createLongTensor(
-            attributes.toLongArray(),
-            attributes.size
+            attributes.toLongArray(), attributes.size
         )
     }
 
     private fun tensorFromStateEdges(
-        states: Collection<StateWrapper<*, *, *>>,
-        stateIds: MutableMap<StateId, Int>,
-        vertexIds: MutableMap<Int, Int>
+        states: Collection<StateWrapper<*, *, *>>, stateIds: MutableMap<StateId, Int>, vertexIds: MutableMap<Int, Int>
     ): Triple<OnnxTensor, OnnxTensor, OnnxTensor> {
         val numOfParentOfEdges = states.sumOf { it.children.size }
         val numOfHistoryEdges = states.sumOf { it.history.size }
@@ -173,20 +191,15 @@ class OnnxModelImpl<Block: BasicBlock>(
         }
 
         val parentTensor = createLongTensor(parentOf, 2, numOfParentOfEdges)
-        val historyIndexTensor =
-            createLongTensor(historyIndexVertexToState, 2, numOfHistoryEdges)
+        val historyIndexTensor = createLongTensor(historyIndexVertexToState, 2, numOfHistoryEdges)
         val historyAttributesTensor = createLongTensor(
-            historyAttributes,
-            numOfHistoryEdges,
-            numOfHistoryEdgeAttributes
+            historyAttributes, numOfHistoryEdges, numOfHistoryEdgeAttributes
         )
         return Triple(parentTensor, historyIndexTensor, historyAttributesTensor)
     }
 
     private fun tensorFromStatePositions(
-        states: Collection<StateWrapper<*, *, *>>,
-        stateIds: MutableMap<StateId, Int>,
-        vertexIds: MutableMap<Int, Int>
+        states: Collection<StateWrapper<*, *, *>>, stateIds: MutableMap<StateId, Int>, vertexIds: MutableMap<Int, Int>
     ): OnnxTensor {
         val totalStates = states.size
         val vertexToState = LongArray(2 * totalStates)

usvm-core/src/main/kotlin/org/usvm/utils/StateWrapper.kt

@@ -13,7 +13,6 @@ data class StateHistoryElement(
 
 class StateWrapper<Statement, State, Block>(
     private val state: State,
-    private val parentPathConditionSize: Int,
     private val parentHistory: MutableMap<Block, StateHistoryElement>,
     private val blockGraph: BlockGraph<*, Block, Statement>,
 ) where State : UState<*, *, Statement, *, *, State>, Block : BasicBlock {
@@ -24,7 +23,6 @@ class StateWrapper<Statement, State, Block>(
     private var visitedStatement: Statement? = null
     lateinit var currentBlock: Block
     var position by Delegates.notNull<Int>()
-    var pathConditionSize by Delegates.notNull<Int>()
     var visitedAgainVertices by Delegates.notNull<Int>()
     var visitedNotCoveredVerticesInZone by Delegates.notNull<Int>()
     var visitedNotCoveredVerticesOutOfZone by Delegates.notNull<Int>()
@@ -42,7 +40,6 @@ class StateWrapper<Statement, State, Block>(
             currentBlock.states.add(this@StateWrapper.id)
 
             position = currentBlock.id
-            pathConditionSize = parentPathConditionSize + state.forkPoints.depth
             instructionsVisitedInCurrentBlock = 0
         }
         instructionsVisitedInCurrentBlock++

usvm-ml-gameserver/build.gradle.kts

@@ -5,6 +5,7 @@ plugins {
     id(Plugins.Ktor)
     kotlin("plugin.serialization") version Versions.kotlin
     application
+    kotlin
 }
 
 application {