Update whitelist queries

public-awesome · Nov 5, 2024 · 5f6a681 · 5f6a681
1 parent f832f27
commit 5f6a681
Show file tree

Hide file tree

Showing 6 changed files with 72 additions and 13 deletions.
diff --git a/contracts/whitelists/tiered-whitelist-flex/src/contract.rs b/contracts/whitelists/tiered-whitelist-flex/src/contract.rs
@@ -7,7 +7,7 @@ use crate::helpers::{fetch_active_stage, fetch_active_stage_index, validate_stag
 use crate::msg::{
     AddMembersMsg, ConfigResponse, ExecuteMsg, HasEndedResponse, HasMemberResponse,
     HasStartedResponse, InstantiateMsg, IsActiveResponse, Member, MembersResponse, QueryMsg,
-    RemoveMembersMsg, StageResponse, StagesResponse, UpdateStageConfigMsg,
+    RemoveMembersMsg, StageMemberInfoResponse, StageResponse, StagesResponse, UpdateStageConfigMsg,
 };
 use crate::state::{AdminList, Config, Stage, ADMIN_LIST, CONFIG, WHITELIST_STAGES};
 #[cfg(not(feature = "library"))]
@@ -396,6 +396,9 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
             to_json_binary(&fetch_active_stage_index(deps.storage, &env).map_or(0, |i| i + 1))
         }
         QueryMsg::HasMember { member } => to_json_binary(&query_has_member(deps, env, member)?),
+        QueryMsg::StageMemberInfo { stage_id, member } => {
+            to_json_binary(&query_stage_member_info(deps, stage_id, member)?)
+        }
         QueryMsg::Config {} => to_json_binary(&query_config(deps, env)?),
         QueryMsg::Stage { stage_id } => to_json_binary(&query_stage(deps, stage_id)?),
         QueryMsg::Stages {} => to_json_binary(&query_stage_list(deps)?),
@@ -463,6 +466,19 @@ pub fn query_has_member(deps: Deps, env: Env, member: String) -> StdResult<HasMe
     Ok(HasMemberResponse { has_member })
 }
 
+pub fn query_stage_member_info(
+    deps: Deps,
+    stage_id: u32,
+    member: String,
+) -> StdResult<StageMemberInfoResponse> {
+    let addr = deps.api.addr_validate(&member)?;
+    let mint_count = WHITELIST_STAGES.may_load(deps.storage, (stage_id, addr.clone()))?;
+    Ok(StageMemberInfoResponse {
+        is_member: mint_count.is_some(),
+        per_address_limit: mint_count.unwrap_or(0),
+    })
+}
+
 pub fn query_member(deps: Deps, env: Env, member: String) -> StdResult<Member> {
     let addr = deps.api.addr_validate(&member)?;
     let active_stage_id = fetch_active_stage_index(deps.storage, &env);

diff --git a/contracts/whitelists/tiered-whitelist-flex/src/msg.rs b/contracts/whitelists/tiered-whitelist-flex/src/msg.rs
@@ -81,7 +81,10 @@ pub enum QueryMsg {
     HasMember {
         member: String,
     },
-
+    StageMemberInfo {
+        member: String,
+        stage_id: u32,
+    },
     Member {
         member: String,
     },
@@ -170,3 +173,9 @@ pub struct StageResponse {
 pub struct StagesResponse {
     pub stages: Vec<Stage>,
 }
+
+#[cw_serde]
+pub struct StageMemberInfoResponse {
+    pub is_member: bool,
+    pub per_address_limit: u32,
+}
diff --git a/contracts/whitelists/tiered-whitelist-merkletree/src/contract.rs b/contracts/whitelists/tiered-whitelist-merkletree/src/contract.rs
@@ -220,8 +220,8 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::CanExecute { sender, .. } => to_json_binary(&query_can_execute(deps, &sender)?),
         QueryMsg::Stage { stage_id } => to_json_binary(&query_stage(deps, stage_id)?),
         QueryMsg::Stages {} => to_json_binary(&query_stage_list(deps)?),
-        QueryMsg::MerkleRoot {} => to_json_binary(&query_merkle_root(deps)?),
-        QueryMsg::MerkleTreeURI {} => to_json_binary(&query_merkle_tree_uri(deps)?),
+        QueryMsg::MerkleRoots {} => to_json_binary(&query_merkle_roots(deps)?),
+        QueryMsg::MerkleTreeURIs {} => to_json_binary(&query_merkle_tree_uris(deps)?),
     }
 }
 
@@ -330,15 +330,15 @@ pub fn query_config(deps: Deps, env: Env) -> StdResult<ConfigResponse> {
     }
 }
 
-pub fn query_merkle_root(deps: Deps) -> StdResult<MerkleRootResponse> {
+pub fn query_merkle_roots(deps: Deps) -> StdResult<MerkleRootResponse> {
     Ok(MerkleRootResponse {
-        merkle_root: MERKLE_ROOTS.load(deps.storage)?,
+        merkle_roots: MERKLE_ROOTS.load(deps.storage)?,
     })
 }
 
-pub fn query_merkle_tree_uri(deps: Deps) -> StdResult<MerkleTreeURIResponse> {
+pub fn query_merkle_tree_uris(deps: Deps) -> StdResult<MerkleTreeURIResponse> {
     Ok(MerkleTreeURIResponse {
-        merkle_tree_uri: MERKLE_TREE_URIS.may_load(deps.storage)?,
+        merkle_tree_uris: MERKLE_TREE_URIS.may_load(deps.storage)?,
     })
 }
 
@@ -348,8 +348,10 @@ pub fn query_stage(deps: Deps, stage_id: u32) -> StdResult<StageResponse> {
         stage_id < config.stages.len() as u32,
         StdError::generic_err("Stage not found")
     );
+    let merkle_root = MERKLE_ROOTS.load(deps.storage)?[stage_id as usize].clone();
     Ok(StageResponse {
         stage: config.stages[stage_id as usize].clone(),
+        merkle_root,
     })
 }
 
@@ -359,8 +361,10 @@ pub fn query_stage_list(deps: Deps) -> StdResult<StagesResponse> {
         !config.stages.is_empty(),
         StdError::generic_err("No stages found")
     );
+    let merkle_roots = MERKLE_ROOTS.load(deps.storage)?;
     Ok(StagesResponse {
         stages: config.stages.clone(),
+        merkle_roots,
     })
 }
 

diff --git a/contracts/whitelists/tiered-whitelist-merkletree/src/msg.rs b/contracts/whitelists/tiered-whitelist-merkletree/src/msg.rs
@@ -72,9 +72,9 @@ pub enum QueryMsg {
         msg: CosmosMsg<Empty>,
     },
     #[returns(MerkleRootResponse)]
-    MerkleRoot {},
+    MerkleRoots {},
     #[returns(MerkleTreeURIResponse)]
-    MerkleTreeURI {},
+    MerkleTreeURIs {},
 }
 
 #[cw_serde]
@@ -115,12 +115,12 @@ pub struct ConfigResponse {
 
 #[cw_serde]
 pub struct MerkleRootResponse {
-    pub merkle_root: Vec<String>,
+    pub merkle_roots: Vec<String>,
 }
 
 #[cw_serde]
 pub struct MerkleTreeURIResponse {
-    pub merkle_tree_uri: Option<Vec<String>>,
+    pub merkle_tree_uris: Option<Vec<String>>,
 }
 
 #[cw_serde]
@@ -139,9 +139,11 @@ pub struct CanExecuteResponse {
 #[cw_serde]
 pub struct StageResponse {
     pub stage: Stage,
+    pub merkle_root: String,
 }
 
 #[cw_serde]
 pub struct StagesResponse {
     pub stages: Vec<Stage>,
+    pub merkle_roots: Vec<String>,
 }
diff --git a/contracts/whitelists/tiered-whitelist/src/contract.rs b/contracts/whitelists/tiered-whitelist/src/contract.rs
@@ -7,7 +7,7 @@ use crate::helpers::{fetch_active_stage, fetch_active_stage_index, validate_stag
 use crate::msg::{
     AddMembersMsg, ConfigResponse, ExecuteMsg, HasEndedResponse, HasMemberResponse,
     HasStartedResponse, InstantiateMsg, IsActiveResponse, MembersResponse, QueryMsg,
-    RemoveMembersMsg, StageResponse, StagesResponse, UpdateStageConfigMsg,
+    RemoveMembersMsg, StageMemberInfoResponse, StageResponse, StagesResponse, UpdateStageConfigMsg,
 };
 use crate::state::{AdminList, Config, Stage, ADMIN_LIST, CONFIG, WHITELIST_STAGES};
 #[cfg(not(feature = "library"))]
@@ -396,6 +396,9 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
             to_json_binary(&fetch_active_stage_index(deps.storage, &env).map_or(0, |i| i + 1))
         }
         QueryMsg::HasMember { member } => to_json_binary(&query_has_member(deps, env, member)?),
+        QueryMsg::StageMemberInfo { stage_id, member } => {
+            to_json_binary(&query_stage_member_info(deps, stage_id, member)?)
+        }
         QueryMsg::Config {} => to_json_binary(&query_config(deps, env)?),
         QueryMsg::Stage { stage_id } => to_json_binary(&query_stage(deps, stage_id)?),
         QueryMsg::Stages {} => to_json_binary(&query_stage_list(deps)?),
@@ -456,6 +459,21 @@ pub fn query_has_member(deps: Deps, env: Env, member: String) -> StdResult<HasMe
     Ok(HasMemberResponse { has_member })
 }
 
+pub fn query_stage_member_info(
+    deps: Deps,
+    stage_id: u32,
+    member: String,
+) -> StdResult<StageMemberInfoResponse> {
+    let addr = deps.api.addr_validate(&member)?;
+    let stage = CONFIG.load(deps.storage)?;
+    let stage = stage.stages[stage_id as usize].clone();
+    let is_member = WHITELIST_STAGES.has(deps.storage, (stage_id, addr.clone()));
+    Ok(StageMemberInfoResponse {
+        is_member,
+        per_address_limit: stage.per_address_limit,
+    })
+}
+
 pub fn query_config(deps: Deps, env: Env) -> StdResult<ConfigResponse> {
     let config = CONFIG.load(deps.storage)?;
     let active_stage = fetch_active_stage(deps.storage, &env);

diff --git a/contracts/whitelists/tiered-whitelist/src/msg.rs b/contracts/whitelists/tiered-whitelist/src/msg.rs
@@ -76,6 +76,10 @@ pub enum QueryMsg {
     HasMember {
         member: String,
     },
+    StageMemberInfo {
+        member: String,
+        stage_id: u32,
+    },
     Config {},
 
     Stage {
@@ -155,3 +159,9 @@ pub struct StageResponse {
 pub struct StagesResponse {
     pub stages: Vec<Stage>,
 }
+
+#[cw_serde]
+pub struct StageMemberInfoResponse {
+    pub is_member: bool,
+    pub per_address_limit: u32,
+}