From 62dd98f30521d661e4e991bc0249f204cb03d903 Mon Sep 17 00:00:00 2001
From: Eric Eilebrecht <eric@certora.com>
Date: Fri, 15 Dec 2023 14:56:47 -0800
Subject: [PATCH] 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.
---
 .../com/certora/collect/MutableMapEntry.kt    | 18 ++++++--
 .../com/certora/collect/MapEntryTest.kt       | 44 +++++++++++++++++++
 2 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 collect/src/test/kotlin/com/certora/collect/MapEntryTest.kt

diff --git a/collect/src/main/kotlin/com/certora/collect/MutableMapEntry.kt b/collect/src/main/kotlin/com/certora/collect/MutableMapEntry.kt
index 060c3e4..205ba23 100644
--- a/collect/src/main/kotlin/com/certora/collect/MutableMapEntry.kt
+++ b/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
+    }
 }
diff --git a/collect/src/test/kotlin/com/certora/collect/MapEntryTest.kt b/collect/src/test/kotlin/com/certora/collect/MapEntryTest.kt
new file mode 100644
index 0000000..3afa8a4
--- /dev/null
+++ b/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])
+    }
+}