Skip to content

Commit

Permalink
Update stage member count logic
Browse files Browse the repository at this point in the history
  • Loading branch information
MightOfOaks committed Nov 6, 2024
1 parent b5608fb commit f1ff48c
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 13 deletions.
33 changes: 20 additions & 13 deletions contracts/whitelists/tiered-whitelist/src/contract.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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)?;
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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,
Expand All @@ -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()
Expand Down Expand Up @@ -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)?;
Expand Down Expand Up @@ -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),
})
}

Expand All @@ -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)
}
2 changes: 2 additions & 0 deletions contracts/whitelists/tiered-whitelist/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");

0 comments on commit f1ff48c

Please sign in to comment.