Fix set union performance (#226)

* Fix set union

* Fix detekt issues

Author: Saloed

usvm-core/src/main/kotlin/org/usvm/collection/set/primitive/USymbolicSetUnionAdapter.kt

@@ -35,9 +35,25 @@ sealed class USymbolicSetUnionAdapter<
     override fun includesConcretely(key: DstKey) =
         includesSymbolically(key, composer = null).isTrue
 
+    private var lastIncludesSymbolicallyCheck: IncludesSymbolicallyCache<SrcKey>? = null
+
     override fun includesSymbolically(key: DstKey, composer: UComposer<*, *>?): UBoolExpr {
         val srcKey = convert(key, composer)
-        return setOfKeys.read(srcKey, composer)
+
+        /**
+         * In the case of deep set union hierarchy we have multiple checks of the same key.
+         * We can cache the last checked key to overcome this issue
+         * */
+        val prevIncludesSymbolicallyCache = lastIncludesSymbolicallyCheck
+        if (prevIncludesSymbolicallyCache != null) {
+            if (prevIncludesSymbolicallyCache.key == srcKey) {
+                return prevIncludesSymbolicallyCache.result
+            }
+        }
+
+        return setOfKeys.read(srcKey, composer).also {
+            lastIncludesSymbolicallyCheck = IncludesSymbolicallyCache(srcKey, it)
+        }
     }
 
     override fun isIncludedByUpdateConcretely(
@@ -49,6 +65,8 @@ sealed class USymbolicSetUnionAdapter<
         "(union $collection)"
 
     abstract override fun collectSetElements(elements: USymbolicSetElementsCollector.Elements<DstKey>)
+
+    private data class IncludesSymbolicallyCache<Key>(val key: Key, val result: UBoolExpr)
 }
 
 class UAllocatedToAllocatedSymbolicSetUnionAdapter<SetType, ElemSort : USort>(

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

@@ -1,19 +1,26 @@
 package org.usvm.memory
 
+import io.ksmt.expr.KBitVecValue
+import io.ksmt.sort.KBv32Sort
 import io.mockk.every
 import io.mockk.mockk
 import org.junit.jupiter.api.BeforeEach
 import org.usvm.Type
+import org.usvm.UBoolExpr
 import org.usvm.UBv32SizeExprProvider
+import org.usvm.UBv32Sort
 import org.usvm.UComponents
 import org.usvm.UContext
 import org.usvm.UExpr
 import org.usvm.UHeapRef
 import org.usvm.USizeSort
+import org.usvm.UTransformer
 import org.usvm.api.refSetAddElement
 import org.usvm.api.refSetRemoveElement
 import org.usvm.api.setAddElement
 import org.usvm.api.setRemoveElement
+import org.usvm.apply
+import org.usvm.collection.set.primitive.USetEntryLValue
 import org.usvm.collection.set.primitive.setEntries
 import org.usvm.collection.set.primitive.setUnion
 import org.usvm.collection.set.ref.refSetEntries
@@ -23,6 +30,7 @@ import org.usvm.constraints.UEqualityConstraints
 import org.usvm.constraints.UTypeConstraints
 import org.usvm.isTrue
 import org.usvm.memory.key.USizeExprKeyInfo
+import org.usvm.regions.SetRegion
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
@@ -104,4 +112,58 @@ class SetEntriesTest {
         val definitelyInSet = setEntries.entries.filter { heap.read(it).isTrue }.mapTo(hashSetOf()) { it.setElement }
         assertEquals(expectedDefInSet, definitelyInSet)
     }
+
+    private object SetBv32ElementInfo :
+        USymbolicCollectionKeyInfo<UExpr<UBv32Sort>, SetRegion<KBitVecValue<KBv32Sort>>> {
+        override fun mapKey(key: UExpr<UBv32Sort>, transformer: UTransformer<*, *>?): UExpr<UBv32Sort> =
+            transformer.apply(key)
+
+        override fun eqSymbolic(ctx: UContext<*>, key1: UExpr<UBv32Sort>, key2: UExpr<UBv32Sort>): UBoolExpr =
+            ctx.mkEq(key1, key2)
+
+        override fun eqConcrete(key1: UExpr<UBv32Sort>, key2: UExpr<UBv32Sort>): Boolean =
+            key1 == key2
+
+        override fun keyToRegion(key: UExpr<UBv32Sort>): SetRegion<KBitVecValue<KBv32Sort>> =
+            if (key is KBitVecValue<UBv32Sort>) {
+                SetRegion.singleton(key)
+            } else {
+                SetRegion.universe()
+            }
+
+        override fun topRegion(): SetRegion<KBitVecValue<KBv32Sort>> = SetRegion.universe()
+        override fun bottomRegion(): SetRegion<KBitVecValue<KBv32Sort>> = SetRegion.empty()
+
+        override fun cmpSymbolicLe(ctx: UContext<*>, key1: UExpr<UBv32Sort>, key2: UExpr<UBv32Sort>) = error("unused")
+        override fun cmpConcreteLe(key1: UExpr<UBv32Sort>, key2: UExpr<UBv32Sort>) = error("unused")
+        override fun keyRangeRegion(from: UExpr<UBv32Sort>, to: UExpr<UBv32Sort>) = error("unused")
+    }
+
+    @Test
+    fun testSetUnion(): Unit = with(ctx) {
+        val maxElement = 100
+        val sets = List(maxElement + 1) { idx ->
+            val setRef = mkConcreteHeapRef(addressCounter.freshAllocatedAddress())
+            val setElement = mkBv(idx)
+            heap.setAddElement(setRef, setElement, setType, SetBv32ElementInfo, trueExpr)
+            setRef
+        }
+
+        sets.reduce { left, right ->
+            heap.setUnion(left, right, setType, bv32Sort, SetBv32ElementInfo, trueExpr)
+            right
+        }
+
+        val resultSet = sets.last()
+
+        for (i in 0..maxElement) {
+            val lvalue = USetEntryLValue(bv32Sort, resultSet, mkBv(i), setType, SetBv32ElementInfo)
+            val elementContains = heap.read(lvalue)
+            assertEquals(trueExpr, elementContains)
+        }
+
+        val lvalue = USetEntryLValue(bv32Sort, resultSet, mkBv(maxElement + 1), setType, SetBv32ElementInfo)
+        val elementContains = heap.read(lvalue)
+        assertEquals(falseExpr, elementContains)
+    }
 }