From 0a2968029d15c34db50cad58c955fe4bf9fa81ab Mon Sep 17 00:00:00 2001 From: Sam Blackshear Date: Sat, 25 Jun 2022 15:57:04 -0700 Subject: [PATCH] [sui framework] add VecMap collection type It's comment for Move programmers to want a collection with a map-like API, but the standard library does not have one. This attempts to close that gap by providing a `VecMap`, which provides this API by using naive equality checking on keys. As the comments on the collection suggest, this is not intended to be used for large collections--instead, the idea is to use this to replace some of the handrolling done in code like `validator_set::find_validator`. I expect there will be lots of third-party code that looks like this and want to provide a standard collection that pre-empts it. --- .../tests/sui/coin_transfer.exp | 2 +- crates/sui-framework/sources/vec_map.move | 152 ++++++++++++++++++ crates/sui-framework/tests/vec_map_tests.move | 99 ++++++++++++ 3 files changed, 252 insertions(+), 1 deletion(-) create mode 100644 crates/sui-framework/sources/vec_map.move create mode 100644 crates/sui-framework/tests/vec_map_tests.move diff --git a/crates/sui-adapter-transactional-tests/tests/sui/coin_transfer.exp b/crates/sui-adapter-transactional-tests/tests/sui/coin_transfer.exp index d9ee7d59e832a..31f572dc96da8 100644 --- a/crates/sui-adapter-transactional-tests/tests/sui/coin_transfer.exp +++ b/crates/sui-adapter-transactional-tests/tests/sui/coin_transfer.exp @@ -28,4 +28,4 @@ Contents: sui::coin::Coin {id: sui::id::VersionedID {id: sui::id: task 7 'view-object'. lines 20-20: Owner: Account Address ( B ) -Contents: sui::coin::Coin {id: sui::id::VersionedID {id: sui::id::UniqueID {id: sui::id::ID {bytes: fake(107)}}, version: 1u64}, balance: sui::balance::Balance {value: 99592u64}} +Contents: sui::coin::Coin {id: sui::id::VersionedID {id: sui::id::UniqueID {id: sui::id::ID {bytes: fake(107)}}, version: 1u64}, balance: sui::balance::Balance {value: 99570u64}} diff --git a/crates/sui-framework/sources/vec_map.move b/crates/sui-framework/sources/vec_map.move new file mode 100644 index 0000000000000..12e2204ddafa9 --- /dev/null +++ b/crates/sui-framework/sources/vec_map.move @@ -0,0 +1,152 @@ +// Copyright (c) 2022, Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +module sui::vec_map { + use std::option::{Self, Option}; + use std::vector; + + /// This key already exists in the map + const EKeyAlreadyExists: u64 = 0; + + /// This key does not exist in the map + const EKeyDoesNotExist: u64 = 1; + + /// Trying to destroy a map that is not empty + const EMapNotEmpty: u64 = 2; + + /// Trying to access an element of the map at an invalid index + const EIndexOutOfBounds: u64 = 3; + + /// A map data structure backed by a vector. The map is guaranteed not to contain duplicate keys, but entries + /// are *not* sorted by key--entries are included in insertion order. + /// All operations are O(N) in the size of the map--the intention of this data structure is only to provide + /// the convenience of programming against a map API. + /// Large maps should use handwritten parent/child relationships instead. + /// Maps that need sorted iteration rather than insertion order iteration should also be handwritten. + struct VecMap has copy, drop, store { + contents: vector>, + } + + /// An entry in the map + struct Entry has copy, drop, store { + key: K, + value: V, + } + + /// Create an empty `VecMap` + public fun empty(): VecMap { + VecMap { contents: vector::empty() } + } + + /// Insert the entry `key` |-> `value` into self. + /// Aborts if `key` is already bound in `self`. + public fun insert(self: &mut VecMap, key: K, value: V) { + assert!(!contains(self, &key), EKeyAlreadyExists); + vector::push_back(&mut self.contents, Entry { key, value }) + } + + /// Remove the entry `key` |-> `value` from self. Aborts if `key` is not bound in `self`. + public fun remove(self: &mut VecMap, key: &K): (K, V) { + let idx = get_idx(self, key); + let Entry { key, value } = vector::remove(&mut self.contents, idx); + (key, value) + } + + /// Get a mutable reference to the value bound to `key` in `self`. + /// Aborts if `key` is not bound in `self`. + public fun get_mut(self: &mut VecMap, key: &K): &mut V { + let idx = get_idx(self, key); + let entry = vector::borrow_mut(&mut self.contents, idx); + &mut entry.value + } + + /// Get a reference to the value bound to `key` in `self`. + /// Aborts if `key` is not bound in `self`. + public fun get(self: &VecMap, key: &K): &V { + let idx = get_idx(self, key); + let entry = vector::borrow(&self.contents, idx); + &entry.value + } + + /// Return true if `self` contains an entry for `key`, false otherwise + public fun contains(self: &VecMap, key: &K): bool { + option::is_some(&get_idx_opt(self, key)) + } + + /// Return the number of entries in `self` + public fun size(self: &VecMap): u64 { + vector::length(&self.contents) + } + + /// Return true if `self` has 0 elements, false otherwise + public fun is_empty(self: &VecMap): bool { + size(self) == 0 + } + + /// Destroy an empty map. Aborts if `self` is not empty + public fun destroy_empty(self: VecMap) { + let VecMap { contents } = self; + assert!(vector::is_empty(&contents), EMapNotEmpty); + vector::destroy_empty(contents) + } + + /// Unpack `self` into vectors of its keys and values. + /// The output keys and values are stored in insertion order, *not* sorted by key. + public fun into_keys_values(self: VecMap): (vector, vector) { + let VecMap { contents } = self; + // reverse the vector so the output keys and values will appear in insertion order + vector::reverse(&mut contents); + let i = 0; + let n = vector::length(&contents); + let keys = vector::empty(); + let values = vector::empty(); + while (i < n) { + let Entry { key, value } = vector::pop_back(&mut contents); + vector::push_back(&mut keys, key); + vector::push_back(&mut values, value); + i = i + 1; + }; + vector::destroy_empty(contents); + (keys, values) + } + + /// Find the index of `key` in `self. Return `None` if `key` is not in `self`. + /// Note that map entries are stored in insertion order, *not* sorted by key. + public fun get_idx_opt(self: &VecMap, key: &K): Option { + let i = 0; + let n = size(self); + while (i < n) { + if (&vector::borrow(&self.contents, i).key == key) { + return option::some(i) + }; + i = i + 1; + }; + option::none() + } + + /// Find the index of `key` in `self. Aborts if `key` is not in `self`. + /// Note that map entries are stored in insertion order, *not* sorted by key. + public fun get_idx(self: &VecMap, key: &K): u64 { + let idx_opt = get_idx_opt(self, key); + assert!(option::is_some(&idx_opt), EKeyDoesNotExist); + option::destroy_some(idx_opt) + } + + /// Return a reference to the `idx`th entry of `self`. This gives direct access into the backing array of the map--use with caution. + /// Note that map entries are stored in insertion order, *not* sorted by key. + /// Aborts if `idx` is greater than or equal to `size(self)` + public fun get_entry_by_idx(self: &VecMap, idx: u64): (&K, &V) { + assert!(idx < size(self), EIndexOutOfBounds); + let entry = vector::borrow(&self.contents, idx); + (&entry.key, &entry.value) + } + + /// Return a mutable reference to the `idx`th entry of `self`. This gives direct access into the backing array of the map--use with caution. + /// Note that map entries are stored in insertion order, *not* sorted by key. + /// Aborts if `idx` is greater than or equal to `size(self)` + public fun get_entry_by_idx_mut(self: &mut VecMap, idx: u64): (&K, &mut V) { + assert!(idx < size(self), EIndexOutOfBounds); + let entry = vector::borrow_mut(&mut self.contents, idx); + (&entry.key, &mut entry.value) + } +} diff --git a/crates/sui-framework/tests/vec_map_tests.move b/crates/sui-framework/tests/vec_map_tests.move new file mode 100644 index 0000000000000..efa3ca11c0b30 --- /dev/null +++ b/crates/sui-framework/tests/vec_map_tests.move @@ -0,0 +1,99 @@ +// Copyright (c) 2022, Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +#[test_only] +module sui::vec_map_tests { + use std::vector; + use sui::vec_map::{Self, VecMap}; + + #[test] + #[expected_failure(abort_code = 0)] + fun duplicate_key_abort() { + let m = vec_map::empty(); + vec_map::insert(&mut m, 1, true); + vec_map::insert(&mut m, 1, false); + } + + #[test] + #[expected_failure(abort_code = 1)] + fun nonexistent_key_get() { + let m = vec_map::empty(); + vec_map::insert(&mut m, 1, true); + let k = 2; + let _v = vec_map::get(&m, &k); + } + + #[test] + #[expected_failure(abort_code = 1)] + fun nonexistent_key_get_idx_or_abort() { + let m = vec_map::empty(); + vec_map::insert(&mut m, 1, true); + let k = 2; + let _idx = vec_map::get_idx(&m, &k); + } + + #[test] + #[expected_failure(abort_code = 1)] + fun nonexistent_key_get_entry_by_idx() { + let m = vec_map::empty(); + vec_map::insert(&mut m, 1, true); + let k = 2; + let _idx = vec_map::get_idx(&m, &k); + } + + #[test] + #[expected_failure(abort_code = 2)] + fun destroy_non_empty() { + let m = vec_map::empty(); + vec_map::insert(&mut m, 1, true); + vec_map::destroy_empty(m) + } + + #[test] + fun destroy_empty() { + let m: VecMap = vec_map::empty(); + assert!(vec_map::is_empty(&m), 0); + vec_map::destroy_empty(m) + } + + #[test] + fun smoke() { + let m = vec_map::empty(); + let i = 0; + while (i < 10) { + let k = i + 2; + let v = i + 5; + vec_map::insert(&mut m, k, v); + i = i + 1; + }; + assert!(!vec_map::is_empty(&m), 0); + assert!(vec_map::size(&m) == 10, 1); + let i = 0; + // make sure the elements are as expected in all of the getter APIs we expose + while (i < 10) { + let k = i + 2; + assert!(vec_map::contains(&m, &k), 2); + let v = *vec_map::get(&m, &k); + assert!(v == i + 5, 3); + assert!(vec_map::get_idx(&m, &k) == i, 4); + let (other_k, other_v) = vec_map::get_entry_by_idx(&m, i); + assert!(*other_k == k, 5); + assert!(*other_v == v, 6); + i = i + 1; + }; + // remove all the elements + let (keys, values) = vec_map::into_keys_values(copy m); + let i = 0; + while (i < 10) { + let k = i + 2; + let (other_k, v) = vec_map::remove(&mut m, &k); + assert!(k == other_k, 7); + assert!(v == i + 5, 8); + assert!(*vector::borrow(&keys, i) == k, 9); + assert!(*vector::borrow(&values, i) == v, 10); + + i = i + 1; + } + } + +}