From 803e1b367ee2185b923b398b96e8eef39ead3e63 Mon Sep 17 00:00:00 2001
From: Arnaud Mimart <33665250+amimart@users.noreply.github.com>
Date: Mon, 27 Feb 2023 11:32:28 +0100
Subject: [PATCH] fix(storage): use Uint128 instead of primitive in state

---
 contracts/cw-storage/src/contract.rs | 23 ++++++++++++-----------
 contracts/cw-storage/src/state.rs    | 10 +++++-----
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/contracts/cw-storage/src/contract.rs b/contracts/cw-storage/src/contract.rs
index c9e4b899..ba93ba8b 100644
--- a/contracts/cw-storage/src/contract.rs
+++ b/contracts/cw-storage/src/contract.rs
@@ -47,6 +47,7 @@ pub fn execute(
 pub mod execute {
     use super::*;
     use crate::state::Limits;
+    use cosmwasm_std::Uint128;
 
     pub fn store_object(
         deps: DepsMut,
@@ -54,31 +55,29 @@ pub mod execute {
         data: Binary,
         pin: bool,
     ) -> Result<Response, ContractError> {
-        let size = data.len() as u128;
+        let size = (data.len() as u128).into();
         BUCKET.update(deps.storage, |mut bucket| -> Result<_, ContractError> {
             bucket.size += size;
-            bucket.object_count += 1;
+            bucket.object_count += Uint128::one();
             match bucket.limits {
                 Limits {
                     max_object_size: Some(max),
                     ..
-                } if size > max.u128() => Err(BucketError::MaxObjectSizeLimitExceeded.into()),
+                } if size > max => Err(BucketError::MaxObjectSizeLimitExceeded.into()),
                 Limits {
                     max_objects: Some(max),
                     ..
-                } if bucket.object_count > max.u128() => {
-                    Err(BucketError::MaxObjectsLimitExceeded.into())
-                }
+                } if bucket.object_count > max => Err(BucketError::MaxObjectsLimitExceeded.into()),
                 Limits {
                     max_object_pins: Some(max),
                     ..
-                } if pin && max.u128() < 1u128 => {
+                } if pin && max < Uint128::one() => {
                     Err(BucketError::MaxObjectPinsLimitExceeded.into())
                 }
                 Limits {
                     max_total_size: Some(max),
                     ..
-                } if bucket.size > max.u128() => Err(BucketError::MaxTotalSizeLimitExceeded.into()),
+                } if bucket.size > max => Err(BucketError::MaxTotalSizeLimitExceeded.into()),
                 _ => Ok(bucket),
             }
         })?;
@@ -144,7 +143,7 @@ pub mod query {
             .load(deps.storage, id)
             .map(|object| ObjectResponse {
                 id: object.id.clone(),
-                size: object.size.into(),
+                size: object.size,
                 owner: object.owner.into(),
                 is_pinned: pins()
                     .idx
@@ -306,7 +305,7 @@ mod tests {
             let created = objects().load(&deps.storage, String::from(obj.2)).unwrap();
             assert_eq!(created.id, obj.2);
             assert_eq!(created.owner, info.clone().sender);
-            assert_eq!(created.size, obj.3);
+            assert_eq!(created.size.u128(), obj.3);
 
             assert_eq!(
                 pins().has(&deps.storage, (String::from(obj.2), info.clone().sender)),
@@ -314,7 +313,9 @@ mod tests {
             );
         }
 
-        assert_eq!(BUCKET.load(&deps.storage).unwrap().size, obj1.3 + obj2.3);
+        let bucket = BUCKET.load(&deps.storage).unwrap();
+        assert_eq!(bucket.size.u128(), obj1.3 + obj2.3);
+        assert_eq!(bucket.object_count.u128(), 2);
 
         let msg = ExecuteMsg::StoreObject {
             data: Binary::from_base64(obj1.0.as_str()).unwrap(),
diff --git a/contracts/cw-storage/src/state.rs b/contracts/cw-storage/src/state.rs
index c80e2f4d..996b70c7 100644
--- a/contracts/cw-storage/src/state.rs
+++ b/contracts/cw-storage/src/state.rs
@@ -15,9 +15,9 @@ pub struct Bucket {
     /// The limits of the bucket.
     pub limits: Limits,
     /// The total size of the objects contained in the bucket.
-    pub size: u128,
+    pub size: Uint128,
     /// The number of objects in the bucket.
-    pub object_count: u128,
+    pub object_count: Uint128,
 }
 
 impl Bucket {
@@ -30,8 +30,8 @@ impl Bucket {
         Ok(Self {
             name: n,
             limits,
-            size: 0u128,
-            object_count: 0u128,
+            size: Uint128::zero(),
+            object_count: Uint128::zero(),
         })
     }
 }
@@ -71,7 +71,7 @@ pub struct Object {
     /// The owner of the object.
     pub owner: Addr,
     /// The size of the object.
-    pub size: u128,
+    pub size: Uint128,
 }
 
 pub struct ObjectIndexes<'a> {