diff --git a/aptos-move/framework/supra-stdlib/sources/enumerable_map.move b/aptos-move/framework/supra-stdlib/sources/enumerable_map.move index ed47b4aea0f3c..0a86f139a61f9 100644 --- a/aptos-move/framework/supra-stdlib/sources/enumerable_map.move +++ b/aptos-move/framework/supra-stdlib/sources/enumerable_map.move @@ -15,7 +15,7 @@ module supra_std::enumerable_map { const EVECTOR_EMPTY: u64 = 3; /// Enumerable Map to store the key value pairs - struct EnumerableMap has store { + struct EnumerableMap has store { /// List of all keys list: vector, /// Key mapped to a tuple containing the (position of key in list and value corresponding to the key) @@ -23,31 +23,31 @@ module supra_std::enumerable_map { } /// Tuple to store the position of key in list and value corresponding to the key - struct Tuple has store, copy, drop { + struct Tuple has store, copy, drop { position: u64, value: V, } /// Return type - struct KeyValue has store, copy, drop { + struct KeyValue has store, copy, drop { key: K, value: V, } /// To create an empty enum map - public fun new_map(): EnumerableMap { + public fun new_map(): EnumerableMap { return EnumerableMap { list: vector::empty(), map: table::new>() } } /// Add Single Key in the Enumerable Map - public fun add_value(map: &mut EnumerableMap, key: K, value: V) { + public fun add_value(map: &mut EnumerableMap, key: K, value: V) { assert!(!contains(map, key), error::already_exists(EKEY_ALREADY_ADDED)); table::add(&mut map.map, key, Tuple { position: vector::length(&map.list), value }); vector::push_back(&mut map.list, key); } /// Add Multiple Keys in the Enumerable Map - public fun add_value_bulk( + public fun add_value_bulk( map: &mut EnumerableMap, keys: vector, values: vector @@ -82,7 +82,7 @@ module supra_std::enumerable_map { } /// Remove single Key from the Enumerable Map - public fun remove_value(map: &mut EnumerableMap, key: K) { + public fun remove_value(map: &mut EnumerableMap, key: K) { assert!(contains(map, key), error::not_found(EKEY_ABSENT)); let map_last_index = vector::length(&map.list) - 1; @@ -96,7 +96,7 @@ module supra_std::enumerable_map { } /// Remove Multiple Keys from the Enumerable Map - public fun remove_value_bulk( + public fun remove_value_bulk( map: &mut EnumerableMap, keys: vector ): vector { @@ -123,44 +123,66 @@ module supra_std::enumerable_map { } /// Will clear the entire data from the Enumerable Map - public fun clear(map: &mut EnumerableMap) { + public fun clear(map: &mut EnumerableMap) { let list = get_map_list(map); remove_value_bulk(map, list); } /// Returns the value of a key that is present in Enumerable Map - public fun get_value(map: & EnumerableMap, key: K): V { + public fun get_value(map: & EnumerableMap, key: K): V { table::borrow(&map.map, key).value } /// Returns reference to the value of a key that is present in Enumerable Map - public fun get_value_ref(map: & EnumerableMap, key: K): &V { + public fun get_value_ref(map: & EnumerableMap, key: K): &V { &table::borrow(&map.map, key).value } /// Returns the value of a key that is present in Enumerable Map - public fun get_value_mut(map: &mut EnumerableMap, key: K): &mut V { + public fun get_value_mut(map: &mut EnumerableMap, key: K): &mut V { &mut table::borrow_mut(&mut map.map, key).value } /// Returns the list of keys that the Enumerable Map contains - public fun get_map_list(map: &EnumerableMap): vector { + public fun get_map_list(map: &EnumerableMap): vector { return map.list } /// Check whether Key is present into the Enumerable map or not - public fun contains(map: &EnumerableMap, key: K): bool { + public fun contains(map: &EnumerableMap, key: K): bool { table::contains(&map.map, key) } /// Return current length of the EnumerableSetRing - public fun length(set: &EnumerableMap): u64 { + public fun length(set: &EnumerableMap): u64 { return vector::length(&set.list) } + /// Apply the function to each element in the EnumerableMap. + public inline fun for_each_value(set: &EnumerableMap, f: |V|) { + let i = 0; + let len = length(set); + while (i < len) { + let key = *vector::borrow(&set.list, i); + f(table::borrow(&set.map, key).value); + i = i + 1 + } + } + + /// Apply the function to a reference of each element in the EnumerableMap. + public inline fun for_each_value_ref(set: &EnumerableMap, f: |&V|) { + let i = 0; + let len = length(set); + while (i < len) { + let key = *vector::borrow(&set.list, i); + f(&table::borrow(&set.map, key).value); + i = i + 1 + } + } + #[test_only] - struct EnumerableMapTest has key { + struct EnumerableMapTest has key { e: EnumerableMap } @@ -280,4 +302,28 @@ module supra_std::enumerable_map { move_to(owner, EnumerableMapTest { e: enum_map }) } + + #[test(owner= @0x1111)] + public fun test_for_each_value_and_ref(owner: &signer) { + let enum_map = new_map(); + add_value(&mut enum_map, 1, 1); + add_value(&mut enum_map, 2, 2); + add_value(&mut enum_map, 3, 3); + add_value(&mut enum_map, 4, 4); + add_value(&mut enum_map, 5, 5); + add_value(&mut enum_map, 6, 6); + + let i = 1; + for_each_value(&enum_map, |v| { + assert!(i == v, 100); + i = i + 1; + }); + + let j = 1; + for_each_value_ref(&enum_map, |v| { + assert!(&j == v, 200); + j = j + 1; + }); + move_to(owner, EnumerableMapTest { e: enum_map }) + } }