Move TrieNode from kotlinx.collections (#194)

* move immutable collections from kotlinx

* use `UPersistentHashMap`; introduce ownership

* add ownership in memory

* numeric constraints: reorder remove and get

* linter warnings fixes

* fix ownership flow; add ownership in type constraints

* use default ownership in composer memory

* use memory ownership in memcpy, union, merge

* fix ownership inheritance

* various fixes

* add ownership in composer

* fixes

* change pc.ownership in clone

* remove scope.ownership

* fixes

* fix set iterator

* fix iterator

* iterate sequence only once in JcStaticFieldsRegion

* fix mutable iterators in NumericConstraints

* add equals, hashCode, toString

* drop RegionTree

* fix hashCode, equals; fix union in map, set

* move some logic in extensions

* toList in equalityConstrints

* fix compile err

* add toList in some iterators

* style

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

@@ -15,14 +15,16 @@ import org.usvm.collection.set.primitive.UInputSetReading
 import org.usvm.collection.set.ref.UAllocatedRefSetWithInputElementsReading
 import org.usvm.collection.set.ref.UInputRefSetWithAllocatedElementsReading
 import org.usvm.collection.set.ref.UInputRefSetWithInputElementsReading
+import org.usvm.collections.immutable.internal.MutabilityOwnership
 import org.usvm.memory.UReadOnlyMemory
 import org.usvm.memory.USymbolicCollectionId
 import org.usvm.regions.Region
 
 @Suppress("MemberVisibilityCanBePrivate")
 open class UComposer<Type, USizeSort : USort>(
     ctx: UContext<USizeSort>,
-    val memory: UReadOnlyMemory<Type>
+    val memory: UReadOnlyMemory<Type>,
+    val ownership: MutabilityOwnership
 ) : UExprTransformer<Type, USizeSort>(ctx) {
     open fun <Sort : USort> compose(expr: UExpr<Sort>): UExpr<Sort> = apply(expr)
 

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

@@ -42,6 +42,7 @@ import org.usvm.collection.set.ref.UInputRefSetWithAllocatedElements
 import org.usvm.collection.set.ref.UInputRefSetWithAllocatedElementsReading
 import org.usvm.collection.set.ref.UInputRefSetWithInputElements
 import org.usvm.collection.set.ref.UInputRefSetWithInputElementsReading
+import org.usvm.collections.immutable.internal.MutabilityOwnership
 import org.usvm.memory.UAddressCounter
 import org.usvm.memory.UReadOnlyMemory
 import org.usvm.memory.splitUHeapRef
@@ -61,11 +62,12 @@ open class UContext<USizeSort : USort>(
     private val solver by lazy { components.mkSolver(this) }
     private val typeSystem by lazy { components.mkTypeSystem(this) }
     private val softConstraintsProvider by lazy { components.mkSoftConstraintsProvider(this) }
-    private val composerBuilder: (UReadOnlyMemory<*>) -> UComposer<*, USizeSort> by lazy {
+    private val composerBuilder: (UReadOnlyMemory<*>, MutabilityOwnership) -> UComposer<*, USizeSort> by lazy {
         @Suppress("UNCHECKED_CAST")
-        components.mkComposer(this) as (UReadOnlyMemory<*>) -> UComposer<*, USizeSort>
+        components.mkComposer(this) as (UReadOnlyMemory<*>, MutabilityOwnership) -> UComposer<*, USizeSort>
     }
 
+    val defaultOwnership = MutabilityOwnership()
     val sizeExprs by lazy { components.mkSizeExprProvider(this) }
     val statesForkProvider by lazy { components.mkStatesForkProvider() }
 
@@ -86,7 +88,8 @@ open class UContext<USizeSort : USort>(
 
     fun <Type> softConstraintsProvider(): USoftConstraintsProvider<Type, USizeSort> = softConstraintsProvider.cast()
 
-    fun <Type> composer(memory: UReadOnlyMemory<Type>): UComposer<Type, USizeSort> = composerBuilder(memory).cast()
+    fun <Type> composer(memory: UReadOnlyMemory<Type>, ownership: MutabilityOwnership): UComposer<Type, USizeSort> =
+        composerBuilder(memory, ownership).cast()
 
     val addressSort: UAddressSort = mkUninterpretedSort("Address")
     val nullRef: UNullRef = UNullRef(this)

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

@@ -2,9 +2,11 @@ package org.usvm
 
 import io.ksmt.utils.cast
 import kotlinx.collections.immutable.PersistentList
-import kotlinx.collections.immutable.PersistentMap
-import kotlinx.collections.immutable.persistentHashMapOf
 import kotlinx.collections.immutable.persistentListOf
+import org.usvm.collections.immutable.getOrDefault
+import org.usvm.collections.immutable.persistentHashMapOf
+import org.usvm.collections.immutable.implementations.immutableMap.UPersistentHashMap
+import org.usvm.collections.immutable.internal.MutabilityOwnership
 import org.usvm.merging.MergeGuard
 import org.usvm.merging.UMergeable
 
@@ -19,48 +21,58 @@ interface UMocker<Method> : UMockEvaluator {
         method: Method,
         args: Sequence<UExpr<out USort>>,
         sort: Sort,
+        ownership: MutabilityOwnership,
     ): UMockSymbol<Sort>
-    val trackedLiterals: Collection<TrackedLiteral>
+    val trackedLiterals: Sequence<TrackedLiteral>
 
-    fun <Sort : USort> createMockSymbol(trackedLiteral: TrackedLiteral?, sort: Sort): UExpr<Sort>
+    fun <Sort : USort> createMockSymbol(
+        trackedLiteral: TrackedLiteral?,
+        sort: Sort,
+        ownership: MutabilityOwnership,
+    ): UExpr<Sort>
 
     fun getTrackedExpression(trackedLiteral: TrackedLiteral): UExpr<USort>
 
     fun clone(): UMocker<Method>
 }
 
 class UIndexedMocker<Method>(
-    private var methodMockClauses: PersistentMap<Method, PersistentList<UMockSymbol<out USort>>> = persistentHashMapOf(),
-    private var trackedSymbols: PersistentMap<TrackedLiteral, UExpr<out USort>> = persistentHashMapOf(),
+    private var methodMockClauses: UPersistentHashMap<Method, PersistentList<UMockSymbol<out USort>>> = persistentHashMapOf(),
+    private var trackedSymbols: UPersistentHashMap<TrackedLiteral, UExpr<out USort>> = persistentHashMapOf(),
     private var untrackedSymbols: PersistentList<UExpr<out USort>> = persistentListOf(),
 ) : UMocker<Method>, UMergeable<UIndexedMocker<Method>, MergeGuard> {
     override fun <Sort : USort> call(
         method: Method,
         args: Sequence<UExpr<out USort>>,
         sort: Sort,
+        ownership: MutabilityOwnership,
     ): UMockSymbol<Sort> {
         val currentClauses = methodMockClauses.getOrDefault(method, persistentListOf())
         val index = currentClauses.size
         val const = sort.uctx.mkIndexedMethodReturnValue(method, index, sort)
-        methodMockClauses = methodMockClauses.put(method, currentClauses.add(const))
+        methodMockClauses = methodMockClauses.put(method, currentClauses.add(const), ownership)
 
         return const
     }
 
     override fun <Sort : USort> eval(symbol: UMockSymbol<Sort>): UExpr<Sort> = symbol
 
-    override val trackedLiterals: Collection<TrackedLiteral>
+    override val trackedLiterals: Sequence<TrackedLiteral>
         get() = trackedSymbols.keys
 
     /**
      * Creates a mock symbol. If [trackedLiteral] is not null, created expression
      * can be retrieved later by this [trackedLiteral] using [getTrackedExpression] method.
      */
-    override fun <Sort : USort> createMockSymbol(trackedLiteral: TrackedLiteral?, sort: Sort): UExpr<Sort> {
+    override fun <Sort : USort> createMockSymbol(
+        trackedLiteral: TrackedLiteral?,
+        sort: Sort,
+        ownership: MutabilityOwnership,
+    ): UExpr<Sort> {
         val const = sort.uctx.mkTrackedSymbol(sort)
 
         if (trackedLiteral != null) {
-            trackedSymbols = trackedSymbols.put(trackedLiteral, const)
+            trackedSymbols = trackedSymbols.put(trackedLiteral, const, ownership)
         } else {
             untrackedSymbols = untrackedSymbols.add(const)
         }
@@ -71,10 +83,11 @@ class UIndexedMocker<Method>(
     override fun getTrackedExpression(trackedLiteral: TrackedLiteral): UExpr<USort> {
         if (trackedLiteral !in trackedSymbols) error("Access by unregistered track literal $trackedLiteral")
 
-        return trackedSymbols.getValue(trackedLiteral).cast()
+        return trackedSymbols[trackedLiteral]!!.cast()
     }
 
-    override fun clone(): UIndexedMocker<Method> = UIndexedMocker(methodMockClauses, trackedSymbols, untrackedSymbols)
+    override fun clone(): UIndexedMocker<Method> =
+        UIndexedMocker(methodMockClauses, trackedSymbols, untrackedSymbols)
 
     /**
      * Check if this [UIndexedMocker] can be merged with [other] indexed mocker.
@@ -83,7 +96,10 @@ class UIndexedMocker<Method>(
      *
      * @return the merged indexed mocker.
      */
-    override fun mergeWith(other: UIndexedMocker<Method>, by: MergeGuard): UIndexedMocker<Method>? {
+    override fun mergeWith(
+        other: UIndexedMocker<Method>,
+        by: MergeGuard,
+    ): UIndexedMocker<Method>? {
         if (methodMockClauses !== other.methodMockClauses
             || trackedSymbols !== other.trackedSymbols
             || untrackedSymbols !== other.untrackedSymbols

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

@@ -1,5 +1,6 @@
 package org.usvm
 
+import org.usvm.collections.immutable.internal.MutabilityOwnership
 import org.usvm.constraints.UPathConstraints
 import org.usvm.memory.UMemory
 import org.usvm.merging.UMergeable
@@ -12,6 +13,7 @@ typealias StateId = UInt
 abstract class UState<Type, Method, Statement, Context, Target, State>(
     // TODO: add interpreter-specific information
     val ctx: Context,
+    initOwnership: MutabilityOwnership,
     open val callStack: UCallStack<Method, Statement>,
     open val pathConstraints: UPathConstraints<Type>,
     open val memory: UMemory<Type, Method>,
@@ -33,6 +35,9 @@ abstract class UState<Type, Method, Statement, Context, Target, State>(
      */
     val id: StateId = ctx.getNextStateId()
 
+    open var ownership = initOwnership
+        protected set
+
     /**
      * Creates new state structurally identical to this.
      * If [newConstraints] is null, clones [pathConstraints]. Otherwise, uses [newConstraints] in cloned state.

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

@@ -1,5 +1,6 @@
 package org.usvm
 
+import org.usvm.collections.immutable.internal.MutabilityOwnership
 import org.usvm.model.UModelBase
 import org.usvm.solver.USatResult
 import org.usvm.solver.UUnknownResult
@@ -155,7 +156,11 @@ object WithSolverStateForker : StateForker {
         val satResult = solver.check(constraintsToCheck)
 
         return when (satResult) {
-            is UUnsatResult -> null
+            is UUnsatResult -> {
+                // rollback previous ownership
+                state.pathConstraints.changeOwnership(state.ownership)
+                null
+            }
 
             is USatResult -> {
                 // Note that we cannot extract common code here due to
@@ -177,6 +182,8 @@ object WithSolverStateForker : StateForker {
             }
 
             is UUnknownResult -> {
+                // rollback previous ownership
+                state.pathConstraints.changeOwnership(state.ownership)
                 state.pathConstraints += if (stateToCheck) newConstraintToOriginalState else newConstraintToForkedState
 
                 null
@@ -192,11 +199,11 @@ object NoSolverStateForker : StateForker {
     ): ForkResult<T> {
         val (trueModels, falseModels, _) = splitModelsByCondition(state.models, condition)
         val notCondition = state.ctx.mkNot(condition)
-
         val clonedPathConstraints = state.pathConstraints.clone()
         clonedPathConstraints += condition
 
         val (posState, negState) = if (clonedPathConstraints.isFalse) {
+            // changing ownership is unnecessary
             state.pathConstraints += notCondition
             state.models = falseModels
 
@@ -225,8 +232,7 @@ object NoSolverStateForker : StateForker {
         val result = mutableListOf<T?>()
         for (condition in conditions) {
             val (trueModels, _) = splitModelsByCondition(curState.models, condition)
-
-            val clonedConstraints = curState.pathConstraints.clone()
+            val clonedConstraints = curState.pathConstraints.clone(MutabilityOwnership(), MutabilityOwnership())
             clonedConstraints += condition
 
             if (clonedConstraints.isFalse) {

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

@@ -1,5 +1,6 @@
 package org.usvm
 
+import org.usvm.collections.immutable.internal.MutabilityOwnership
 import org.usvm.memory.UReadOnlyMemory
 import org.usvm.model.ULazyModelDecoder
 import org.usvm.model.UModelDecoder
@@ -34,8 +35,8 @@ interface UComponents<Type, USizeSort : USort> {
 
     fun <Context : UContext<USizeSort>> mkComposer(
         ctx: Context,
-    ): (UReadOnlyMemory<Type>) -> UComposer<Type, USizeSort> =
-        { memory: UReadOnlyMemory<Type> -> UComposer(ctx, memory) }
+    ): (UReadOnlyMemory<Type>, MutabilityOwnership) -> UComposer<Type, USizeSort> =
+        { memory: UReadOnlyMemory<Type>, ownership: MutabilityOwnership -> UComposer(ctx, memory, ownership) }
 
     fun mkStatesForkProvider(): StateForker = if (useSolverForForks) WithSolverStateForker else NoSolverStateForker
 

usvm-core/src/main/kotlin/org/usvm/api/MemoryApi.kt

@@ -6,11 +6,12 @@ import org.usvm.UContext
 import org.usvm.UExpr
 import org.usvm.UHeapRef
 import org.usvm.USort
-import org.usvm.memory.UMemory
 import org.usvm.memory.UReadOnlyMemory
 import org.usvm.memory.UWritableMemory
 import org.usvm.collection.array.UArrayIndexLValue
 import org.usvm.collection.array.length.UArrayLengthLValue
+import org.usvm.collection.array.memcpy as memcpyInternal
+import org.usvm.collection.array.memset as memsetInternal
 import org.usvm.collection.field.UFieldLValue
 import org.usvm.collection.set.primitive.USetEntryLValue
 import org.usvm.collection.set.ref.URefSetEntryLValue
@@ -22,8 +23,6 @@ import org.usvm.regions.Region
 import org.usvm.types.UTypeStream
 import org.usvm.uctx
 import org.usvm.withSizeSort
-import org.usvm.collection.array.memcpy as memcpyInternal
-import org.usvm.collection.array.memset as memsetInternal
 import org.usvm.collection.array.allocateArray as allocateArrayInternal
 import org.usvm.collection.array.allocateArrayInitialized as allocateArrayInitializedInternal
 

usvm-core/src/main/kotlin/org/usvm/api/MockApi.kt

@@ -15,7 +15,7 @@ fun <Method, T : USort> UState<*, Method, *, *, *, *>.makeSymbolicPrimitive(
     sort: T
 ): UExpr<T> {
     check(sort != sort.uctx.addressSort) { "$sort is not primitive" }
-    return memory.mocker.createMockSymbol(trackedLiteral = null, sort)
+    return memory.mocker.createMockSymbol(trackedLiteral = null, sort, ownership)
 }
 
 fun <Type, Method, State> StepScope<State, Type, *, *>.makeSymbolicRef(
@@ -34,7 +34,7 @@ fun <Type, Method, State> StepScope<State, Type, *, *>.makeNullableSymbolicRefWi
     mockSymbolicRef { ctx.mkOr(objectTypeEquals(it, representative), ctx.mkEq(it, ctx.nullRef)) }
 
 fun <Method> UState<*, Method, *, *, *, *>.makeSymbolicRefUntyped(): UHeapRef =
-    memory.mocker.createMockSymbol(trackedLiteral = null, ctx.addressSort)
+    memory.mocker.createMockSymbol(trackedLiteral = null, ctx.addressSort, ownership)
 
 private inline fun <Type, Method, State> StepScope<State, Type, *, *>.mockSymbolicRef(
     crossinline mkTypeConstraint: State.(UHeapRef) -> UBoolExpr