Fix treap map/list bugs, tighten Treapability rule, expand tests

collect/src/main/kotlin/com/certora/collect/AbstractTreapMap.kt

@@ -167,7 +167,7 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
         var newThis = this as TreapMap<K, V>
         for ((k, v) in m.entries) {
             if (k in this) {
-                newThis = newThis + (k to merger(k, this[k]!!, v))
+                newThis = newThis + (k to merger(k, storedValue(k), v))
             } else {
                 newThis = newThis + (k to v)
             }
@@ -191,12 +191,16 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
         var newThis = clear()
         for ((k, v) in m.entries) {
             if (k in this) {
-                newThis = newThis + (k to merger(k, this[k]!!, v))
+                newThis = newThis + (k to merger(k, storedValue(k), v))
             }
         }
         return newThis
     }
 
+    // Map.get returns V?; a key that is present may map to null.
+    @Suppress("UNCHECKED_CAST")
+    private fun storedValue(key: K): V = this[key] as V
+
     /**
         Merges the entries in `m` with the entries in this AbstractTreapMap, applying the "merger" function to get the
         new values for each key.
@@ -492,6 +496,19 @@ internal fun <@Treapable K, V, @Treapable S : AbstractTreapMap<K, V, S>> S?.merg
         mergeWithImpl(that, mode, shallowMerge)
     }
 
+/**
+    True when the two treaps together are smaller than about 2^[parallelThresholdLog2].
+
+    Each side is compared to half of that budget. A threshold of 0 compares each side to 2^0, since only empty treaps
+    are smaller than that.
+ */
+private fun <@Treapable T, S : Treap<T, S>> Pair<S?, S?>.isSmallerThanMergeThreshold(
+    parallelThresholdLog2: Int
+): Boolean {
+    val sizeLog2 = if (parallelThresholdLog2 == 0) { 0 } else { parallelThresholdLog2 - 1 }
+    return first.isApproximatelySmallerThanLog2(sizeLog2) && second.isApproximatelySmallerThanLog2(sizeLog2)
+}
+
 internal fun <@Treapable K, V, @Treapable S : AbstractTreapMap<K, V, S>> S?.parallelMergeWith(
     that: S?,
     mode: MergeMode,
@@ -500,10 +517,7 @@ internal fun <@Treapable K, V, @Treapable S : AbstractTreapMap<K, V, S>> S?.para
 ): S? =
     maybeForking(
         this to that,
-        {
-            it.first.isApproximatelySmallerThanLog2(parallelThresholdLog2 - 1) &&
-            it.second.isApproximatelySmallerThanLog2(parallelThresholdLog2 - 1)
-        }
+        { it.isSmallerThanMergeThreshold(parallelThresholdLog2) }
     ) {
         mergeWithImpl(that, mode, shallowMerge)
     }
@@ -565,10 +579,7 @@ internal fun <@Treapable K, V, @Treapable S : AbstractTreapMap<K, V, S>> S?.para
 ): S? =
     maybeForking(
         this to that,
-        {
-            it.first.isApproximatelySmallerThanLog2(parallelThresholdLog2 - 1) &&
-            it.second.isApproximatelySmallerThanLog2(parallelThresholdLog2 - 1)
-        }
+        { it.isSmallerThanMergeThreshold(parallelThresholdLog2) }
     ) {
         unionWithImpl(that, shallowUnion)
     }
@@ -618,10 +629,7 @@ internal fun <@Treapable K, V, @Treapable S : AbstractTreapMap<K, V, S>> S?.para
 ): S? =
     maybeForking(
         this to that,
-        {
-            it.first.isApproximatelySmallerThanLog2(parallelThresholdLog2 - 1) &&
-            it.second.isApproximatelySmallerThanLog2(parallelThresholdLog2 - 1)
-        }
+        { it.isSmallerThanMergeThreshold(parallelThresholdLog2) }
     ) {
         intersectWithImpl(that, shallowIntersect)
     }

collect/src/main/kotlin/com/certora/collect/HashTreapMap.kt

@@ -51,9 +51,10 @@ internal class HashTreapMap<@Treapable K, V>(
         }
         if (mode == MergeMode.UNION) {
             t2?.forEachPair { (k, v2) ->
-                val v1 = t1?.shallowGetValue(k)
-                if (v1 == null) {
-                    val v = merger(k, v1, v2)
+                // shallowGetValue returns null both for a missing key and for a stored null.
+                // Only keys absent from t1 belong in this loop; a stored null was already merged above.
+                if (t1?.shallowContainsKey(k) != true) {
+                    val v = merger(k, null, v2)
                     if (v != null) {
                         newPairs = KeyValuePairList.More(k, v, newPairs)
                     }

collect/src/main/kotlin/com/certora/collect/TreapListBuilder.kt

@@ -3,7 +3,7 @@ package com.certora.collect
 internal class TreapListBuilder<E>(
     private var list: TreapList<E>
 ) : TreapList.Builder<E>, java.io.Serializable, AbstractMutableList<E>() {
-    override val size: Int = list.size
+    override val size: Int get() = list.size
     override fun isEmpty(): Boolean = list.isEmpty()
 
     override fun get(index: Int): E = list[index]
@@ -57,23 +57,27 @@ internal class TreapListBuilder<E>(
             val i = nextIndex()
             add(i, element)
             listIterator = list.listIterator(i + 1)
+            lastReturnedIndex = -1
         }
 
         override fun set(element: E) {
             val i = lastReturnedIndex
             if (i == -1) {
-                throw IllegalStateException("set() called before next()")
+                throw IllegalStateException("set() called before next() or previous(), or after add() or remove()")
             }
+            val cursor = listIterator.nextIndex()
             list = list.set(i, element)
+            listIterator = list.listIterator(cursor)
         }
 
         override fun remove() {
             val i = lastReturnedIndex
             if (i == -1) {
-                throw IllegalStateException("remove() called before next()")
+                throw IllegalStateException("remove() called before next() or previous(), or after add() or remove()")
             }
             removeAt(i)
             listIterator = list.listIterator(i)
+            lastReturnedIndex = -1
         }
     }
 

collect/src/main/kotlin/com/certora/collect/TreapListNode.kt

@@ -27,7 +27,7 @@ internal class TreapListNode<E> private constructor(
     override val size: Int = addExact(1, addExact(left?.size?:0, right?.size?:0))
 ) : TreapList<E>, java.io.Serializable {
 
-    constructor(elem: E) : this(elem, priority = Random.Default.nextInt(), size = 1)
+    constructor(elem: E) : this(elem, priority = nextPriority(), size = 1)
 
     override fun hashCode() = computeHashCode(initial = 1)
     override fun toString() = joinToString(", ", "[", "]")
@@ -387,6 +387,32 @@ internal class TreapListNode<E> private constructor(
     }
 
     companion object {
+        /**
+            Per-thread replacement for [Random.Default] when drawing node priorities.
+
+            With nothing installed on this thread, the production setting, priorities come from [Random.Default].
+            Tests install a seeded [Random] on the test thread, so one test cannot see another test's draws. A
+            fork/join worker does not inherit the installation.
+         */
+        private val priorityRandom = ThreadLocal<Random?>()
+
+        internal fun setTreapListPriorityRandom(random: Random?) {
+            if (random == null) {
+                priorityRandom.remove()
+            } else {
+                priorityRandom.set(random)
+            }
+        }
+
+        private fun nextPriority(): Int {
+            val override = priorityRandom.get()
+            return if (override == null) {
+                Random.Default.nextInt()
+            } else {
+                override.nextInt()
+            }
+        }
+
         private infix fun <E> TreapListNode<E>?.append(that: TreapListNode<E>?): TreapListNode<E>? = when {
             this == null -> that
             that == null -> this
@@ -417,7 +443,7 @@ internal class TreapListNode<E> private constructor(
 
                 // Start the next element, and loop until we find one with a higher priority than [upperPri]
                 var nextElem = elems.next()
-                var nextPri = Random.Default.nextInt()
+                var nextPri = nextPriority()
                 while (true) {
                     when {
                         nextPri > upperPri -> {
@@ -432,7 +458,7 @@ internal class TreapListNode<E> private constructor(
                                 return Result(thisNode)
                             }
                             nextElem = elems.next()
-                            nextPri = Random.Default.nextInt()
+                            nextPri = nextPriority()
                         }
                         else -> {
                             // The priority is below the current node; construct a subtree on the right.
@@ -450,7 +476,7 @@ internal class TreapListNode<E> private constructor(
             }
 
             // Build the whole list
-            return buildLowerPri(Int.MAX_VALUE, elems.next(), Random.Default.nextInt()).node
+            return buildLowerPri(Int.MAX_VALUE, elems.next(), nextPriority()).node
         }
 
         internal tailrec fun <E> TreapListNode<E>?.isApproximatelySmallerThanLog2(sizeLog2: Int): Boolean = when {

collect/src/test/kotlin/com/certora/collect/HashTreapMapTest.kt

@@ -2,6 +2,7 @@ package com.certora.collect
 
 import com.certora.collect.*
 import kotlinx.serialization.DeserializationStrategy
+import kotlin.test.*
 
 /** Tests for [HashTreapMap]. */
 class HashTreapMapTest: TreapMapTest() {
@@ -15,4 +16,78 @@ class HashTreapMapTest: TreapMapTest() {
     override fun makeBaselineOfInts(): MutableMap<Int?, Int?> = HashMap()
     override fun getBaseDeserializer(): DeserializationStrategy<*>? = null
     override fun getDeserializer(): DeserializationStrategy<*>? = null
+
+    @Test
+    fun hashCollisions() {
+        val a = HashTestKey(1, code = 7)
+        val b = HashTestKey(2, code = 7)
+        val c = HashTestKey(3, code = 8)
+        val d = HashTestKey(4, code = 7)
+        val left = treapMapOf(a to 1, b to 2, c to 3)
+        val right = treapMapOf(b to 20, c to 30, d to 4)
+
+        assertEquals(mapOf(a to 1, b to 22, c to 33, d to 4), left.merge(right) { _, v1, v2 -> (v1 ?: 0) + (v2 ?: 0) })
+        assertEquals(mapOf(b to 2, c to 3), left.intersect(right) { _, v1, _ -> v1 })
+        assertEquals(mapOf(a to 1, b to 20, c to 30, d to 4), left.union(right) { _, _, v2 -> v2 })
+        assertEquals(
+            setOf(MapEntry(a, 1 to null), MapEntry(b, 2 to 20), MapEntry(c, 3 to 30), MapEntry(d, null to 4)),
+            left.zip(right).toSet()
+        )
+
+        val updated = left.updateEntry(b, 5) { current, extra -> current!! + extra }
+        assertEquals(7, updated[b])
+        assertEquals(1, updated[a])
+        val removed = left.updateEntry(a, 0) { _, _ -> null }
+        assertFalse(a in removed)
+        assertEquals(2, removed[b])
+        val inserted = left.updateEntry(d, 9) { current, extra ->
+            assertNull(current)
+            extra
+        }
+        assertEquals(9, inserted[d])
+
+        assertEquals(mapOf(b to 20, c to 30), left.updateValues { key, value ->
+            if (key == a) {
+                null
+            } else {
+                value * 10
+            }
+        })
+        assertFalse(b in left.remove(b, 2))
+        assertEquals(left, left.remove(b, 99))
+
+        val onlyLeft = HashTestKey(8, code = 42)
+        val onlyRight = HashTestKey(9, code = 99)
+        val withDisjointHashes = treapMapOf(a to 1, onlyLeft to 5).merge(treapMapOf(onlyRight to 6, b to 2)) { _, v1, v2 ->
+            (v1 ?: 0) + (v2 ?: 0)
+        }
+        assertEquals(mapOf(a to 1, onlyLeft to 5, onlyRight to 6, b to 2), withDisjointHashes)
+
+        val sameValue = treapMapOf(a to "x", b to "x")
+        assertEquals(mapOf(b to "x"), sameValue.remove(a, "x"))
+        assertEquals(mapOf(a to "x"), sameValue.remove(b, "x"))
+
+        assertSame(left, left.updateEntry(a, 0) { current, _ -> current })
+        assertSame(left, left.updateEntry(b, 0) { current, _ -> current })
+        val missingInBucket = HashTestKey(5, code = 7)
+        assertSame(left, left.updateEntry(missingInBucket, 1) { _, _ -> null })
+
+        assertEquals(6, left.mapReduce({ _, value -> value }, Int::plus))
+    }
+
+    @Test
+    fun mergeUnionWithNullValueInCollidingBucketMergesOnce() {
+        val key = HashTestKey(1, code = 7)
+        val other = HashTestKey(2, code = 7)
+        var calls = 0
+        val merged = treapMapOf(key to null as String?, other to "keep").merge(
+            treapMapOf(key to "b")
+        ) { _, left, right ->
+            calls += 1
+            right ?: left
+        }
+        assertEquals(2, calls)
+        val expected: Map<HashTestKey, String?> = mapOf(key to "b", other to "keep")
+        assertEquals(expected, merged)
+    }
 }

collect/src/test/kotlin/com/certora/collect/HashTreapSetTest.kt

@@ -2,11 +2,29 @@ package com.certora.collect
 
 import com.certora.collect.*
 import kotlinx.serialization.DeserializationStrategy
+import kotlin.test.*
 
 /** Tests for [HashTreapSet]. */
 class HashTreapSetTest : TreapSetTest() {
     override fun makeKey(value: Int, code: Int) = HashTestKey(value, code)
     override fun makeBaseline(): MutableSet<TestKey?> = HashSet()
     override fun getBaseDeserializer(): DeserializationStrategy<*>? = null
     override fun getDeserializer(): DeserializationStrategy<*>? = null
+
+    @Test
+    fun hashCollisions() {
+        val a = HashTestKey(1, code = 7)
+        val b = HashTestKey(2, code = 7)
+        val c = HashTestKey(3, code = 8)
+        val d = HashTestKey(4, code = 7)
+        val left = treapSetOf(a, b, c)
+        val right = treapSetOf(b, c, d)
+        assertEquals(setOf(a, b, c, d), left union right)
+        assertEquals(setOf(b, c), left intersect right)
+        assertEquals(setOf(a), left subtract right)
+        assertTrue(left.containsAny(listOf(d, b)))
+        assertSame(b, left.findEqual(HashTestKey(2, code = 7)))
+        assertFalse(b in left.remove(b))
+        assertEquals(setOf(a, c), left.removeAll { it == b })
+    }
 }

collect/src/test/kotlin/com/certora/collect/SortedTreapMapTest.kt

@@ -73,5 +73,16 @@ class SortedTreapMapTest: TreapMapTest() {
         assertEquals(1000, map.lowerKey(1001))
         assertNull(map.ceilingKey(1001))
         assertNull(map.higherKey(1001))
+
+        assertEquals(0 to 1, map.firstEntry()!!.toPair())
+        assertEquals(1000 to 1001, map.lastEntry()!!.toPair())
+        assertEquals(4 to 5, map.floorEntry(4)!!.toPair())
+        assertEquals(4 to 5, map.floorEntry(5)!!.toPair())
+        assertEquals(2 to 3, map.lowerEntry(4)!!.toPair())
+        assertEquals(4 to 5, map.ceilingEntry(4)!!.toPair())
+        assertEquals(6 to 7, map.ceilingEntry(5)!!.toPair())
+        assertEquals(6 to 7, map.higherEntry(4)!!.toPair())
+        assertNull(treapMapOf<Int, Int>().firstEntry())
+        assertNull(treapMapOf<Int, Int>().lastEntry())
     }
 }

collect/src/test/kotlin/com/certora/collect/SortedTreapSetTest.kt

@@ -3,6 +3,7 @@ package com.certora.collect
 import com.certora.collect.*
 import kotlinx.serialization.DeserializationStrategy
 import java.util.TreeSet
+import kotlin.test.*
 
 /** Tests for [SortedTreapSet]. */
 @Suppress("UNCHECKED_CAST")
@@ -12,4 +13,10 @@ class SortedTreapSetTest : TreapSetTest() {
     override fun makeBaseline(): MutableSet<TestKey?> = TreeSet()
     override fun getBaseDeserializer(): DeserializationStrategy<*>? = null
     override fun getDeserializer(): DeserializationStrategy<*>? = null
+
+    @Test
+    fun iteratesInSortedOrder() {
+        val set = treapSetOf(makeKey(3), makeKey(1), makeKey(4), makeKey(2))
+        assertEquals(listOf(1, 2, 3, 4).map { makeKey(it) }, set.toList())
+    }
 }

collect/src/test/kotlin/com/certora/collect/TreapListPriorityRule.kt

@@ -0,0 +1,30 @@
+package com.certora.collect
+
+import kotlin.random.Random
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+    Installs a fresh [Random] for list priorities around each test method.
+
+    The generator is stored on the running thread, so test order does not matter and two test threads do not share
+    draws. Any test that builds a [TreapList] needs this rule. Without it, that test still draws from
+    [Random.Default].
+ */
+class TreapListPriorityRule : TestRule {
+    override fun apply(base: Statement, description: Description): Statement = object : Statement() {
+        override fun evaluate() {
+            TreapListNode.setTreapListPriorityRandom(Random(TREAP_LIST_PRIORITY_SEED))
+            try {
+                base.evaluate()
+            } finally {
+                TreapListNode.setTreapListPriorityRandom(null)
+            }
+        }
+    }
+
+    internal companion object {
+        const val TREAP_LIST_PRIORITY_SEED = 1
+    }
+}