diff --git a/aptos-move/framework/aptos-framework/doc/big_ordered_map.md b/aptos-move/framework/aptos-framework/doc/big_ordered_map.md
index 1f97dcba8e4c5..13403c6710107 100644
--- a/aptos-move/framework/aptos-framework/doc/big_ordered_map.md
+++ b/aptos-move/framework/aptos-framework/doc/big_ordered_map.md
@@ -56,6 +56,11 @@ A set of inline utility methods is provided instead, to provide guaranteed valid
- [Function `upsert`](#0x1_big_ordered_map_upsert)
- [Function `remove`](#0x1_big_ordered_map_remove)
- [Function `remove_or_none`](#0x1_big_ordered_map_remove_or_none)
+- [Function `modify`](#0x1_big_ordered_map_modify)
+- [Function `modify_and_return`](#0x1_big_ordered_map_modify_and_return)
+- [Function `modify_or_add`](#0x1_big_ordered_map_modify_or_add)
+- [Function `modify_if_present`](#0x1_big_ordered_map_modify_if_present)
+- [Function `modify_if_present_and_return`](#0x1_big_ordered_map_modify_if_present_and_return)
- [Function `add_all`](#0x1_big_ordered_map_add_all)
- [Function `pop_front`](#0x1_big_ordered_map_pop_front)
- [Function `pop_back`](#0x1_big_ordered_map_pop_back)
@@ -66,6 +71,7 @@ A set of inline utility methods is provided instead, to provide guaranteed valid
- [Function `contains`](#0x1_big_ordered_map_contains)
- [Function `borrow`](#0x1_big_ordered_map_borrow)
- [Function `get`](#0x1_big_ordered_map_get)
+- [Function `get_and_map`](#0x1_big_ordered_map_get_and_map)
- [Function `borrow_mut`](#0x1_big_ordered_map_borrow_mut)
- [Function `borrow_front`](#0x1_big_ordered_map_borrow_front)
- [Function `borrow_back`](#0x1_big_ordered_map_borrow_back)
@@ -85,6 +91,7 @@ A set of inline utility methods is provided instead, to provide guaranteed valid
- [Function `iter_borrow_key`](#0x1_big_ordered_map_iter_borrow_key)
- [Function `iter_borrow`](#0x1_big_ordered_map_iter_borrow)
- [Function `iter_borrow_mut`](#0x1_big_ordered_map_iter_borrow_mut)
+- [Function `iter_modify`](#0x1_big_ordered_map_iter_modify)
- [Function `iter_remove`](#0x1_big_ordered_map_iter_remove)
- [Function `iter_next`](#0x1_big_ordered_map_iter_next)
- [Function `iter_prev`](#0x1_big_ordered_map_iter_prev)
@@ -130,6 +137,7 @@ A set of inline utility methods is provided instead, to provide guaranteed valid
- [Function `internal_find`](#@Specification_1_internal_find)
- [Function `contains`](#@Specification_1_contains)
- [Function `borrow`](#@Specification_1_borrow)
+ - [Function `get_and_map`](#@Specification_1_get_and_map)
- [Function `borrow_mut`](#@Specification_1_borrow_mut)
- [Function `borrow_front`](#@Specification_1_borrow_front)
- [Function `borrow_back`](#@Specification_1_borrow_back)
@@ -143,6 +151,7 @@ A set of inline utility methods is provided instead, to provide guaranteed valid
- [Function `iter_borrow_key`](#@Specification_1_iter_borrow_key)
- [Function `iter_borrow`](#@Specification_1_iter_borrow)
- [Function `iter_borrow_mut`](#@Specification_1_iter_borrow_mut)
+ - [Function `iter_modify`](#@Specification_1_iter_modify)
- [Function `iter_next`](#@Specification_1_iter_next)
- [Function `iter_prev`](#@Specification_1_iter_prev)
- [Function `validate_dynamic_size_and_init_max_degrees`](#@Specification_1_validate_dynamic_size_and_init_max_degrees)
@@ -1140,6 +1149,151 @@ Returns none if there is no entry for key.
+
+
+
+
+## Function `modify`
+
+Modifies the element in the map via calling f.
+Aborts if element doesn't exist
+
+
+public fun modify<K: copy, drop, store, V: store>(self: &mut big_ordered_map::BigOrderedMap<K, V>, key: &K, f: |&mut V|)
+
+
+
+
+
+Implementation
+
+
+public inline fun modify<K: drop + copy + store, V: store>(self: &mut BigOrderedMap<K, V>, key: &K, f: |&mut V|) {
+ self.modify_and_return(key, |v| { f(v); true});
+}
+
+
+
+
+
+
+
+
+## Function `modify_and_return`
+
+Modifies the element in the map via calling f, and propagates the return value of the function.
+Aborts if element doesn't exist
+
+This function cannot be inline, due to iter_modify requiring actual function value.
+This also is why we return a value
+
+
+public fun modify_and_return<K: copy, drop, store, V: store, R>(self: &mut big_ordered_map::BigOrderedMap<K, V>, key: &K, f: |&mut V|R): R
+
+
+
+
+
+Implementation
+
+
+public inline fun modify_and_return<K: drop + copy + store, V: store, R>(self: &mut BigOrderedMap<K, V>, key: &K, f: |&mut V|R): R {
+ let iter = self.internal_find(key);
+ assert!(!iter.iter_is_end(self), error::invalid_argument(EKEY_NOT_FOUND));
+ iter.iter_modify(self, |v| f(v))
+}
+
+
+
+
+
+
+
+
+## Function `modify_or_add`
+
+Modifies element by calling modify_f if it exists, or calling add_f to add if it doesn't.
+Returns true if element already existed.
+
+
+public fun modify_or_add<K: copy, drop, store, V: store>(self: &mut big_ordered_map::BigOrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop, add_f: ||V has drop): bool
+
+
+
+
+
+Implementation
+
+
+public inline fun modify_or_add<K: drop + copy + store, V: store>(self: &mut BigOrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop, add_f: ||V has drop): bool {
+ let exists = self.modify_if_present_and_return(key, |v| { modify_f(v); true }).is_some();
+ if (!exists) {
+ self.add(*key, add_f());
+ };
+ exists
+}
+
+
+
+
+
+
+
+
+## Function `modify_if_present`
+
+
+
+public fun modify_if_present<K: copy, drop, store, V: store>(self: &mut big_ordered_map::BigOrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop): bool
+
+
+
+
+
+Implementation
+
+
+public inline fun modify_if_present<K: drop + copy + store, V: store>(self: &mut BigOrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop): bool {
+ self.modify_if_present_and_return(key, |v| { modify_f(v); true }).is_some()
+}
+
+
+
+
+
+
+
+
+## Function `modify_if_present_and_return`
+
+Modifies the element in the map via calling modify_f, and propagates the return value of the function.
+Returns None if not present.
+
+Function value cannot be inlined, due to iter_modify requiring actual function value.
+This also is why we return a value
+
+
+public fun modify_if_present_and_return<K: copy, drop, store, V: store, R>(self: &mut big_ordered_map::BigOrderedMap<K, V>, key: &K, modify_f: |&mut V|R has drop): option::Option<R>
+
+
+
+
+
+Implementation
+
+
+public inline fun modify_if_present_and_return<K: drop + copy + store, V: store, R>(self: &mut BigOrderedMap<K, V>, key: &K, modify_f: |&mut V|R has drop): Option<R> {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(iter.iter_modify(self, |v| modify_f(v)))
+ }
+}
+
+
+
+
@@ -1462,6 +1616,35 @@ Returns a reference to the element with its key, aborts if the key is not found.
+
+
+
+
+## Function `get_and_map`
+
+
+
+public fun get_and_map<K: copy, drop, store, V: copy, store, R>(self: &big_ordered_map::BigOrderedMap<K, V>, key: &K, f: |&V|R has drop): option::Option<R>
+
+
+
+
+
+Implementation
+
+
+public fun get_and_map<K: drop + copy + store, V: copy + store, R>(self: &BigOrderedMap<K, V>, key: &K, f: |&V|R has drop): Option<R> {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(f(iter.iter_borrow(self)))
+ }
+}
+
+
+
+
@@ -2044,6 +2227,44 @@ Note: Requires that the map is not changed after the input iterator is generated
+
+
+
+
+## Function `iter_modify`
+
+
+
+public fun iter_modify<K: drop, store, V: store, R>(self: big_ordered_map::IteratorPtr<K>, map: &mut big_ordered_map::BigOrderedMap<K, V>, f: |&mut V|R): R
+
+
+
+
+
+Implementation
+
+
+public fun iter_modify<K: drop + store, V: store, R>(self: IteratorPtr<K>, map: &mut BigOrderedMap<K, V>, f: |&mut V|R): R {
+ assert!(!self.iter_is_end(map), error::invalid_argument(EITER_OUT_OF_BOUNDS));
+ let IteratorPtr::Some { node_index, child_iter, key } = self;
+ let children = &mut map.borrow_node_mut(node_index).children;
+ let value_mut = &mut child_iter.iter_borrow_mut(children).value;
+ let result = f(value_mut);
+
+ if (map.constant_kv_size) {
+ return result;
+ };
+
+ // validate that after modifications size invariants hold
+ let key_size = bcs::serialized_size(&key);
+ let value_size = bcs::serialized_size(value_mut);
+ map.validate_size_and_init_max_degrees(key_size, value_size);
+ result
+}
+
+
+
+
@@ -3724,6 +3945,23 @@ Given a path to node (excluding the node itself), which is currently stored unde
+
+
+### Function `get_and_map`
+
+
+public fun get_and_map<K: copy, drop, store, V: copy, store, R>(self: &big_ordered_map::BigOrderedMap<K, V>, key: &K, f: |&V|R has drop): option::Option<R>
+
+
+
+
+
+pragma opaque;
+pragma verify = false;
+
+
+
+
### Function `borrow_mut`
@@ -3965,6 +4203,23 @@ std::cmp::compare(key, k) == std::cmp::Ordering::Greater);
+pragma opaque;
+pragma verify = false;
+
+
+
+
+
+
+### Function `iter_modify`
+
+
+public fun iter_modify<K: drop, store, V: store, R>(self: big_ordered_map::IteratorPtr<K>, map: &mut big_ordered_map::BigOrderedMap<K, V>, f: |&mut V|R): R
+
+
+
+
+
pragma opaque;
pragma verify = false;
diff --git a/aptos-move/framework/aptos-framework/doc/ordered_map.md b/aptos-move/framework/aptos-framework/doc/ordered_map.md
index 6158b381b9dd1..54d6e950974b8 100644
--- a/aptos-move/framework/aptos-framework/doc/ordered_map.md
+++ b/aptos-move/framework/aptos-framework/doc/ordered_map.md
@@ -45,9 +45,13 @@ A set of inline utility methods is provided instead, to provide guaranteed valid
- [Function `upsert`](#0x1_ordered_map_upsert)
- [Function `remove`](#0x1_ordered_map_remove)
- [Function `remove_or_none`](#0x1_ordered_map_remove_or_none)
+- [Function `modify_or_add`](#0x1_ordered_map_modify_or_add)
+- [Function `modify_if_present`](#0x1_ordered_map_modify_if_present)
- [Function `contains`](#0x1_ordered_map_contains)
- [Function `borrow`](#0x1_ordered_map_borrow)
- [Function `borrow_mut`](#0x1_ordered_map_borrow_mut)
+- [Function `get`](#0x1_ordered_map_get)
+- [Function `get_and_map`](#0x1_ordered_map_get_and_map)
- [Function `replace_key_inplace`](#0x1_ordered_map_replace_key_inplace)
- [Function `add_all`](#0x1_ordered_map_add_all)
- [Function `upsert_all`](#0x1_ordered_map_upsert_all)
@@ -99,6 +103,8 @@ A set of inline utility methods is provided instead, to provide guaranteed valid
- [Function `contains`](#@Specification_1_contains)
- [Function `borrow`](#@Specification_1_borrow)
- [Function `borrow_mut`](#@Specification_1_borrow_mut)
+ - [Function `get`](#@Specification_1_get)
+ - [Function `get_and_map`](#@Specification_1_get_and_map)
- [Function `replace_key_inplace`](#@Specification_1_replace_key_inplace)
- [Function `add_all`](#@Specification_1_add_all)
- [Function `upsert_all`](#@Specification_1_upsert_all)
@@ -554,6 +560,68 @@ Returns none if key doesn't exist.
+
+
+
+
+## Function `modify_or_add`
+
+Modifies element by calling modify_f if it exists, or calling add_f to add if it doesn't.
+Returns true if element already existed.
+
+
+public fun modify_or_add<K: copy, drop, store, V: store>(self: &mut ordered_map::OrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop, add_f: ||V has drop): bool
+
+
+
+
+
+Implementation
+
+
+public inline fun modify_or_add<K: drop + copy + store, V: store>(self: &mut OrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop, add_f: ||V has drop): bool {
+ let exists = self.modify_if_present(key, |v| modify_f(v));
+ if (!exists) {
+ self.add(*key, add_f());
+ };
+ exists
+}
+
+
+
+
+
+
+
+
+## Function `modify_if_present`
+
+Modifies element by calling modify_f if it exists.
+Returns true if element already existed.
+
+
+public fun modify_if_present<K: copy, drop, store, V: store>(self: &mut ordered_map::OrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop): bool
+
+
+
+
+
+Implementation
+
+
+public inline fun modify_if_present<K: drop + copy + store, V: store>(self: &mut OrderedMap<K, V>, key: &K, modify_f: |&mut V| has drop): bool {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ false
+ } else {
+ modify_f(iter.iter_borrow_mut(self));
+ true
+ }
+}
+
+
+
+
@@ -627,6 +695,64 @@ Returns whether map contains a given key.
+
+
+
+
+## Function `get`
+
+
+
+public fun get<K: copy, drop, store, V: copy, store>(self: &ordered_map::OrderedMap<K, V>, key: &K): option::Option<V>
+
+
+
+
+
+Implementation
+
+
+public fun get<K: drop + copy + store, V: copy + store>(self: &OrderedMap<K, V>, key: &K): Option<V> {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(*iter.iter_borrow(self))
+ }
+}
+
+
+
+
+
+
+
+
+## Function `get_and_map`
+
+
+
+public fun get_and_map<K: copy, drop, store, V: copy, store, R>(self: &ordered_map::OrderedMap<K, V>, key: &K, f: |&V|R has drop): option::Option<R>
+
+
+
+
+
+Implementation
+
+
+public fun get_and_map<K: drop + copy + store, V: copy + store, R>(self: &OrderedMap<K, V>, key: &K, f: |&V|R has drop): Option<R> {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(f(iter.iter_borrow(self)))
+ }
+}
+
+
+
+
@@ -2189,6 +2315,40 @@ Apply the function to a mutable reference of each key-value pair in the map.
+
+
+### Function `get`
+
+
+public fun get<K: copy, drop, store, V: copy, store>(self: &ordered_map::OrderedMap<K, V>, key: &K): option::Option<V>
+
+
+
+
+
+pragma opaque;
+pragma verify = false;
+
+
+
+
+
+
+### Function `get_and_map`
+
+
+public fun get_and_map<K: copy, drop, store, V: copy, store, R>(self: &ordered_map::OrderedMap<K, V>, key: &K, f: |&V|R has drop): option::Option<R>
+
+
+
+
+
+pragma opaque;
+pragma verify = false;
+
+
+
+
### Function `replace_key_inplace`
diff --git a/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.move b/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.move
index ae79022d5b200..7fbd1feac6f60 100644
--- a/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.move
+++ b/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.move
@@ -399,6 +399,68 @@ module aptos_std::big_ordered_map {
}
}
+ /// Modifies the element in the map via calling f.
+ /// Aborts if element doesn't exist
+ public inline fun modify(self: &mut BigOrderedMap, key: &K, f: |&mut V|) {
+ self.modify_and_return(key, |v| { f(v); true});
+ }
+
+ /// Modifies the element in the map via calling f, and propagates the return value of the function.
+ /// Aborts if element doesn't exist
+ ///
+ /// This function cannot be inline, due to iter_modify requiring actual function value.
+ /// This also is why we return a value
+ public inline fun modify_and_return(self: &mut BigOrderedMap, key: &K, f: |&mut V|R): R {
+ let iter = self.internal_find(key);
+ assert!(!iter.iter_is_end(self), error::invalid_argument(EKEY_NOT_FOUND));
+ iter.iter_modify(self, |v| f(v))
+ }
+
+ /// Modifies element by calling modify_f if it exists, or calling add_f to add if it doesn't.
+ /// Returns true if element already existed.
+ public inline fun modify_or_add(self: &mut BigOrderedMap, key: &K, modify_f: |&mut V| has drop, add_f: ||V has drop): bool {
+ let exists = self.modify_if_present_and_return(key, |v| { modify_f(v); true }).is_some();
+ if (!exists) {
+ self.add(*key, add_f());
+ };
+ exists
+ }
+
+ public inline fun modify_if_present(self: &mut BigOrderedMap, key: &K, modify_f: |&mut V| has drop): bool {
+ self.modify_if_present_and_return(key, |v| { modify_f(v); true }).is_some()
+ }
+
+ /// Modifies the element in the map via calling modify_f, and propagates the return value of the function.
+ /// Returns None if not present.
+ ///
+ /// Function value cannot be inlined, due to iter_modify requiring actual function value.
+ /// This also is why we return a value
+ public inline fun modify_if_present_and_return(self: &mut BigOrderedMap, key: &K, modify_f: |&mut V|R has drop): Option {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(iter.iter_modify(self, |v| modify_f(v)))
+ }
+ }
+
+ // /// If value exists, calls modify_f on it, which returns tuple (to_keep, result).
+ // /// If to_keep is false, value is deleted from the map, and option::some(result) is returned.
+ // /// This function cannot be inline, due to iter_modify requiring actual function value.
+ // /// This also is why we return a value
+ // public fun modify_or_remove_if_present_and_return(self: &mut BigOrderedMap, key: &K, modify_f: |&mut V|(R, bool) has drop): Option {
+ // let iter = self.find(key);
+ // if (iter.iter_is_end(self)) {
+ // option::none()
+ // } else {
+ // let (result, keep) = iter.iter_modify(self, modify_f);
+ // if (!keep) {
+ // iter.iter_remove(self);
+ // };
+ // option::some(result)
+ // }
+ // }
+
/// Add multiple key/value pairs to the map. The keys must not already exist.
/// Aborts with EKEY_ALREADY_EXISTS if key already exist, or duplicate keys are passed in.
public fun add_all(self: &mut BigOrderedMap, keys: vector, values: vector) {
@@ -524,6 +586,15 @@ module aptos_std::big_ordered_map {
}
}
+ public fun get_and_map(self: &BigOrderedMap, key: &K, f: |&V|R has drop): Option {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(f(iter.iter_borrow(self)))
+ }
+ }
+
/// Returns a mutable reference to the element with its key at the given index, aborts if the key is not found.
/// Aborts with EBORROW_MUT_REQUIRES_CONSTANT_VALUE_SIZE if KV size doesn't have constant size,
/// because if it doesn't we cannot assert invariants on the size.
@@ -731,6 +802,24 @@ module aptos_std::big_ordered_map {
&mut child_iter.iter_borrow_mut(children).value
}
+ public fun iter_modify(self: IteratorPtr, map: &mut BigOrderedMap, f: |&mut V|R): R {
+ assert!(!self.iter_is_end(map), error::invalid_argument(EITER_OUT_OF_BOUNDS));
+ let IteratorPtr::Some { node_index, child_iter, key } = self;
+ let children = &mut map.borrow_node_mut(node_index).children;
+ let value_mut = &mut child_iter.iter_borrow_mut(children).value;
+ let result = f(value_mut);
+
+ if (map.constant_kv_size) {
+ return result;
+ };
+
+ // validate that after modifications size invariants hold
+ let key_size = bcs::serialized_size(&key);
+ let value_size = bcs::serialized_size(value_mut);
+ map.validate_size_and_init_max_degrees(key_size, value_size);
+ result
+ }
+
/// Removes the entry from BigOrderedMap and returns the value which `key` maps to.
/// Aborts if there is no entry for `key`.
public fun iter_remove(self: IteratorPtrWithPath, map: &mut BigOrderedMap): V {
@@ -1829,6 +1918,30 @@ module aptos_std::big_ordered_map {
map.destroy(|_v| {});
}
+ #[test]
+ fun test_modify_and_get() {
+ let map = new_with_config(4, 3, false);
+ map.add_all(vector[1, 2, 3], vector[1, 2, 3]);
+ map.modify(&2, |v| *v += 10);
+ assert!(map.get(&2) == option::some(12));
+ assert!(map.get(&4) == option::none());
+
+ assert!(map.get_and_map(&2, |v| *v + 5) == option::some(17));
+ assert!(map.get_and_map(&4, |v| *v + 5) == option::none());
+
+ map.modify_or_add(&3, |v| *v += 10, || 20);
+ assert!(map.get(&3) == option::some(13));
+ map.modify_or_add(&4, |v| *v += 10, || 20);
+ assert!(map.get(&4) == option::some(20));
+
+ assert!(option::some(7) == map.modify_if_present_and_return(&4, |v| { *v += 10; 7}));
+ assert!(map.get(&4) == option::some(30));
+
+ assert!(option::none() == map.modify_if_present_and_return(&5, |v| { *v += 10; 7}));
+
+ map.destroy(|_v| {});
+ }
+
#[test]
fun test_contains() {
let map = new_with_config(4, 3, false);
diff --git a/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.spec.move b/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.spec.move
index 86e4aa195ff1d..5dc342278a367 100644
--- a/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.spec.move
+++ b/aptos-move/framework/aptos-framework/sources/datastructures/big_ordered_map.spec.move
@@ -252,5 +252,13 @@ spec aptos_std::big_ordered_map {
ensures [abstract] result == spec_len(self);
}
+ spec get_and_map {
+ pragma opaque;
+ pragma verify = false;
+ }
+ spec iter_modify {
+ pragma opaque;
+ pragma verify = false;
+ }
}
diff --git a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move
index db35394c18ac1..1b0289910fa23 100644
--- a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move
+++ b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move
@@ -148,6 +148,28 @@ module aptos_std::ordered_map {
}
}
+ /// Modifies element by calling modify_f if it exists, or calling add_f to add if it doesn't.
+ /// Returns true if element already existed.
+ public inline fun modify_or_add(self: &mut OrderedMap, key: &K, modify_f: |&mut V| has drop, add_f: ||V has drop): bool {
+ let exists = self.modify_if_present(key, |v| modify_f(v));
+ if (!exists) {
+ self.add(*key, add_f());
+ };
+ exists
+ }
+
+ /// Modifies element by calling modify_f if it exists.
+ /// Returns true if element already existed.
+ public inline fun modify_if_present(self: &mut OrderedMap, key: &K, modify_f: |&mut V| has drop): bool {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ false
+ } else {
+ modify_f(iter.iter_borrow_mut(self));
+ true
+ }
+ }
+
/// Returns whether map contains a given key.
public fun contains(self: &OrderedMap, key: &K): bool {
!self.internal_find(key).iter_is_end(self)
@@ -161,6 +183,24 @@ module aptos_std::ordered_map {
self.internal_find(key).iter_borrow_mut(self)
}
+ public fun get(self: &OrderedMap, key: &K): Option {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(*iter.iter_borrow(self))
+ }
+ }
+
+ public fun get_and_map(self: &OrderedMap, key: &K, f: |&V|R has drop): Option {
+ let iter = self.internal_find(key);
+ if (iter.iter_is_end(self)) {
+ option::none()
+ } else {
+ option::some(f(iter.iter_borrow(self)))
+ }
+ }
+
/// Changes the key, while keeping the same value attached to it
/// Aborts with EKEY_NOT_FOUND if `old_key` doesn't exist.
/// Aborts with ENEW_KEY_NOT_IN_ORDER if `new_key` doesn't keep the order `old_key` was in.
@@ -851,6 +891,23 @@ module aptos_std::ordered_map {
assert!(map.values() == vector[1, 2], 0);
}
+ #[test]
+ fun test_modify_and_get() {
+ let map = new();
+ map.add_all(vector[1, 2, 3], vector[1, 2, 3]);
+ assert!(true == map.modify_if_present(&2, |v| *v += 10));
+ assert!(map.get(&2) == option::some(12));
+ assert!(map.get(&4) == option::none());
+
+ assert!(map.get_and_map(&2, |v| *v + 5) == option::some(17));
+ assert!(map.get_and_map(&4, |v| *v + 5) == option::none());
+
+ map.modify_or_add(&3, |v| {*v += 10}, || 20);
+ assert!(map.get(&3) == option::some(13));
+ map.modify_or_add(&4, |v| {*v += 10}, || 20);
+ assert!(map.get(&4) == option::some(20));
+ }
+
#[test]
fun test_for_each_variants() {
let keys = vector[1, 3, 5];
diff --git a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move
index f36e6f6a6250e..b6166d87c63ca 100644
--- a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move
+++ b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move
@@ -294,4 +294,14 @@ spec aptos_std::ordered_map {
pragma opaque;
pragma verify = false;
}
+
+ spec get {
+ pragma opaque;
+ pragma verify = false;
+ }
+
+ spec get_and_map {
+ pragma opaque;
+ pragma verify = false;
+ }
}