From ef7bf3da5523e6cfaacb2ee464fdafd7375638d9 Mon Sep 17 00:00:00 2001
From: Helehex <Helehex@gmail.com>
Date: Fri, 9 Aug 2024 16:12:01 -0700
Subject: [PATCH] [External] [stdlib] mark dict entry as destroyed in
 `Dict.pop()` (#44897)

[External] [stdlib] mark dict entry as destroyed in `Dict.pop()`

Fixes https://github.com/modularml/mojo/issues/2756

Co-authored-by: Helehex <Helehex@gmail.com>
Closes modularml/mojo#2796
MODULAR_ORIG_COMMIT_REV_ID: d9ee05077d11b2b4b54bdd48d1a67100ab624eee

Signed-off-by: Manuel Saelices <msaelices@gmail.com>
---
 stdlib/src/collections/dict.mojo       | 11 ++++++++++-
 stdlib/test/collections/test_dict.mojo | 16 ++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/stdlib/src/collections/dict.mojo b/stdlib/src/collections/dict.mojo
index bd91fd17bc..4d37544f55 100644
--- 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
index 579f8bd887..92972b6215 100644
--- 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()