Fix splitting for range updates

usvm-core/src/main/kotlin/org/usvm/memory/SymbolicCollectionUpdates.kt

@@ -1,7 +1,7 @@
 package org.usvm.memory
 
+import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentMapOf
-import kotlinx.collections.immutable.toPersistentMap
 import org.usvm.UBoolExpr
 import org.usvm.UComposer
 import org.usvm.UExpr
@@ -317,44 +317,52 @@ data class UTreeUpdates<Key, Reg : Region<Reg>, Sort : USort>(
         guardBuilder: GuardBuilder,
         composer: UComposer<*, *>?,
     ): UTreeUpdates<Key, Reg, Sort> {
-        var restRecordsAreIrrelevant = false
-        val symbolicSubtreesStack =
-            mutableListOf<Pair<Reg, Pair<UUpdateNode<Key, Sort>, RegionTree<Reg, UUpdateNode<Key, Sort>>>>>()
-
-        fun traverseTreeWhilePredicateSatisfied(tree: RegionTree<Reg, UUpdateNode<Key, Sort>>) {
+        // returns fully symbolic tree entries after splitting
+        fun traverseTreeWhilePredicateSatisfied(
+            tree: RegionTree<Reg, UUpdateNode<Key, Sort>>,
+        ): List<Pair<Reg, Pair<UUpdateNode<Key, Sort>, RegionTree<Reg, UUpdateNode<Key, Sort>>>>> {
+            val symbolicEntries =
+                mutableListOf<Pair<Reg, Pair<UUpdateNode<Key, Sort>, RegionTree<Reg, UUpdateNode<Key, Sort>>>>>()
             // process nodes of the same level from newest to oldest
-            for ((reg, entry) in tree.entries.toList().reversed()) {
+            for ((reg, entry) in tree.entries.toList().asReversed()) {
                 val (node, subtree) = entry
                 val splitUpdate = node.split(key, predicate, matchingWrites, guardBuilder, composer)
-                // range nodes are considered to belong to invariant tree since they can contain concretes
-                if (splitUpdate === null || node is URangedUpdateNode<*, *, Key, Sort>) {
-                    traverseTreeWhilePredicateSatisfied(subtree)
-                    if (restRecordsAreIrrelevant) break
-                } else {
-                    // since [splitWrite] preserves invariant that records satisfying predicate are always
-                    // at the top of the tree  we can just add rest branch to splitTreeEntries
-                    symbolicSubtreesStack += reg to entry
+                when (splitUpdate) {
+                    null -> {
+                        val restSymbolicTree = traverseTreeWhilePredicateSatisfied(subtree)
+                        symbolicEntries += restSymbolicTree
+                    }
+                    // range nodes are considered to belong to concrete subtree since they can contain concretes
+                    is URangedUpdateNode<*, *, Key, Sort> -> {
+                        val restSymbolicEntries = traverseTreeWhilePredicateSatisfied(subtree)
+                        val restTree = RegionTree(restSymbolicEntries.asReversed().toPersistentMap())
+                        symbolicEntries.add(reg to (splitUpdate to restTree))
+                    }
+
+                    else -> {
+                        // since [splitWrite] preserves invariant that records satisfying predicate are always
+                        // at the top of the tree  we can just add rest branch to splitTreeEntries
+                        symbolicEntries.add(reg to entry)
+                    }
                 }
 
-                if (guardBuilder.isFalse) {
-                    restRecordsAreIrrelevant = true
-                    break
-                }
+                if (guardBuilder.isFalse) break
             }
+
+            return symbolicEntries
         }
 
-        traverseTreeWhilePredicateSatisfied(updates)
+        val symbolicEntries = traverseTreeWhilePredicateSatisfied(updates)
 
-        val splitTreeEntries =
-            mutableMapOf<Reg, Pair<UUpdateNode<Key, Sort>, RegionTree<Reg, UUpdateNode<Key, Sort>>>>()
+        return this.copy(updates = RegionTree(symbolicEntries.asReversed().toPersistentMap()))
+    }
 
-        // reconstruct region tree, including all updates unsatisfying `predicate(update.value(key))` in the same order
-        while (symbolicSubtreesStack.isNotEmpty()) {
-            val (reg, entry) = symbolicSubtreesStack.removeLast()
-            splitTreeEntries[reg] = entry
+    private fun <K, V> List<Pair<K, V>>.toPersistentMap(): PersistentMap<K, V> {
+        val builder = persistentMapOf<K, V>().builder()
+        for ((reg, tree) in this) {
+            builder[reg] = tree
         }
-
-        return this.copy(updates = RegionTree(splitTreeEntries.toPersistentMap()))
+        return builder.build()
     }
 
     override fun splitWrite(

usvm-core/src/test/kotlin/org/usvm/memory/HeapRefSplittingTest.kt

@@ -445,4 +445,41 @@ class HeapRefSplittingTest {
 
         return array to ref
     }
+
+    @Test
+    fun testRangedUpdateWithMixedWrite() = with(ctx) {
+        val ref = allocateConcreteRef()
+        val idx0 = mkSizeExpr(0)
+        val idx1 = mkSizeExpr(1)
+        val idx2 = mkSizeExpr(2)
+
+        val concreteValue1 = allocateConcreteRef()
+        val symbolicValue2 = mkRegisterReading(0, addressSort)
+        val concreteValue3 = allocateConcreteRef()
+
+        heap.writeArrayIndex(ref, idx0, arrayDescr.first, arrayDescr.second, concreteValue1, trueExpr)
+        heap.writeArrayIndex(ref, idx1, arrayDescr.first, arrayDescr.second, symbolicValue2, trueExpr)
+        heap.writeArrayIndex(ref, idx2, arrayDescr.first, arrayDescr.second, concreteValue3, trueExpr)
+
+        heap.memcpy(ref, ref, arrayDescr.first, arrayDescr.second, idx1, idx0, idx1)
+
+        val read = heap.readArrayIndex(ref, idx0, arrayDescr.first, arrayDescr.second)
+        assertSame(symbolicValue2, read)
+    }
+
+    @Test
+    fun testClone() = with(ctx) {
+        val ref = allocateConcreteRef()
+        val idx = mkSizeExpr(2)
+
+        val concreteValue = allocateConcreteRef()
+
+        heap.writeArrayIndex(ref, idx, arrayDescr.first, arrayDescr.second, concreteValue, trueExpr)
+
+        val copyRef = allocateConcreteRef()
+        heap.memcpy(ref, copyRef, arrayDescr.first, arrayDescr.second, mkBv(0), mkBv(0), mkBv(3))
+
+        val read = heap.readArrayIndex(copyRef, idx, arrayDescr.first, arrayDescr.second)
+        assertSame(read, concreteValue)
+    }
 }