Add VecSet in Move framework (MystenLabs#2871)

crates/sui-framework/sources/bag.move

@@ -12,33 +12,30 @@
 module sui::bag {
     use std::errors;
     use std::option::{Self, Option};
-    use std::vector::Self;
     use sui::id::{Self, ID, VersionedID};
     use sui::transfer::{Self, ChildRef};
     use sui::tx_context::{Self, TxContext};
+    use sui::vec_set::VecSet;
+    use sui::vec_set;
 
     // Error codes
-    /// When removing an object from the collection, EObjectNotFound
-    /// will be triggered if the object is not owned by the collection.
-    const EObjectNotFound: u64 = 0;
-
     /// Adding the same object to the collection twice is not allowed.
-    const EObjectDoubleAdd: u64 = 1;
+    const EObjectDoubleAdd: u64 = 0;
 
     /// The max capacity set for the collection cannot exceed the hard limit
     /// which is DEFAULT_MAX_CAPACITY.
-    const EInvalidMaxCapacity: u64 = 2;
+    const EInvalidMaxCapacity: u64 = 1;
 
     /// Trying to add object to the collection when the collection is
     /// already at its maximum capacity.
-    const EMaxCapacityExceeded: u64 = 3;
+    const EMaxCapacityExceeded: u64 = 2;
 
     // TODO: this is a placeholder number
     const DEFAULT_MAX_CAPACITY: u64 = 65536;
 
     struct Bag has key {
         id: VersionedID,
-        objects: vector<ID>,
+        objects: VecSet<ID>,
         max_capacity: u64,
     }
 
@@ -55,7 +52,7 @@ module sui::bag {
         );
         Bag {
             id: tx_context::new_id(ctx),
-            objects: vector::empty(),
+            objects: vec_set::empty(),
             max_capacity,
         }
     }
@@ -67,7 +64,7 @@ module sui::bag {
 
     /// Returns the size of the Bag.
     public fun size(c: &Bag): u64 {
-        vector::length(&c.objects)
+        vec_set::size(&c.objects)
     }
 
     /// Add an object to the Bag.
@@ -83,7 +80,7 @@ module sui::bag {
         if (contains(c, id)) {
             abort EObjectDoubleAdd
         };
-        vector::push_back(&mut c.objects, *id);
+        vec_set::insert(&mut c.objects, *id);
         transfer::transfer_to_object_unsafe(object, old_child_ref, c);
     }
 
@@ -103,17 +100,13 @@ module sui::bag {
     /// Check whether the Bag contains a specific object,
     /// identified by the object id in bytes.
     public fun contains(c: &Bag, id: &ID): bool {
-        option::is_some(&find(c, id))
+        vec_set::contains(&c.objects, id)
     }
 
     /// Remove and return the object from the Bag.
     /// Abort if the object is not found.
     public fun remove<T: key + store>(c: &mut Bag, object: T): T {
-        let idx = find(c, id::id(&object));
-        if (option::is_none(&idx)) {
-            abort EObjectNotFound
-        };
-        vector::remove(&mut c.objects, *option::borrow(&idx));
+        vec_set::remove(&mut c.objects, id::id(&object));
         object
     }
 
@@ -134,18 +127,4 @@ module sui::bag {
     ): (VersionedID, ChildRef<Bag>) {
         transfer::transfer_to_object_id(obj, owner_id)
     }
-
-    /// Look for the object identified by `id_bytes` in the Bag.
-    /// Returns the index if found, none if not found.
-    fun find(c: &Bag, id: &ID): Option<u64> {
-        let i = 0;
-        let len = size(c);
-        while (i < len) {
-            if (vector::borrow(&c.objects, i) == id) {
-                return option::some(i)
-            };
-            i = i + 1;
-        };
-        return option::none()
-    }
 }

crates/sui-framework/sources/vec_set.move

@@ -0,0 +1,85 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+module sui::vec_set {
+    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;
+
+    /// A set data structure backed by a vector. The set is guaranteed not to contain duplicate keys.
+    /// All operations are O(N) in the size of the set--the intention of this data structure is only to provide
+    /// the convenience of programming against a set API.
+    /// Sets that need sorted iteration rather than insertion order iteration should be handwritten.
+    struct VecSet<K: copy + drop> has copy, drop, store {
+        contents: vector<K>,
+    }
+
+    /// Create an empty `VecSet`
+    public fun empty<K: copy + drop>(): VecSet<K> {
+        VecSet { contents: vector::empty() }
+    }
+
+    /// Insert a `key` into self.
+    /// Aborts if `key` is already present in `self`.
+    public fun insert<K: copy + drop>(self: &mut VecSet<K>, key: K) {
+        assert!(!contains(self, &key), EKeyAlreadyExists);
+        vector::push_back(&mut self.contents, key)
+    }
+
+    /// Remove the entry `key` from self. Aborts if `key` is not present in `self`.
+    public fun remove<K: copy + drop>(self: &mut VecSet<K>, key: &K) {
+        let idx = get_idx(self, key);
+        vector::remove(&mut self.contents, idx);
+    }
+
+    /// Return true if `self` contains an entry for `key`, false otherwise
+    public fun contains<K: copy + drop>(self: &VecSet<K>, key: &K): bool {
+        option::is_some(&get_idx_opt(self, key))
+    }
+
+    /// Return the number of entries in `self`
+    public fun size<K: copy + drop>(self: &VecSet<K>): u64 {
+        vector::length(&self.contents)
+    }
+
+    /// Return true if `self` has 0 elements, false otherwise
+    public fun is_empty<K: copy + drop>(self: &VecSet<K>): bool {
+        size(self) == 0
+    }
+
+    /// Unpack `self` into vectors of keys.
+    /// The output keys are stored in insertion order, *not* sorted.
+    public fun into_keys<K: copy + drop>(self: VecSet<K>): vector<K> {
+        let VecSet { contents } = self;
+        contents
+    }
+
+    // == Helper functions ==
+
+    /// Find the index of `key` in `self. Return `None` if `key` is not in `self`.
+    /// Note that keys are stored in insertion order, *not* sorted.
+    fun get_idx_opt<K: copy + drop>(self: &VecSet<K>, key: &K): Option<u64> {
+        let i = 0;
+        let n = size(self);
+        while (i < n) {
+            if (vector::borrow(&self.contents, i) == 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.
+    fun get_idx<K: copy + drop>(self: &VecSet<K>, key: &K): u64 {
+        let idx_opt = get_idx_opt(self, key);
+        assert!(option::is_some(&idx_opt), EKeyDoesNotExist);
+        option::destroy_some(idx_opt)
+    }
+}

crates/sui-framework/tests/bag_tests.move

@@ -54,7 +54,7 @@ module sui::bag_tests {
     }
 
     #[test]
-    #[expected_failure(abort_code = 520)]
+    #[expected_failure(abort_code = 264)]
     fun test_init_with_invalid_max_capacity() {
         let ctx = tx_context::dummy();
         // Sui::bag::DEFAULT_MAX_CAPACITY is not readable outside the module
@@ -64,15 +64,15 @@ module sui::bag_tests {
     }
 
     #[test]
-    #[expected_failure(abort_code = 520)]
+    #[expected_failure(abort_code = 264)]
     fun test_init_with_zero() {
         let ctx = tx_context::dummy();
         let bag = bag::new_with_max_capacity(&mut ctx, 0);
         bag::transfer(bag, tx_context::sender(&ctx));
     }
 
     #[test]
-    #[expected_failure(abort_code = 776)]
+    #[expected_failure(abort_code = 520)]
     fun test_exceed_max_capacity() {
         let ctx = tx_context::dummy();
         let bag = bag::new_with_max_capacity(&mut ctx, 1);

crates/sui-framework/tests/vec_set_tests.move

@@ -0,0 +1,61 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+#[test_only]
+module sui::vec_set_tests {
+    use std::vector;
+    use sui::vec_set;
+
+    fun key_contains() {
+        let m = vec_set::empty();
+        vec_set::insert(&mut m, 1);
+        assert!(vec_set::contains(&m, &1), 0);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 0)]
+    fun duplicate_key_abort() {
+        let m = vec_set::empty();
+        vec_set::insert(&mut m, 1);
+        vec_set::insert(&mut m, 1);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 1)]
+    fun nonexistent_key_remove() {
+        let m = vec_set::empty();
+        vec_set::insert(&mut m, 1);
+        let k = 2;
+        vec_set::remove(&mut m, &k);
+    }
+
+    #[test]
+    fun smoke() {
+        let m = vec_set::empty();
+        let i = 0;
+        while (i < 10) {
+            let k = i + 2;
+            vec_set::insert(&mut m, k);
+            i = i + 1;
+        };
+        assert!(!vec_set::is_empty(&m), 0);
+        assert!(vec_set::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_set::contains(&m, &k), 2);
+            i = i + 1;
+        };
+        // remove all the elements
+        let keys = vec_set::into_keys(copy m);
+        let i = 0;
+        while (i < 10) {
+            let k = i + 2;
+            vec_set::remove(&mut m, &k);
+            assert!(*vector::borrow(&keys, i) == k, 9);
+            i = i + 1;
+        }
+    }
+
+}