Add upsert and destroy also to vector

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 `remove`](#0x1_simple_map_remove)
 -  [Function `find`](#0x1_simple_map_find)
 -  [Specification](#@Specification_1)
@@ -320,6 +321,44 @@ Map key is not found
 
 
 
+</details>
+
+<a name="0x1_simple_map_upsert"></a>
+
+## Function `upsert`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="simple_map.md#0x1_simple_map_upsert">upsert</a>&lt;Key: <b>copy</b>, drop, store, Value: <b>copy</b>, drop, 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)
+</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 + drop + <b>copy</b>, Value: store + drop + <b>copy</b>&gt;(
+    map: &<b>mut</b> <a href="simple_map.md#0x1_simple_map_SimpleMap">SimpleMap</a>&lt;Key, Value&gt;,
+    key: &Key,
+    value: &Value
+) {
+    <b>let</b> len = std::vector::length(&map.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_mut">vector::borrow_mut</a>(&<b>mut</b> map.data, i);
+        <b>if</b> (&element.key == key) {
+            element.value = *value;
+            <b>return</b>
+        };
+        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: *key, value: *value });
+}
+</code></pre>
+
+
+
 </details>
 
 <a name="0x1_simple_map_remove"></a>

aptos-move/framework/aptos-stdlib/sources/simple_map.move

@@ -77,6 +77,24 @@ module aptos_std::simple_map {
         vector::push_back(&mut map.data, Element { key, value });
     }
 
+    public fun upsert<Key: store + drop + copy, Value: store + drop + copy>(
+        map: &mut SimpleMap<Key, Value>,
+        key: &Key,
+        value: &Value
+    ) {
+        let len = std::vector::length(&map.data);
+        let i = 0;
+        while (i < len) {
+            let element = vector::borrow_mut(&mut map.data, i);
+            if (&element.key == key) {
+                element.value = *value;
+                return
+            };
+            i = i + 1;
+        };
+        vector::push_back(&mut map.data, Element { key: *key, value: *value });
+    }
+
     public fun remove<Key: store, Value: store>(
         map: &mut SimpleMap<Key, Value>,
         key: &Key,
@@ -158,4 +176,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);
+    }
 }