Simpler approach

Certora · Dec 22, 2024 · 3b93541 · 3b93541
1 parent 1ea99a7
commit 3b93541
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 9 deletions.
diff --git a/collect/src/main/kotlin/com/certora/collect/AbstractTreapMap.kt b/collect/src/main/kotlin/com/certora/collect/AbstractTreapMap.kt
@@ -139,14 +139,6 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
             override fun iterator() = entrySequence().iterator()
         }
 
-    override val keys: ImmutableSet<K>
-        get() = object: AbstractSet<K>(), ImmutableSet<K> {
-            override val size get() = this@AbstractTreapMap.size
-            override fun isEmpty() = this@AbstractTreapMap.isEmpty()
-            override operator fun contains(element: K) = containsKey(element)
-            override operator fun iterator() = entrySequence().map { it.key }.iterator()
-        }
-
     override val values: ImmutableCollection<V>
         get() = object: AbstractCollection<V>(), ImmutableCollection<V> {
             override val size get() = this@AbstractTreapMap.size

diff --git a/collect/src/main/kotlin/com/certora/collect/EmptyTreapMap.kt b/collect/src/main/kotlin/com/certora/collect/EmptyTreapMap.kt
@@ -66,7 +66,7 @@ internal class EmptyTreapMap<@Treapable K, V> private constructor() : TreapMap<K
         m.asSequence().map { MapEntry(it.key, null to it.value) }
 
     override val entries: ImmutableSet<Map.Entry<K, V>> get() = persistentSetOf<Map.Entry<K, V>>()
-    override val keys: ImmutableSet<K> get() = persistentSetOf<K>()
+    override val keys: TreapSet<K> get() = treapSetOf<K>()
     override val values: ImmutableCollection<V> get() = persistentSetOf<V>()
 
     @Suppress("Treapability", "UNCHECKED_CAST")

diff --git a/collect/src/main/kotlin/com/certora/collect/HashTreapMap.kt b/collect/src/main/kotlin/com/certora/collect/HashTreapMap.kt
@@ -350,6 +350,8 @@ internal class HashTreapMap<@Treapable K, V>(
         forEachPair { (k, v) -> action(MapEntry(k, v)) }
         right?.forEachEntry(action)
     }
+
+    override val keys: HashTreapSet<K> get() = HashTreapSet(key, next?.toKeyList(), left?.keys, right?.keys)
 }
 
 internal interface KeyValuePairList<K, V> {
@@ -359,6 +361,8 @@ internal interface KeyValuePairList<K, V> {
     operator fun component1() = key
     operator fun component2() = value
 
+    fun toKeyList(): ElementList.More<K> = ElementList.More(key, next?.toKeyList())
+
     class More<K, V>(
         override val key: K,
         override val value: V,

diff --git a/collect/src/main/kotlin/com/certora/collect/SortedTreapMap.kt b/collect/src/main/kotlin/com/certora/collect/SortedTreapMap.kt
@@ -163,4 +163,6 @@ internal class SortedTreapMap<@Treapable K, V>(
         action(this.asEntry())
         right?.forEachEntry(action)
     }
+
+    override val keys: SortedTreapSet<K> get() = SortedTreapSet(key, left?.keys, right?.keys)
 }
diff --git a/collect/src/main/kotlin/com/certora/collect/TreapMap.kt b/collect/src/main/kotlin/com/certora/collect/TreapMap.kt
@@ -12,6 +12,7 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
     override fun remove(key: K, value: @UnsafeVariance V): TreapMap<K, V>
     override fun putAll(m: Map<out K, @UnsafeVariance V>): TreapMap<K, V>
     override fun clear(): TreapMap<K, V>
+    override val keys: TreapSet<K>
 
     /**
         A [PersistentMap.Builder] that produces a [TreapMap].