[Feature] Simple map

build/StarcoinFramework/docs/README.md

@@ -93,6 +93,7 @@ This is the root document for the Move StarcoinFramework module documentation. T
 -  [`0x1::Signature`](Signature.md#0x1_Signature)
 -  [`0x1::SignedInteger64`](SignedInteger64.md#0x1_SignedInteger64)
 -  [`0x1::Signer`](Signer.md#0x1_Signer)
+-  [`0x1::SimpleMap`](SimpleMap.md#0x1_SimpleMap)
 -  [`0x1::SnapshotUtil`](SnapshotUtil.md#0x1_SnapshotUtil)
 -  [`0x1::StakeToSBTPlugin`](StakeToSBTPlugin.md#0x1_StakeToSBTPlugin)
 -  [`0x1::StarcoinDAO`](StarcoinDAO.md#0x1_StarcoinDAO)

sources/SimpleMap.move

@@ -0,0 +1,229 @@
+/// This module provides a solution for sorted maps, that is it has the properties that
+/// 1) Keys point to Values
+/// 2) Each Key must be unique
+/// 3) A Key can be found within O(Log N) time
+/// 4) The data is stored as sorted by Key
+/// 5) Adds and removals take O(N) time
+module StarcoinFramework::SimpleMap {
+
+    use StarcoinFramework::Errors;
+    use StarcoinFramework::Option;
+    use StarcoinFramework::Vector;
+    use StarcoinFramework::Compare;
+    use StarcoinFramework::BCS;
+
+    spec module {
+        pragma verify = false;
+        pragma aborts_if_is_strict = true;
+    }
+
+    /// Map key already exists
+    const EKEY_ALREADY_EXISTS: u64 = 1;
+    /// Map key is not found
+    const EKEY_NOT_FOUND: u64 = 2;
+
+    struct SimpleMap<Key, Value> has copy, drop, store {
+        data: vector<Element<Key, Value>>,
+    }
+
+    struct Element<Key, Value> has copy, drop, store {
+        key: Key,
+        value: Value,
+    }
+
+    public fun length<Key: store, Value: store>(map: &SimpleMap<Key, Value>): u64 {
+        Vector::length(&map.data)
+    }
+
+    public fun create<Key: store, Value: store>(): SimpleMap<Key, Value> {
+        SimpleMap {
+            data: Vector::empty(),
+        }
+    }
+
+    public fun borrow<Key: store, Value: store>(
+        map: &SimpleMap<Key, Value>,
+        key: &Key,
+    ): &Value {
+        let (maybe_idx, _) = find(map, key);
+        assert!(Option::is_some(&maybe_idx), Errors::invalid_argument(EKEY_NOT_FOUND));
+        let idx = Option::extract(&mut maybe_idx);
+        &Vector::borrow(&map.data, idx).value
+    }
+
+    public fun borrow_mut<Key: store, Value: store>(
+        map: &mut SimpleMap<Key, Value>,
+        key: &Key,
+    ): &mut Value {
+        let (maybe_idx, _) = find(map, key);
+        assert!(Option::is_some(&maybe_idx), Errors::invalid_argument(EKEY_NOT_FOUND));
+        let idx = Option::extract(&mut maybe_idx);
+        &mut Vector::borrow_mut(&mut map.data, idx).value
+    }
+
+    public fun contains_key<Key: store, Value: store>(
+        map: &SimpleMap<Key, Value>,
+        key: &Key,
+    ): bool {
+        let (maybe_idx, _) = find(map, key);
+        Option::is_some(&maybe_idx)
+    }
+
+    public fun destroy_empty<Key: store, Value: store>(map: SimpleMap<Key, Value>) {
+        let SimpleMap { data } = map;
+        Vector::destroy_empty(data);
+    }
+
+    public fun add<Key: store, Value: store>(
+        map: &mut SimpleMap<Key, Value>,
+        key: Key,
+        value: Value,
+    ) {
+        let (maybe_idx, maybe_placement) = find(map, &key);
+        assert!(Option::is_none(&maybe_idx), Errors::invalid_argument(EKEY_ALREADY_EXISTS));
+
+        // Append to the end and then swap elements until the list is ordered again
+        Vector::push_back(&mut map.data, Element { key, value });
+
+        let placement = Option::extract(&mut maybe_placement);
+        let end = Vector::length(&map.data) - 1;
+        while (placement < end) {
+            Vector::swap(&mut map.data, placement, end);
+            placement = placement + 1;
+        };
+    }
+
+    public fun remove<Key: store, Value: store>(
+        map: &mut SimpleMap<Key, Value>,
+        key: &Key,
+    ): (Key, Value) {
+        let (maybe_idx, _) = find(map, key);
+        assert!(Option::is_some(&maybe_idx), Errors::invalid_argument(EKEY_NOT_FOUND));
+
+        let placement = Option::extract(&mut maybe_idx);
+        let end = Vector::length(&map.data) - 1;
+
+        while (placement < end) {
+            Vector::swap(&mut map.data, placement, placement + 1);
+            placement = placement + 1;
+        };
+
+        let Element { key, value } = Vector::pop_back(&mut map.data);
+        (key, value)
+    }
+
+    fun find<Key: store, Value: store>(
+        map: &SimpleMap<Key, Value>,
+        key: &Key,
+    ): (Option::Option<u64>, Option::Option<u64>) {
+        let length = Vector::length(&map.data);
+
+        if (length == 0) {
+            return (Option::none(), Option::some(0))
+        };
+
+        let left = 0;
+        let right = length;
+
+        while (left != right) {
+            let mid = left + (right - left) / 2;
+            let potential_key = &Vector::borrow(&map.data, mid).key;
+            if (Compare::is_less_than(Compare::cmp_bytes(&BCS::to_bytes(potential_key), &BCS::to_bytes(key)))) {
+                left = mid + 1;
+            } else {
+                right = mid;
+            };
+        };
+
+        if (left != length && key == &Vector::borrow(&map.data, left).key) {
+            (Option::some(left), Option::none())
+        } else {
+            (Option::none(), Option::some(left))
+        }
+    }
+
+    #[test]
+    public fun add_remove_many() {
+        let map = create<u64, u64>();
+
+        assert!(length(&map) == 0, 0);
+        assert!(!contains_key(&map, &3), 1);
+        add(&mut map, 3, 1);
+        assert!(length(&map) == 1, 2);
+        assert!(contains_key(&map, &3), 3);
+        assert!(borrow(&map, &3) == &1, 4);
+        *borrow_mut(&mut map, &3) = 2;
+        assert!(borrow(&map, &3) == &2, 5);
+
+        assert!(!contains_key(&map, &2), 6);
+        add(&mut map, 2, 5);
+        assert!(length(&map) == 2, 7);
+        assert!(contains_key(&map, &2), 8);
+        assert!(borrow(&map, &2) == &5, 9);
+        *borrow_mut(&mut map, &2) = 9;
+        assert!(borrow(&map, &2) == &9, 10);
+
+        remove(&mut map, &2);
+        assert!(length(&map) == 1, 11);
+        assert!(!contains_key(&map, &2), 12);
+        assert!(borrow(&map, &3) == &2, 13);
+
+        remove(&mut map, &3);
+        assert!(length(&map) == 0, 14);
+        assert!(!contains_key(&map, &3), 15);
+
+        destroy_empty(map);
+    }
+
+    #[test]
+    public fun test_several() {
+        let map = create<u64, u64>();
+        add(&mut map, 6, 6);
+        add(&mut map, 1, 1);
+        add(&mut map, 5, 5);
+        add(&mut map, 2, 2);
+        add(&mut map, 3, 3);
+        add(&mut map, 0, 0);
+        add(&mut map, 7, 7);
+        add(&mut map, 4, 4);
+
+        let idx = 0;
+        while (idx < Vector::length(&map.data)) {
+            assert!(idx == Vector::borrow(&map.data, idx).key, idx);
+            idx = idx + 1;
+        };
+
+        remove(&mut map, &0);
+        remove(&mut map, &1);
+        remove(&mut map, &2);
+        remove(&mut map, &3);
+        remove(&mut map, &4);
+        remove(&mut map, &5);
+        remove(&mut map, &6);
+        remove(&mut map, &7);
+
+        destroy_empty(map);
+    }
+
+    #[test]
+    #[expected_failure]
+    public fun add_twice() {
+        let map = create<u64, u64>();
+        add(&mut map, 3, 1);
+        add(&mut map, 3, 1);
+
+        remove(&mut map, &3);
+        destroy_empty(map);
+    }
+
+    #[test]
+    #[expected_failure]
+    public fun remove_twice() {
+        let map = create<u64, u64>();
+        add(&mut map, 3, 1);
+        remove(&mut map, &3);
+        remove(&mut map, &3);
+
+        destroy_empty(map);
+    }
+}