From 0eb577ac0173e4777972efd60118c342b165d8b3 Mon Sep 17 00:00:00 2001 From: gerben-stavenga <54682677+gerben-stavenga@users.noreply.github.com> Date: Tue, 14 Mar 2023 17:25:59 -0700 Subject: [PATCH] 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 --- .../framework/aptos-stdlib/doc/simple_map.md | 44 +++++++++++++++++ .../aptos-stdlib/sources/simple_map.move | 47 +++++++++++++++++++ 2 files changed, 91 insertions(+) diff --git a/aptos-move/framework/aptos-stdlib/doc/simple_map.md b/aptos-move/framework/aptos-stdlib/doc/simple_map.md index 022c772cf7bb8..561e00eb0a00a 100644 --- 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 + + + + +## Function `upsert` + +Insert key/value pair or update an existing key to a new value + + +
public fun upsert<Key: store, Value: store>(map: &mut simple_map::SimpleMap<Key, Value>, key: Key, value: Value): (option::Option<Key>, option::Option<Value>)
+
+ + + +
+Implementation + + +
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())
+}
+
+ + +
diff --git a/aptos-move/framework/aptos-stdlib/sources/simple_map.move b/aptos-move/framework/aptos-stdlib/sources/simple_map.move index 0d9f78100163f..70a6cd07931e4 100644 --- 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( + map: &mut SimpleMap, + key: Key, + value: Value + ): (std::option::Option, std::option::Option) { + 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( @@ -181,4 +204,28 @@ module aptos_std::simple_map { destroy_empty(map); } + + #[test] + public fun upsert_test() { + let map = create(); + // test adding 3 elements using upsert + upsert(&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); + } }