From cbad8dc60704aa870d840bf42dae03a1eaa6ef0b Mon Sep 17 00:00:00 2001
From: Valentyn Sobol <napster_1997@mail.ru>
Date: Tue, 24 Dec 2024 16:19:55 +0300
Subject: [PATCH] Fix composer handling in set union adapter cache (#228)

* Fix composer handling in set union adapter cache

* minor
---
 .../set/primitive/USymbolicSetUnionAdapter.kt | 32 ++++++++++++++++---
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/usvm-core/src/main/kotlin/org/usvm/collection/set/primitive/USymbolicSetUnionAdapter.kt b/usvm-core/src/main/kotlin/org/usvm/collection/set/primitive/USymbolicSetUnionAdapter.kt
index 20dc6bdf9d..cb30021940 100644
--- a/usvm-core/src/main/kotlin/org/usvm/collection/set/primitive/USymbolicSetUnionAdapter.kt
+++ b/usvm-core/src/main/kotlin/org/usvm/collection/set/primitive/USymbolicSetUnionAdapter.kt
@@ -19,8 +19,9 @@ import org.usvm.memory.USymbolicCollectionKeyInfo
 import org.usvm.memory.UUpdateNode
 import org.usvm.memory.UWritableMemory
 import org.usvm.memory.key.UHeapRefKeyInfo
-import org.usvm.uctx
 import org.usvm.regions.Region
+import org.usvm.uctx
+import java.lang.ref.WeakReference
 
 sealed class USymbolicSetUnionAdapter<
     SetType, SrcKey, DstKey,
@@ -46,13 +47,13 @@ sealed class USymbolicSetUnionAdapter<
          * */
         val prevIncludesSymbolicallyCache = lastIncludesSymbolicallyCheck
         if (prevIncludesSymbolicallyCache != null) {
-            if (prevIncludesSymbolicallyCache.key == srcKey) {
+            if (prevIncludesSymbolicallyCache.containsCachedValue(srcKey, composer)) {
                 return prevIncludesSymbolicallyCache.result
             }
         }
 
         return setOfKeys.read(srcKey, composer).also {
-            lastIncludesSymbolicallyCheck = IncludesSymbolicallyCache(srcKey, it)
+            lastIncludesSymbolicallyCheck = IncludesSymbolicallyCache(srcKey, composer, it)
         }
     }
 
@@ -66,7 +67,30 @@ sealed class USymbolicSetUnionAdapter<
 
     abstract override fun collectSetElements(elements: USymbolicSetElementsCollector.Elements<DstKey>)
 
-    private data class IncludesSymbolicallyCache<Key>(val key: Key, val result: UBoolExpr)
+    private data class IncludesSymbolicallyCache<Key>(
+        val key: WeakReference<Key>,
+        val composer: WeakReference<UComposer<*, *>>?,
+        val result: UBoolExpr
+    ) {
+        constructor(key: Key, composer: UComposer<*, *>?, result: UBoolExpr) :
+                this(WeakReference(key), composer?.let { WeakReference(it) }, result)
+
+        fun containsCachedValue(key: Key, composer: UComposer<*, *>?): Boolean {
+            if (!this.key.equalTo(key)) return false
+
+            val thisComposer = this.composer ?: return composer == null
+            if (composer == null) return false
+
+            return thisComposer.equalTo(composer)
+        }
+
+        companion object {
+            private fun <T> WeakReference<T>.equalTo(other: T): Boolean {
+                val value = get() ?: return false
+                return value == other
+            }
+        }
+    }
 }
 
 class UAllocatedToAllocatedSymbolicSetUnionAdapter<SetType, ElemSort : USort>(