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);
+ }
}