From e2337916afb482ba1380bcbb8161c332c1fe3a11 Mon Sep 17 00:00:00 2001
From: Arnaud Mimart <33665250+amimart@users.noreply.github.com>
Date: Sun, 26 Feb 2023 19:25:54 +0100
Subject: [PATCH] feat(storage): persist object count in bucket state

---
 contracts/cw-storage/src/contract.rs | 7 ++-----
 contracts/cw-storage/src/state.rs    | 3 +++
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/contracts/cw-storage/src/contract.rs b/contracts/cw-storage/src/contract.rs
index dc450274..681b25cb 100644
--- a/contracts/cw-storage/src/contract.rs
+++ b/contracts/cw-storage/src/contract.rs
@@ -55,12 +55,9 @@ pub mod execute {
         pin: bool,
     ) -> Result<Response, ContractError> {
         let size = data.len() as u128;
-        // TODO: store object count in bucket instead of computing it?
-        let object_count = objects()
-            .keys_raw(deps.storage, None, None, Order::Ascending)
-            .count();
         BUCKET.update(deps.storage, |mut bucket| -> Result<_, ContractError> {
             bucket.size += size;
+            bucket.object_count += 1;
             match bucket.limits {
                 Limits {
                     max_object_size: Some(max),
@@ -69,7 +66,7 @@ pub mod execute {
                 Limits {
                     max_objects: Some(max),
                     ..
-                } if object_count as u128 >= max.u128() => {
+                } if bucket.object_count > max.u128() => {
                     Err(BucketError::MaxObjectsLimitExceeded.into())
                 }
                 Limits {
diff --git a/contracts/cw-storage/src/state.rs b/contracts/cw-storage/src/state.rs
index 9b4a77fc..727e6621 100644
--- a/contracts/cw-storage/src/state.rs
+++ b/contracts/cw-storage/src/state.rs
@@ -16,6 +16,8 @@ pub struct Bucket {
     pub limits: Limits,
     /// The total size of the objects contained in the bucket.
     pub size: u128,
+    /// The number of objects in the bucket.
+    pub object_count: u128,
 }
 
 impl Bucket {
@@ -29,6 +31,7 @@ impl Bucket {
             name: n,
             limits,
             size: 0u128,
+            object_count: 0u128,
         })
     }
 }