From 7655cbd5d0b22392fa0add124a47a3724017807c Mon Sep 17 00:00:00 2001
From: Eric Eilebrecht <eric@certora.com>
Date: Thu, 16 May 2024 11:37:13 -0700
Subject: [PATCH] More test

---
 .../com/certora/collect/TreapSetTest.kt       | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/collect/src/test/kotlin/com/certora/collect/TreapSetTest.kt b/collect/src/test/kotlin/com/certora/collect/TreapSetTest.kt
index c6d7065..f3af578 100644
--- a/collect/src/test/kotlin/com/certora/collect/TreapSetTest.kt
+++ b/collect/src/test/kotlin/com/certora/collect/TreapSetTest.kt
@@ -251,6 +251,44 @@ abstract class TreapSetTest {
         assertEquals(s, setOf<TestKey?>(HashTestKey(0), makeKey(0)))
     }
 
+    @Test
+    fun nullQueries() {
+        val withNull = makeSet()
+        withNull.add(makeKey(0))
+        withNull.add(null)
+
+        val withoutNull = makeSet()
+        withNull.add(makeKey(0))
+
+        assertTrue(null in withNull)
+        assertFalse(null in withoutNull)
+
+        assertTrue(withNull.containsAll(treapSetOf(null)))
+        assertFalse(withoutNull.containsAll(treapSetOf(null)))
+
+        assertTrue(withNull.remove(null))
+        assertFalse(withoutNull.remove(null))
+    }
+
+    @Test
+    fun hashQueries() {
+        val withHash = makeSet()
+        withHash.add(makeKey(0))
+        withHash.add(HashTestKey(0))
+
+        val withoutHash = makeSet()
+        withHash.add(makeKey(0))
+
+        assertTrue(HashTestKey(0) in withHash)
+        assertFalse(HashTestKey(0) in withoutHash)
+
+        assertTrue(withHash.containsAll(treapSetOf(HashTestKey(0))))
+        assertFalse(withoutHash.containsAll(treapSetOf(HashTestKey(0))))
+
+        assertTrue(withHash.remove(HashTestKey(0)))
+        assertFalse(withoutHash.remove(HashTestKey(0)))
+    }
+
     @Test
     fun retainAll() {
         val b = makeBaseline()