[External] [stdlib] mark dict entry as destroyed in Dict.pop() (#44…

…897) [External] [stdlib] mark dict entry as destroyed in `Dict.pop()` Fixes modular#2756 Co-authored-by: Helehex <[email protected]> Closes modular#2796 MODULAR_ORIG_COMMIT_REV_ID: d9ee05077d11b2b4b54bdd48d1a67100ab624eee Signed-off-by: Manuel Saelices <[email protected]>
msaelices · Sep 10, 2024 · ef7bf3d · ef7bf3d
1 parent 762e225
commit ef7bf3d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/stdlib/src/collections/dict.mojo b/stdlib/src/collections/dict.mojo
@@ -217,6 +217,15 @@ struct DictEntry[K: KeyElement, V: CollectionElement](
         self.key = other.key
         self.value = other.value
 
+    fn reap_value(owned self) -> V:
+        """Take the value from an owned entry.
+
+        Returns:
+            The value of the entry.
+        """
+        __mlir_op.`lit.ownership.mark_destroyed`(__get_mvalue_as_litref(self))
+        return self.value^
+
 
 alias _EMPTY = -1
 alias _REMOVED = -2
@@ -814,7 +823,7 @@ struct Dict[K: KeyElement, V: CollectionElement](
             var entry_value = entry[].unsafe_take()
             entry[] = None
             self.size -= 1
-            return entry_value.value^
+            return entry_value^.reap_value()
         raise "KeyError"
 
     fn popitem(inout self) raises -> DictEntry[K, V]:

diff --git a/stdlib/test/collections/test_dict.mojo b/stdlib/test/collections/test_dict.mojo
@@ -538,6 +538,21 @@ def test_dict_popitem():
         _ = dict.popitem()
 
 
+def test_pop_string_values():
+    var dict = Dict[String, String]()
+    dict["mojo"] = "lang"
+    dict["max"] = "engine"
+    dict["a"] = ""
+    dict[""] = "a"
+
+    assert_equal(dict.pop("mojo"), "lang")
+    assert_equal(dict.pop("max"), "engine")
+    assert_equal(dict.pop("a"), "")
+    assert_equal(dict.pop(""), "a")
+    with assert_raises(contains="KeyError"):
+        _ = dict.pop("absent")
+
+
 fn test_clear() raises:
     var some_dict = Dict[String, Int]()
     some_dict["key"] = 1
@@ -594,6 +609,7 @@ def main():
     test_owned_kwargs_dict()
     test_bool_conversion()
     test_find_get()
+    test_pop_string_values()
     test_clear()
     test_init_initial_capacity()
     test_dict_setdefault()