Add capacity to Collection

sui_programmability/framework/sources/Collection.move

@@ -1,24 +1,42 @@
 module FastX::Collection {
+    use Std::Errors;
     use Std::Option::{Self, Option};
     use Std::Vector::Self;
     use FastX::Address::{Self, Address};
     use FastX::ID::{Self, ID, IDBytes};
     use FastX::Transfer;
     use FastX::TxContext::{Self, TxContext};
 
-    const OBJECT_NOT_FOUND: u64 = 0;
-    const OBJECT_DOUBLE_ADD: u64 = 1;
+    // Error codes
+    const EOBJECT_NOT_FOUND: u64 = 0;
+    const EOBJECT_DOUBLE_ADD: u64 = 1;
+    const EINVALID_MAX_CAPACITY: u64 = 2;
+    const EMAX_CAPACITY_EXCEEDED: u64 = 3;
+
+    // TODO: this is a placeholder number
+    const DEFAULT_MAX_CAPACITY: u64 = 65536;
 
     struct Collection has key {
         id: ID,
         objects: vector<IDBytes>,
+        max_capacity: u64,
     }
 
     /// Create a new Collection and return it.
     public fun new(ctx: &mut TxContext): Collection {
+        new_with_max_capacity(ctx, DEFAULT_MAX_CAPACITY)
+    }
+
+    /// Create a new Collection with custom size limit and return it.
+    public fun new_with_max_capacity(ctx: &mut TxContext, max_capacity: u64): Collection {
+        assert!(
+            max_capacity <= DEFAULT_MAX_CAPACITY && max_capacity > 0 ,
+            Errors::limit_exceeded(EINVALID_MAX_CAPACITY)
+        );
         Collection {
             id: TxContext::new_id(ctx),
             objects: Vector::empty<IDBytes>(),
+            max_capacity,
         }
     }
 
@@ -35,9 +53,13 @@ module FastX::Collection {
     /// Add a new object to the collection.
     /// Abort if the object is already in the collection.
     public fun add<T: key>(c: &mut Collection, object: T) {
+        assert!(
+            size(c) + 1 <= c.max_capacity,
+            Errors::limit_exceeded(EMAX_CAPACITY_EXCEEDED)
+        );
         let id_bytes = ID::get_id_bytes(&object);
         if (contains(c, id_bytes)) {
-            abort OBJECT_DOUBLE_ADD
+            abort EOBJECT_DOUBLE_ADD
         };
         Vector::push_back(&mut c.objects, *id_bytes);
         Transfer::transfer_to_object(object, c);
@@ -54,7 +76,7 @@ module FastX::Collection {
     public fun remove<T: key>(c: &mut Collection, object: T): T {
         let idx = find(c, ID::get_id_bytes(&object));
         if (Option::is_none(&idx)) {
-            abort OBJECT_NOT_FOUND
+            abort EOBJECT_DOUBLE_ADD
         };
         Vector::remove(&mut c.objects, *Option::borrow(&idx));
         object

sui_programmability/framework/tests/CollectionTests.move

@@ -4,8 +4,8 @@ module FastX::CollectionTests {
     use FastX::ID::{Self, ID};
     use FastX::TxContext;
 
-    const COLLECTION_SIZE_MISMATCH: u64 = 0;
-    const OBJECT_NOT_FOUND: u64 = 1;
+    const ECOLLECTION_SIZE_MISMATCH: u64 = 0;
+    const EOBJECT_NOT_FOUND: u64 = 1;
 
     struct Object has key {
         id: ID,
@@ -15,7 +15,7 @@ module FastX::CollectionTests {
     fun test_collection_add() {
         let ctx = TxContext::dummy();
         let collection = Collection::new(&mut ctx);
-        assert!(Collection::size(&collection) == 0, COLLECTION_SIZE_MISMATCH);
+        assert!(Collection::size(&collection) == 0, ECOLLECTION_SIZE_MISMATCH);
 
         let obj1 = Object { id: TxContext::new_id(&mut ctx) };
         let id_bytes1 = *ID::get_id_bytes(&obj1);
@@ -24,11 +24,42 @@ module FastX::CollectionTests {
 
         Collection::add(&mut collection, obj1);
         Collection::add(&mut collection, obj2);
-        assert!(Collection::size(&collection) == 2, COLLECTION_SIZE_MISMATCH);
+        assert!(Collection::size(&collection) == 2, ECOLLECTION_SIZE_MISMATCH);
 
-        assert!(Collection::contains(&collection, &id_bytes1), OBJECT_NOT_FOUND);
-        assert!(Collection::contains(&collection, &id_bytes2), OBJECT_NOT_FOUND);
+        assert!(Collection::contains(&collection, &id_bytes1), EOBJECT_NOT_FOUND);
+        assert!(Collection::contains(&collection, &id_bytes2), EOBJECT_NOT_FOUND);
 
         Collection::transfer(collection, TxContext::get_signer_address(&ctx));
     }
+
+    #[test]
+    #[expected_failure(abort_code = 520)] 
+    fun test_init_with_invalid_max_capacity() {
+        let ctx = TxContext::dummy();
+        // FastX::Collection::DEFAULT_MAX_CAPACITY is not readable outside the module
+        let max_capacity = 65536;
+        let collection = Collection::new_with_max_capacity(&mut ctx, max_capacity + 1);
+        Collection::transfer(collection, TxContext::get_signer_address(&ctx));
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 520)] 
+    fun test_init_with_zero() {
+        let ctx = TxContext::dummy();
+        let collection = Collection::new_with_max_capacity(&mut ctx, 0);
+        Collection::transfer(collection, TxContext::get_signer_address(&ctx));
+    }
+
+    #[test]
+    #[expected_failure(abort_code = 776)] 
+    fun test_exceed_max_capacity() {
+        let ctx = TxContext::dummy();
+        let collection = Collection::new_with_max_capacity(&mut ctx, 1);
+
+        let obj1 = Object { id: TxContext::new_id(&mut ctx) };
+        Collection::add(&mut collection, obj1);
+        let obj2 = Object { id: TxContext::new_id(&mut ctx) };
+        Collection::add(&mut collection, obj2);
+        Collection::transfer(collection, TxContext::get_signer_address(&ctx));
+    }
 }