Add member_count to StageResponse for non-merkle-tree WLs

public-awesome · Nov 6, 2024 · b5608fb · b5608fb
1 parent 6521fde
commit b5608fb
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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]