Add upsert and destroy also to vector (#6860)

* Add upsert and destroy also to vector * Improve libs * Remove drop version * remove destroy * remove to_vec, split_element * generate docs * generate docs * generate docs * remove std * Remove drop and copy * Add drop constrain for key * Drop drop constrain for key * generate doc
aptos-labs · Mar 15, 2023 · 0eb577a · 0eb577a
1 parent 00d4a05
commit 0eb577a
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 0 deletions.
diff --git a/aptos-move/framework/aptos-stdlib/doc/simple_map.md b/aptos-move/framework/aptos-stdlib/doc/simple_map.md
@@ -21,6 +21,7 @@ This module provides a solution for sorted maps, that is it has the properties t
 -  [Function `contains_key`](#0x1_simple_map_contains_key)
 -  [Function `destroy_empty`](#0x1_simple_map_destroy_empty)
 -  [Function `add`](#0x1_simple_map_add)
+-  [Function `upsert`](#0x1_simple_map_upsert)
 -  [Function `to_vec_pair`](#0x1_simple_map_to_vec_pair)
 -  [Function `remove`](#0x1_simple_map_remove)
 -  [Function `find`](#0x1_simple_map_find)
@@ -321,6 +322,49 @@ Map key is not found
 
 
 
+</details>
+
+<a name="0x1_simple_map_upsert"></a>
+
+## Function `upsert`
+
+Insert key/value pair or update an existing key to a new value
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="simple_map.md#0x1_simple_map_upsert">upsert</a>&lt;Key: store, Value: store&gt;(map: &<b>mut</b> <a href="simple_map.md#0x1_simple_map_SimpleMap">simple_map::SimpleMap</a>&lt;Key, Value&gt;, key: Key, value: Value): (<a href="../../move-stdlib/doc/option.md#0x1_option_Option">option::Option</a>&lt;Key&gt;, <a href="../../move-stdlib/doc/option.md#0x1_option_Option">option::Option</a>&lt;Value&gt;)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="simple_map.md#0x1_simple_map_upsert">upsert</a>&lt;Key: store, Value: store&gt;(
+    map: &<b>mut</b> <a href="simple_map.md#0x1_simple_map_SimpleMap">SimpleMap</a>&lt;Key, Value&gt;,
+    key: Key,
+    value: Value
+): (std::option::Option&lt;Key&gt;, std::option::Option&lt;Value&gt;) {
+    <b>let</b> data = &<b>mut</b> map.data;
+    <b>let</b> len = <a href="../../move-stdlib/doc/vector.md#0x1_vector_length">vector::length</a>(data);
+    <b>let</b> i = 0;
+    <b>while</b> (i &lt; len) {
+        <b>let</b> element = <a href="../../move-stdlib/doc/vector.md#0x1_vector_borrow">vector::borrow</a>(data, i);
+        <b>if</b> (&element.key == &key) {
+            <a href="../../move-stdlib/doc/vector.md#0x1_vector_push_back">vector::push_back</a>(data, <a href="simple_map.md#0x1_simple_map_Element">Element</a> { key, value});
+            <a href="../../move-stdlib/doc/vector.md#0x1_vector_swap">vector::swap</a>(data, i, len);
+            <b>let</b> <a href="simple_map.md#0x1_simple_map_Element">Element</a> { key, value } = <a href="../../move-stdlib/doc/vector.md#0x1_vector_pop_back">vector::pop_back</a>(data);
+            <b>return</b> (std::option::some(key), std::option::some(value))
+        };
+        i = i + 1;
+    };
+    <a href="../../move-stdlib/doc/vector.md#0x1_vector_push_back">vector::push_back</a>(&<b>mut</b> map.data, <a href="simple_map.md#0x1_simple_map_Element">Element</a> { key, value });
+    (std::option::none(), std::option::none())
+}
+</code></pre>
+
+
+
 </details>
 
 <a name="0x1_simple_map_to_vec_pair"></a>

diff --git a/aptos-move/framework/aptos-stdlib/sources/simple_map.move b/aptos-move/framework/aptos-stdlib/sources/simple_map.move
@@ -77,6 +77,29 @@ module aptos_std::simple_map {
         vector::push_back(&mut map.data, Element { key, value });
     }
 
+    /// Insert key/value pair or update an existing key to a new value
+    public fun upsert<Key: store, Value: store>(
+        map: &mut SimpleMap<Key, Value>,
+        key: Key,
+        value: Value
+    ): (std::option::Option<Key>, std::option::Option<Value>) {
+        let data = &mut map.data;
+        let len = vector::length(data);
+        let i = 0;
+        while (i < len) {
+            let element = vector::borrow(data, i);
+            if (&element.key == &key) {
+                vector::push_back(data, Element { key, value});
+                vector::swap(data, i, len);
+                let Element { key, value } = vector::pop_back(data);
+                return (std::option::some(key), std::option::some(value))
+            };
+            i = i + 1;
+        };
+        vector::push_back(&mut map.data, Element { key, value });
+        (std::option::none(), std::option::none())
+    }
+
     /// Transform the map into two vectors with the keys and values respectively
     /// Primarily used to destroy a map
     public fun to_vec_pair<Key: store, Value: store>(
@@ -181,4 +204,28 @@ module aptos_std::simple_map {
 
         destroy_empty(map);
     }
+
+    #[test]
+    public fun upsert_test() {
+        let map = create<u64, u64>();
+        // test adding 3 elements using upsert
+        upsert<u64, u64>(&mut map, 1, 1 );
+        upsert(&mut map, 2, 2 );
+        upsert(&mut map, 3, 3 );
+
+        assert!(length(&map) == 3, 0);
+        assert!(contains_key(&map, &1), 1);
+        assert!(contains_key(&map, &2), 2);
+        assert!(contains_key(&map, &3), 3);
+        assert!(borrow(&map, &1) == &1, 4);
+        assert!(borrow(&map, &2) == &2, 5);
+        assert!(borrow(&map, &3) == &3, 6);
+
+        // change mapping 1->1 to 1->4
+        upsert(&mut map, 1, 4 );
+
+        assert!(length(&map) == 3, 7);
+        assert!(contains_key(&map, &1), 8);
+        assert!(borrow(&map, &1) == &4, 9);
+    }
 }