Better diagnostics for map iteration while mutating

If one obtains a `MutableMap.Entry` from a mutable map / builder, and then removes the underlying value from the map, we currently return `null` from the entry's `value` property, rather than throwing.  Either behavior is allowed by the `MutableMap.Entry` spec (which specifies `IllegalStateException` *or* undefined behavior in this case), but it's much easier to debug this if an exception is thrown.

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

@@ -5,8 +5,18 @@ public class MutableMapEntry<K, V>(
     private val map: MutableMap<K, V>,
     override val key: K
 ) : AbstractMapEntry<K, V>(), MutableMap.MutableEntry<K, V> {
-    @Suppress("UNCHECKED_CAST")
-    override val value: @UnsafeVariance V get() = map.get(key) as V
-    @Suppress("UNCHECKED_CAST")
-    override fun setValue(newValue: V): @UnsafeVariance V = map.put(key, newValue) as V
+
+    override val value: @UnsafeVariance V get() {
+        return map.get(key) ?: run {
+            check(key in map) { "Key '$key' was removed from the map" }
+            @Suppress("UNCHECKED_CAST")
+            null as V
+        }
+    }
+
+    override fun setValue(newValue: V): @UnsafeVariance V {
+        check(key in map) { "Key '$key' was removed from the map" }
+        @Suppress("UNCHECKED_CAST")
+        return map.put(key, newValue) as V
+    }
 }

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

@@ -0,0 +1,44 @@
+package com.certora.collect
+
+import kotlin.test.*
+
+/** Tests for map entries. */
+class MapEntryTest {
+    @Test
+    fun getValue() {
+        val map = treapMapBuilderOf(1 to 2, 3 to 4)
+        val e = map.entries.first()
+        assertEquals(1, e.key)
+        assertEquals(2, e.value)
+        map.remove(1)
+        assertEquals(1, e.key)
+        assertFailsWith<IllegalStateException> { e.value }
+    }
+
+    @Test
+    fun setValue() {
+        val map = treapMapBuilderOf(1 to 2, 3 to 4)
+        val e = map.entries.first()
+        assertEquals(2, e.setValue(5))
+        assertEquals(1, e.key)
+        assertEquals(5, e.value)
+        assertEquals(5, map[1])
+        map.remove(1)
+        assertEquals(1, e.key)
+        assertFailsWith<IllegalStateException> { e.setValue(10) }
+    }
+
+    @Test
+    fun getAndSetNullValue() {
+        val map = treapMapBuilderOf(1 to null, 3 to 4)
+        val e = map.entries.first()
+        assertEquals(1, e.key)
+        assertEquals(null, e.value)
+        assertEquals(null, e.setValue(5))
+        assertEquals(5, e.value)
+        assertEquals(5, map[1])
+        assertEquals(5, e.setValue(null))
+        assertEquals(null, e.value)
+        assertEquals(null, map[1])
+    }
+}