From 3b9354199aae8ffcee056d645e623b932678f258 Mon Sep 17 00:00:00 2001
From: Eric Eilebrecht <eric@certora.com>
Date: Sun, 22 Dec 2024 07:37:53 -0800
Subject: [PATCH] Simpler approach

---
 .../main/kotlin/com/certora/collect/AbstractTreapMap.kt   | 8 --------
 .../src/main/kotlin/com/certora/collect/EmptyTreapMap.kt  | 2 +-
 .../src/main/kotlin/com/certora/collect/HashTreapMap.kt   | 4 ++++
 .../src/main/kotlin/com/certora/collect/SortedTreapMap.kt | 2 ++
 collect/src/main/kotlin/com/certora/collect/TreapMap.kt   | 1 +
 5 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/collect/src/main/kotlin/com/certora/collect/AbstractTreapMap.kt b/collect/src/main/kotlin/com/certora/collect/AbstractTreapMap.kt
index b0dcf0b..cea55aa 100644
--- 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
index a99cc4a..d877184 100644
--- 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
index cea3982..21b31e7 100644
--- 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
index 0c3617a..eccbda7 100644
--- 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
index 7931468..5559ec4 100644
--- 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].