From f1ff48cb56bfa8bf3a20f200e647e3f62cdeb222 Mon Sep 17 00:00:00 2001
From: Serkan Reis <serkanreis@gmail.com>
Date: Wed, 6 Nov 2024 18:26:29 +0300
Subject: [PATCH] Update stage member count logic

---
 .../tiered-whitelist/src/contract.rs          | 33 +++++++++++--------
 .../whitelists/tiered-whitelist/src/state.rs  |  2 ++
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/contracts/whitelists/tiered-whitelist/src/contract.rs b/contracts/whitelists/tiered-whitelist/src/contract.rs
index 3b9bf1e6c..310b008c4 100644
--- a/contracts/whitelists/tiered-whitelist/src/contract.rs
+++ b/contracts/whitelists/tiered-whitelist/src/contract.rs
@@ -10,7 +10,7 @@ use crate::msg::{
     QueryMsg, RemoveMembersMsg, StageMemberInfoResponse, StageResponse, StagesResponse,
     UpdateStageConfigMsg,
 };
-use crate::state::{AdminList, Config, Stage, ADMIN_LIST, CONFIG, WHITELIST_STAGES};
+use crate::state::{AdminList, Config, Stage, ADMIN_LIST, CONFIG, MEMBER_COUNT, WHITELIST_STAGES};
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
@@ -102,6 +102,11 @@ pub fn instantiate(
     }
 
     for stage in 0..msg.stages.clone().len() {
+        MEMBER_COUNT.save(
+            deps.storage,
+            stage as u32,
+            &(msg.members[stage].len() as u32),
+        )?;
         for member in msg.members[stage].iter() {
             let addr = deps.api.addr_validate(member)?;
             WHITELIST_STAGES.save(deps.storage, (stage as u32, addr), &true)?;
@@ -207,6 +212,9 @@ pub fn execute_add_members(
         }
         members_added += 1;
         WHITELIST_STAGES.save(deps.storage, (msg.stage_id, addr.clone()), &true)?;
+        MEMBER_COUNT.update(deps.storage, msg.stage_id, |count| {
+            Ok::<u32, StdError>(count.unwrap_or(0) + 1)
+        })?;
         config.num_members += 1;
     }
 
@@ -244,6 +252,9 @@ pub fn execute_remove_members(
             return Err(ContractError::NoMemberFound(addr.to_string()));
         }
         WHITELIST_STAGES.remove(deps.storage, (msg.stage_id, addr.clone()));
+        MEMBER_COUNT.update(deps.storage, msg.stage_id, |count| {
+            Ok::<u32, StdError>(count.unwrap_or(0).saturating_sub(1))
+        })?;
         config.num_members -= 1;
     }
 
@@ -275,7 +286,7 @@ pub fn execute_add_stage(
     // remove duplicate members
     members.sort_unstable();
     members.dedup();
-    for add in members.into_iter() {
+    for add in members.clone().into_iter() {
         if config.num_members >= config.member_limit {
             return Err(ContractError::MembersExceeded {
                 expected: config.member_limit,
@@ -289,6 +300,7 @@ pub fn execute_add_stage(
         WHITELIST_STAGES.save(deps.storage, (stage_id, addr.clone()), &true)?;
         config.num_members += 1;
     }
+    MEMBER_COUNT.save(deps.storage, stage_id, &(members.len() as u32))?;
 
     CONFIG.save(deps.storage, &config)?;
     Ok(Response::new()
@@ -328,6 +340,7 @@ pub fn execute_remove_stage(
             WHITELIST_STAGES.remove(deps.storage, (stage, member));
             config.num_members -= 1;
         }
+        MEMBER_COUNT.remove(deps.storage, stage);
     }
 
     CONFIG.save(deps.storage, &config)?;
@@ -552,7 +565,7 @@ pub fn query_stage(deps: Deps, stage_id: u32) -> StdResult<StageResponse> {
     Ok(StageResponse {
         stage_id,
         stage: config.stages[stage_id as usize].clone(),
-        member_count: count_keys_for_stage(deps.storage, stage_id)?,
+        member_count: MEMBER_COUNT.may_load(deps.storage, stage_id)?.unwrap_or(0),
     })
 }
 
@@ -569,17 +582,11 @@ pub fn query_stage_list(deps: Deps) -> StdResult<StagesResponse> {
         .map(|(i, stage)| StageResponse {
             stage_id: i as u32,
             stage: stage.clone(),
-            member_count: count_keys_for_stage(deps.storage, i as u32).unwrap_or(0),
+            member_count: MEMBER_COUNT
+                .may_load(deps.storage, i as u32)
+                .unwrap_or(Some(0u32))
+                .unwrap_or(0u32),
         })
         .collect();
     Ok(StagesResponse { stages })
 }
-
-fn count_keys_for_stage(storage: &dyn cosmwasm_std::Storage, stage: u32) -> StdResult<u32> {
-    let count = WHITELIST_STAGES
-        .prefix(stage)
-        .keys(storage, None, None, Order::Ascending)
-        .count();
-
-    Ok(count as u32)
-}
diff --git a/contracts/whitelists/tiered-whitelist/src/state.rs b/contracts/whitelists/tiered-whitelist/src/state.rs
index 57abd3443..b157e10f6 100644
--- a/contracts/whitelists/tiered-whitelist/src/state.rs
+++ b/contracts/whitelists/tiered-whitelist/src/state.rs
@@ -40,3 +40,5 @@ pub const ADMIN_LIST: Item<AdminList> = Item::new("admin_list");
 pub const CONFIG: Item<Config> = Item::new("config");
 
 pub const WHITELIST_STAGES: Map<(u32, Addr), bool> = Map::new("wl_stages");
+
+pub const MEMBER_COUNT: Map<u32, u32> = Map::new("member_count");