From b5608fbdf431c79512c646dee1fcfdad72c19df0 Mon Sep 17 00:00:00 2001
From: Serkan Reis <serkanreis@gmail.com>
Date: Wed, 6 Nov 2024 13:11:21 +0300
Subject: [PATCH] Add member_count to StageResponse for non-merkle-tree WLs

---
 .../tiered-whitelist-flex/src/contract.rs     | 25 ++++++++++++++++---
 .../tiered-whitelist-flex/src/msg.rs          |  4 ++-
 .../tiered-whitelist/src/contract.rs          | 25 ++++++++++++++++---
 .../whitelists/tiered-whitelist/src/msg.rs    |  4 ++-
 4 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/contracts/whitelists/tiered-whitelist-flex/src/contract.rs b/contracts/whitelists/tiered-whitelist-flex/src/contract.rs
index 31da6a9e0..1b46c73d2 100644
--- a/contracts/whitelists/tiered-whitelist-flex/src/contract.rs
+++ b/contracts/whitelists/tiered-whitelist-flex/src/contract.rs
@@ -570,7 +570,9 @@ pub fn query_stage(deps: Deps, stage_id: u32) -> StdResult<StageResponse> {
         StdError::generic_err("Stage not found")
     );
     Ok(StageResponse {
+        stage_id,
         stage: config.stages[stage_id as usize].clone(),
+        member_count: count_keys_for_stage(deps.storage, stage_id)?,
     })
 }
 
@@ -580,7 +582,24 @@ pub fn query_stage_list(deps: Deps) -> StdResult<StagesResponse> {
         !config.stages.is_empty(),
         StdError::generic_err("No stages found")
     );
-    Ok(StagesResponse {
-        stages: config.stages.clone(),
-    })
+    let stages = config
+        .stages
+        .iter()
+        .enumerate()
+        .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),
+        })
+        .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-flex/src/msg.rs b/contracts/whitelists/tiered-whitelist-flex/src/msg.rs
index c54736065..b760ae4a8 100644
--- a/contracts/whitelists/tiered-whitelist-flex/src/msg.rs
+++ b/contracts/whitelists/tiered-whitelist-flex/src/msg.rs
@@ -169,12 +169,14 @@ pub struct CanExecuteResponse {
 
 #[cw_serde]
 pub struct StageResponse {
+    pub stage_id: u32,
     pub stage: Stage,
+    pub member_count: u32,
 }
 
 #[cw_serde]
 pub struct StagesResponse {
-    pub stages: Vec<Stage>,
+    pub stages: Vec<StageResponse>,
 }
 
 #[cw_serde]
diff --git a/contracts/whitelists/tiered-whitelist/src/contract.rs b/contracts/whitelists/tiered-whitelist/src/contract.rs
index 282ea36ce..3b9bf1e6c 100644
--- a/contracts/whitelists/tiered-whitelist/src/contract.rs
+++ b/contracts/whitelists/tiered-whitelist/src/contract.rs
@@ -550,7 +550,9 @@ pub fn query_stage(deps: Deps, stage_id: u32) -> StdResult<StageResponse> {
         StdError::generic_err("Stage not found")
     );
     Ok(StageResponse {
+        stage_id,
         stage: config.stages[stage_id as usize].clone(),
+        member_count: count_keys_for_stage(deps.storage, stage_id)?,
     })
 }
 
@@ -560,7 +562,24 @@ pub fn query_stage_list(deps: Deps) -> StdResult<StagesResponse> {
         !config.stages.is_empty(),
         StdError::generic_err("No stages found")
     );
-    Ok(StagesResponse {
-        stages: config.stages.clone(),
-    })
+    let stages = config
+        .stages
+        .iter()
+        .enumerate()
+        .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),
+        })
+        .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/msg.rs b/contracts/whitelists/tiered-whitelist/src/msg.rs
index cbe074404..c3e4f42b0 100644
--- a/contracts/whitelists/tiered-whitelist/src/msg.rs
+++ b/contracts/whitelists/tiered-whitelist/src/msg.rs
@@ -155,12 +155,14 @@ pub struct CanExecuteResponse {
 
 #[cw_serde]
 pub struct StageResponse {
+    pub stage_id: u32,
     pub stage: Stage,
+    pub member_count: u32,
 }
 
 #[cw_serde]
 pub struct StagesResponse {
-    pub stages: Vec<Stage>,
+    pub stages: Vec<StageResponse>,
 }
 
 #[cw_serde]