forked from MystenLabs/sui
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[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.
- Loading branch information
1 parent
8a3d078
commit 0a29680
Showing
3 changed files
with
252 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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<K: copy, V> has copy, drop, store { | ||
contents: vector<Entry<K, V>>, | ||
} | ||
|
||
/// An entry in the map | ||
struct Entry<K: copy, V> has copy, drop, store { | ||
key: K, | ||
value: V, | ||
} | ||
|
||
/// Create an empty `VecMap` | ||
public fun empty<K: copy, V>(): VecMap<K,V> { | ||
VecMap { contents: vector::empty() } | ||
} | ||
|
||
/// Insert the entry `key` |-> `value` into self. | ||
/// Aborts if `key` is already bound in `self`. | ||
public fun insert<K: copy, V>(self: &mut VecMap<K,V>, 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<K: copy, V>(self: &mut VecMap<K,V>, 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<K: copy, V>(self: &mut VecMap<K,V>, 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<K: copy, V>(self: &VecMap<K,V>, 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<K: copy, V>(self: &VecMap<K, V>, key: &K): bool { | ||
option::is_some(&get_idx_opt(self, key)) | ||
} | ||
|
||
/// Return the number of entries in `self` | ||
public fun size<K: copy, V>(self: &VecMap<K,V>): u64 { | ||
vector::length(&self.contents) | ||
} | ||
|
||
/// Return true if `self` has 0 elements, false otherwise | ||
public fun is_empty<K: copy, V>(self: &VecMap<K,V>): bool { | ||
size(self) == 0 | ||
} | ||
|
||
/// Destroy an empty map. Aborts if `self` is not empty | ||
public fun destroy_empty<K: copy, V>(self: VecMap<K, V>) { | ||
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<K: copy, V>(self: VecMap<K, V>): (vector<K>, vector<V>) { | ||
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<K: copy, V>(self: &VecMap<K,V>, key: &K): Option<u64> { | ||
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<K: copy, V>(self: &VecMap<K,V>, 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<K: copy, V>(self: &VecMap<K, V>, 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<K: copy, V>(self: &mut VecMap<K, V>, idx: u64): (&K, &mut V) { | ||
assert!(idx < size(self), EIndexOutOfBounds); | ||
let entry = vector::borrow_mut(&mut self.contents, idx); | ||
(&entry.key, &mut entry.value) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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<u64, u64> = 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; | ||
} | ||
} | ||
|
||
} |