From 263973a814b9844d4e02f8464ccfe860aeb2b4fd Mon Sep 17 00:00:00 2001 From: Abhishek-1857 Date: Fri, 4 Oct 2024 12:01:48 +0530 Subject: [PATCH] voting testcases implementaion --- contracts/external/cw4-group/src/contract.rs | 50 +- contracts/external/cw4-group/src/state.rs | 141 ++- .../external/snip721-roles/src/contract.rs | 2 +- contracts/external/snip721-roles/src/state.rs | 153 ++- .../staking/snip20-stake/src/contract.rs | 2 +- contracts/staking/snip20-stake/src/state.rs | 94 +- contracts/voting/dao-voting-cw4/src/lib.rs | 4 +- contracts/voting/dao-voting-cw4/src/tests.rs | 1022 ++++++++++++++--- .../dao-voting-snip721-staked/src/state.rs | 104 +- .../dao-voting-token-staked/src/state.rs | 117 +- 10 files changed, 1268 insertions(+), 421 deletions(-) diff --git a/contracts/external/cw4-group/src/contract.rs b/contracts/external/cw4-group/src/contract.rs index 6f046a5..69195ad 100644 --- a/contracts/external/cw4-group/src/contract.rs +++ b/contracts/external/cw4-group/src/contract.rs @@ -228,7 +228,6 @@ pub fn query_member(deps: Deps, addr: Addr, height: Option) -> StdResult, limit: Option, ) -> StdResult { + // Determine the limit, ensuring it's between DEFAULT_LIMIT and MAX_LIMIT let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT) as usize; - let mut res_members: Vec = Vec::new(); + let mut res_members: Vec = Vec::with_capacity(limit); // Pre-allocate capacity - let mut start = start_after.clone(); // Clone start_after to mutate it if necessary - - let binding = &MEMBERS_PRIMARY; + let binding = MEMBERS_PRIMARY; let iter = binding.iter(deps.storage)?; + + // Convert `start_after` to Option<&str> for comparison without unnecessary cloning + let mut start_found = start_after.is_none(); + for item in iter { let (address, weight) = item?; - if let Some(start_after) = &start { - if &address == start_after { - // If we found the start point, reset it to start iterating - start = None; + + // Skip items until we find the one *after* `start_after` + if let Some(ref start_after_addr) = start_after { + if !start_found { + // Check if the current address matches the start_after value + if &address == start_after_addr { + start_found = true; // We've found the start_after value, start collecting from the next item + } + continue; } } - if start.is_none() { - res_members.push(Member { - addr: address.to_string(), - weight, - }); - if res_members.len() >= limit { - break; // Break out of loop if limit reached - } + + // Once we've found the start or if no `start_after` was provided, collect the items + res_members.push(Member { + addr: address.to_string(), + weight, + }); + + // Break when we've collected enough members + if res_members.len() >= limit { + break; } } - let response = MemberListResponse { + // Return the list of members + Ok(MemberListResponse { members: res_members, - }; - - Ok(response) + }) } pub fn authenticate(deps: Deps, auth: Auth, query_auth: Contract) -> StdResult { diff --git a/contracts/external/cw4-group/src/state.rs b/contracts/external/cw4-group/src/state.rs index 7f48caf..004e3fd 100644 --- a/contracts/external/cw4-group/src/state.rs +++ b/contracts/external/cw4-group/src/state.rs @@ -13,11 +13,12 @@ pub const HOOKS: Hooks = Hooks::new("cw4-hooks"); pub const QUERY_AUTH: Item = Item::new("query_auth"); /// A historic list of members and total voting weights -pub static MEMBERS_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); -pub static MEMBERS_SNAPSHOT: Keymap<(u64, Addr), u64> = Keymap::new(b"staked_balances_snapshot"); -pub static MEMBERS_AT_HEIGHT: Keymap> = Keymap::new(b"user_Staked_at_height"); +pub const MEMBERS_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); +pub const MEMBERS_SNAPSHOT: Keymap<(u64, Addr), u64> = Keymap::new(b"staked_balances_snapshot"); +pub const MEMBERS_AT_HEIGHT: Keymap> = Keymap::new(b"user_staked_at_height"); pub struct MembersStore {} + impl MembersStore { // Function to store a value at a specific block height pub fn save( @@ -27,117 +28,169 @@ impl MembersStore { value: u64, ) -> StdResult<()> { let default: u64 = 0; + + // Load the current primary value, if it exists let primary = MEMBERS_PRIMARY.get(store, &key.clone()); + if primary.is_none() { + // First time staking for this user MEMBERS_PRIMARY.insert(store, &key.clone(), &value)?; MEMBERS_SNAPSHOT.insert(store, &(block_height, key.clone()), &default)?; MEMBERS_AT_HEIGHT.insert(store, &key.clone(), &vec![block_height])?; } else { - let mut user_staked_height = MEMBERS_AT_HEIGHT.get(store, &key.clone()).unwrap(); + // Update staking info for an existing user + let mut user_staked_height = MEMBERS_AT_HEIGHT + .get(store, &key.clone()) + .unwrap_or_default(); + + // Store the old primary value as a snapshot at the current block height MEMBERS_SNAPSHOT.insert(store, &(block_height, key.clone()), &primary.unwrap())?; MEMBERS_PRIMARY.insert(store, &key.clone(), &value)?; - user_staked_height.push(block_height); + + // Add the new block height to the user's list of heights if it's not a duplicate + if !user_staked_height.contains(&block_height) { + user_staked_height.push(block_height); + } MEMBERS_AT_HEIGHT.insert(store, &key.clone(), &user_staked_height)?; } Ok(()) } + // Function to load the current staking balance of a user pub fn load(store: &dyn Storage, key: Addr) -> u64 { MEMBERS_PRIMARY.get(store, &key).unwrap_or_default() } + // Function to load the staking balance of a user at a specific height pub fn may_load_at_height( store: &dyn Storage, key: Addr, height: u64, ) -> StdResult> { let snapshot_key = (height, key.clone()); - let snapshot_value = MEMBERS_SNAPSHOT.get(store, &snapshot_key); - if snapshot_value.is_none() { + + // If there's a snapshot at the exact height, return it + if snapshot_value.is_some() { + return Ok(snapshot_value); + } + + // Otherwise, get the list of heights this user has staked at + let user_staked_heights = MEMBERS_AT_HEIGHT.get(store, &key).unwrap_or_default(); + + // If the user has never staked, return the current primary balance + if user_staked_heights.is_empty() { + return Ok(MEMBERS_PRIMARY.get(store, &key)); + } + + // Find the latest block height before or at the specified height + let index = match user_staked_heights.binary_search(&height) { + Ok(i) => i, // exact match found + Err(i) => i.saturating_sub(1), // closest lower height, handle case where i = 0 + }; + + // If the height is beyond the last checkpoint, return the primary balance + if index == user_staked_heights.len() - 1 { Ok(MEMBERS_PRIMARY.get(store, &key)) } else { - let x = MEMBERS_AT_HEIGHT.get(store, &key).unwrap(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(MEMBERS_PRIMARY.get(store, &key)) - } else { - let snapshot_value = - MEMBERS_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize], key.clone())); - Ok(snapshot_value) - } + // Otherwise, return the snapshot at the closest lower height + let snapshot_height = user_staked_heights[index]; + Ok(MEMBERS_SNAPSHOT.get(store, &(snapshot_height, key))) } } + + // Function to remove a user's staking data pub fn remove(store: &mut dyn Storage, key: Addr) -> StdResult<()> { - // Remove the member's data from all storage maps + // Get the user's staked heights before removing their data + let user_staked_height = MEMBERS_AT_HEIGHT.get(store, &key).unwrap_or_default(); + + // Remove the primary and height data MEMBERS_PRIMARY.remove(store, &key)?; MEMBERS_AT_HEIGHT.remove(store, &key)?; - // Remove all snapshot entries associated with the member - let user_staked_height = MEMBERS_AT_HEIGHT.get(store, &key).unwrap_or_default(); + // Remove all snapshot entries associated with the user for height in user_staked_height { MEMBERS_SNAPSHOT.remove(store, &(height, key.clone()))?; } - // Return Ok(()) if all removals were successful Ok(()) } } /// A historic snapshot of total weight over time pub const TOTAL_PRIMARY: Item = Item::new("staked_balances_primary"); -pub static TOTAL_SNAPSHOT: Keymap = Keymap::new(b"staked_balances_snapshot"); -pub const TOTAL_AT_HEIGHTS: Item> = Item::new("user_Staked_at_height"); +pub const TOTAL_SNAPSHOT: Keymap = Keymap::new(b"staked_balances_snapshot"); +pub const TOTAL_AT_HEIGHTS: Item> = Item::new("user_staked_at_height"); pub struct TotalStore {} + impl TotalStore { // Function to store a value at a specific block height pub fn save(store: &mut dyn Storage, block_height: u64, value: u64) -> StdResult<()> { let default: u64 = 0; + + // Load the current primary value, or default to 0 if not found let primary = TOTAL_PRIMARY.load(store).unwrap_or_default(); + if primary == 0 { + // First time storing total weight TOTAL_PRIMARY.save(store, &value)?; TOTAL_SNAPSHOT.insert(store, &block_height, &default)?; TOTAL_AT_HEIGHTS.save(store, &vec![block_height])?; } else { - let mut user_staked_height = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); + // Update total weight for existing entries + let mut staked_heights = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); + + // Insert the old primary value into the snapshot at the current block height TOTAL_SNAPSHOT.insert(store, &block_height, &primary)?; TOTAL_PRIMARY.save(store, &value)?; - user_staked_height.push(block_height); - TOTAL_AT_HEIGHTS.save(store, &user_staked_height)?; + + // Add the new block height if it's not a duplicate + if !staked_heights.contains(&block_height) { + staked_heights.push(block_height); + } + TOTAL_AT_HEIGHTS.save(store, &staked_heights)?; } Ok(()) } + // Function to load the current total staked value pub fn load(store: &dyn Storage) -> u64 { TOTAL_PRIMARY.load(store).unwrap_or_default() } + // Function to load the total staked value at a specific block height pub fn may_load_at_height(store: &dyn Storage, height: u64) -> StdResult> { - let snapshot_key = height; + let snapshot_value = TOTAL_SNAPSHOT.get(store, &height); + + // If a snapshot exists at the exact height, return it + if snapshot_value.is_some() { + return Ok(snapshot_value); + } + + // Otherwise, load all the heights where total staked was recorded + let staked_heights = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); - let snapshot_value = TOTAL_SNAPSHOT.get(store, &snapshot_key); - if snapshot_value.is_none() { + // If there are no staked heights, return the current total primary value + if staked_heights.is_empty() { + return Ok(Some(TOTAL_PRIMARY.load(store).unwrap_or_default())); + } + + // Find the latest block height before or at the specified height + let index = match staked_heights.binary_search(&height) { + Ok(i) => i, // Exact match found + Err(i) => i.saturating_sub(1), // Closest lower height, handle case where i = 0 + }; + + // If the height is beyond the last checkpoint, return the current primary value + if index == staked_heights.len() - 1 { Ok(Some(TOTAL_PRIMARY.load(store).unwrap_or_default())) } else { - let x = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(Some(TOTAL_PRIMARY.load(store).unwrap_or_default())) - } else { - let snapshot_value = TOTAL_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize])); - Ok(snapshot_value) - } + // Otherwise, return the snapshot at the closest lower height + let snapshot_height = staked_heights[index]; + Ok(TOTAL_SNAPSHOT.get(store, &snapshot_height)) } } } diff --git a/contracts/external/snip721-roles/src/contract.rs b/contracts/external/snip721-roles/src/contract.rs index c3d958c..ba199b2 100644 --- a/contracts/external/snip721-roles/src/contract.rs +++ b/contracts/external/snip721-roles/src/contract.rs @@ -640,7 +640,7 @@ pub fn query_list_members( let mut start = start_after.clone(); // Clone start_after to mutate it if necessary - let binding = &MEMBERS_PRIMARY; + let binding = MEMBERS_PRIMARY; let iter = binding.iter(deps.storage)?; for item in iter { let (address, weight) = item?; diff --git a/contracts/external/snip721-roles/src/state.rs b/contracts/external/snip721-roles/src/state.rs index 59ba888..7975fa7 100644 --- a/contracts/external/snip721-roles/src/state.rs +++ b/contracts/external/snip721-roles/src/state.rs @@ -32,11 +32,12 @@ pub const SNIP721_INFO: Item = Item::new("si"); // ); /// A historic list of members and total voting weights -pub static MEMBERS_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); -pub static MEMBERS_SNAPSHOT: Keymap<(u64, Addr), u64> = Keymap::new(b"staked_balances_snapshot"); -pub static MEMBERS_AT_HEIGHT: Keymap> = Keymap::new(b"user_Staked_at_height"); +pub const MEMBERS_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); +pub const MEMBERS_SNAPSHOT: Keymap<(u64, Addr), u64> = Keymap::new(b"staked_balances_snapshot"); +pub const MEMBERS_AT_HEIGHT: Keymap> = Keymap::new(b"user_staked_at_height"); pub struct MembersStore {} + impl MembersStore { // Function to store a value at a specific block height pub fn save( @@ -46,117 +47,171 @@ impl MembersStore { value: u64, ) -> StdResult<()> { let default: u64 = 0; - let primary = MEMBERS_PRIMARY.get(store, &key.clone()); + + // Load the current primary value + let primary = MEMBERS_PRIMARY.get(store, &key); + if primary.is_none() { - MEMBERS_PRIMARY.insert(store, &key.clone(), &value)?; + // First time staking for this user + MEMBERS_PRIMARY.insert(store, &key, &value)?; MEMBERS_SNAPSHOT.insert(store, &(block_height, key.clone()), &default)?; - MEMBERS_AT_HEIGHT.insert(store, &key.clone(), &vec![block_height])?; + MEMBERS_AT_HEIGHT.insert(store, &key, &vec![block_height])?; } else { - let mut user_staked_height = MEMBERS_AT_HEIGHT.get(store, &key.clone()).unwrap(); + // Update staking info for an existing user + let mut user_staked_height = MEMBERS_AT_HEIGHT.get(store, &key).unwrap_or_default(); + + // Insert the old primary value as a snapshot at the given block height MEMBERS_SNAPSHOT.insert(store, &(block_height, key.clone()), &primary.unwrap())?; - MEMBERS_PRIMARY.insert(store, &key.clone(), &value)?; - user_staked_height.push(block_height); - MEMBERS_AT_HEIGHT.insert(store, &key.clone(), &user_staked_height)?; + MEMBERS_PRIMARY.insert(store, &key, &value)?; + + // Add the block height if it's not already present + if !user_staked_height.contains(&block_height) { + user_staked_height.push(block_height); + } + + MEMBERS_AT_HEIGHT.insert(store, &key, &user_staked_height)?; } Ok(()) } + // Function to load the current staking balance of a user pub fn load(store: &dyn Storage, key: Addr) -> u64 { MEMBERS_PRIMARY.get(store, &key).unwrap_or_default() } + // Function to load the staking balance of a user at a specific block height pub fn may_load_at_height( store: &dyn Storage, key: Addr, height: u64, ) -> StdResult> { let snapshot_key = (height, key.clone()); - let snapshot_value = MEMBERS_SNAPSHOT.get(store, &snapshot_key); - if snapshot_value.is_none() { + + // If there's a snapshot at the exact height, return it + if snapshot_value.is_some() { + return Ok(snapshot_value); + } + + // Load all the heights where the user has staked + let user_staked_heights = MEMBERS_AT_HEIGHT.get(store, &key).unwrap_or_default(); + + // If there are no staked heights, return the current primary balance + if user_staked_heights.is_empty() { + return Ok(MEMBERS_PRIMARY.get(store, &key)); + } + + // Find the latest block height before or at the given height + let index = match user_staked_heights.binary_search(&height) { + Ok(i) => i, // Exact match + Err(i) => i.saturating_sub(1), // Closest lower height, safe from underflow + }; + + // If the height is beyond the last checkpoint, return the current primary balance + if index == user_staked_heights.len() - 1 { Ok(MEMBERS_PRIMARY.get(store, &key)) } else { - let x = MEMBERS_AT_HEIGHT.get(store, &key).unwrap(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(MEMBERS_PRIMARY.get(store, &key)) - } else { - let snapshot_value = - MEMBERS_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize], key.clone())); - Ok(snapshot_value) - } + // Otherwise, return the snapshot at the closest lower height + let snapshot_height = user_staked_heights[index]; + Ok(MEMBERS_SNAPSHOT.get(store, &(snapshot_height, key))) } } + + // Function to remove a user's staking data pub fn remove(store: &mut dyn Storage, key: Addr) -> StdResult<()> { - // Remove the member's data from all storage maps + // Get the user's staked heights before removing their data + let user_staked_heights = MEMBERS_AT_HEIGHT.get(store, &key).unwrap_or_default(); + + // Remove the primary and height data MEMBERS_PRIMARY.remove(store, &key)?; MEMBERS_AT_HEIGHT.remove(store, &key)?; - // Remove all snapshot entries associated with the member - let user_staked_height = MEMBERS_AT_HEIGHT.get(store, &key).unwrap_or_default(); - for height in user_staked_height { + // Remove all snapshot entries associated with the user + for height in user_staked_heights { MEMBERS_SNAPSHOT.remove(store, &(height, key.clone()))?; } - // Return Ok(()) if all removals were successful Ok(()) } } /// A historic snapshot of total weight over time pub const TOTAL_PRIMARY: Item = Item::new("staked_balances_primary"); -pub static TOTAL_SNAPSHOT: Keymap = Keymap::new(b"staked_balances_snapshot"); -pub const TOTAL_AT_HEIGHTS: Item> = Item::new("user_Staked_at_height"); +pub const TOTAL_SNAPSHOT: Keymap = Keymap::new(b"staked_balances_snapshot"); +pub const TOTAL_AT_HEIGHTS: Item> = Item::new("total_staked_at_height"); pub struct TotalStore {} + impl TotalStore { // Function to store a value at a specific block height pub fn save(store: &mut dyn Storage, block_height: u64, value: u64) -> StdResult<()> { let default: u64 = 0; let primary = TOTAL_PRIMARY.load(store).unwrap_or_default(); + if primary == 0 { + // First time total weight is stored TOTAL_PRIMARY.save(store, &value)?; TOTAL_SNAPSHOT.insert(store, &block_height, &default)?; + TOTAL_AT_HEIGHTS.save(store, &vec![block_height])?; } else { - let mut user_staked_height = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); + // Update existing total weight + let mut total_staked_height = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); + + // Insert the old primary value as a snapshot at the given block height TOTAL_SNAPSHOT.insert(store, &block_height, &primary)?; + + // Update primary with the new total weight TOTAL_PRIMARY.save(store, &value)?; - user_staked_height.push(block_height); - TOTAL_AT_HEIGHTS.save(store, &user_staked_height)?; + + // Ensure no duplicate heights are inserted + if !total_staked_height.contains(&block_height) { + total_staked_height.push(block_height); + } + + TOTAL_AT_HEIGHTS.save(store, &total_staked_height)?; } Ok(()) } + // Function to load the current total weight pub fn load(store: &dyn Storage) -> u64 { TOTAL_PRIMARY.load(store).unwrap_or_default() } + // Function to load the total weight at a specific block height pub fn may_load_at_height(store: &dyn Storage, height: u64) -> StdResult> { - let snapshot_key = height; + // Try to fetch a snapshot at the exact height + let snapshot_value = TOTAL_SNAPSHOT.get(store, &height); - let snapshot_value = TOTAL_SNAPSHOT.get(store, &snapshot_key); + // If there's no snapshot at the exact height, return the current primary value if snapshot_value.is_none() { + return Ok(Some(TOTAL_PRIMARY.load(store).unwrap_or_default())); + } + + // Fetch the block heights where total weight was updated + let total_staked_heights = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); + + // If no heights exist, return the primary value + if total_staked_heights.is_empty() { + return Ok(Some(TOTAL_PRIMARY.load(store).unwrap_or_default())); + } + + // Find the closest block height (binary search) + let index = match total_staked_heights.binary_search(&height) { + Ok(i) => i, // Exact match + Err(i) => i.saturating_sub(1), // Closest lower height (safe from underflow) + }; + + // If the given height is beyond the last checkpoint, return the current primary value + if index == total_staked_heights.len() - 1 { Ok(Some(TOTAL_PRIMARY.load(store).unwrap_or_default())) } else { - let x = TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(Some(TOTAL_PRIMARY.load(store).unwrap_or_default())) - } else { - let snapshot_value = TOTAL_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize])); - Ok(snapshot_value) - } + // Return the snapshot at the closest lower height + let snapshot_height = total_staked_heights[index]; + Ok(TOTAL_SNAPSHOT.get(store, &snapshot_height)) } } } diff --git a/contracts/staking/snip20-stake/src/contract.rs b/contracts/staking/snip20-stake/src/contract.rs index e01f06d..acb8dbf 100644 --- a/contracts/staking/snip20-stake/src/contract.rs +++ b/contracts/staking/snip20-stake/src/contract.rs @@ -482,7 +482,7 @@ pub fn query_list_stakers( let mut start = start_after.clone(); // Clone start_after to mutate it if necessary - let binding = &STAKED_BALANCES_PRIMARY; + let binding = STAKED_BALANCES_PRIMARY; let iter = binding.iter(deps.storage)?; for item in iter { let (address, stake) = item?; diff --git a/contracts/staking/snip20-stake/src/state.rs b/contracts/staking/snip20-stake/src/state.rs index a8fc61c..54c2817 100644 --- a/contracts/staking/snip20-stake/src/state.rs +++ b/contracts/staking/snip20-stake/src/state.rs @@ -28,38 +28,47 @@ pub const MAX_CLAIMS: u64 = 100; pub const CLAIMS: Claims = Claims::new("claims"); -pub static STAKED_TOTAL_AT_HEIGHT: Keymap = Keymap::new(b"staked_total"); +pub const STAKED_TOTAL_AT_HEIGHT: Keymap = Keymap::new(b"staked_total"); +pub const TOTAL_BALANCE: Item = Item::new("total_balance"); pub struct StakedTotalStore {} + impl StakedTotalStore { // Function to store a value at a specific block height pub fn save(store: &mut dyn Storage, block_height: u64, value: Uint128) -> StdResult<()> { + // Insert the total staked value at the given block height STAKED_TOTAL_AT_HEIGHT.insert(store, &block_height, &value)?; + // Also update the total current balance + TOTAL_BALANCE.save(store, &value)?; Ok(()) } + // Load the most recent total staked balance pub fn load(store: &dyn Storage) -> Uint128 { - BALANCE.load(store).unwrap_or_default() + TOTAL_BALANCE.load(store).unwrap_or_default() } + // Load the staked total at a specific block height, falling back to the most recent total if not found pub fn may_load_at_height(store: &dyn Storage, height: u64) -> StdResult> { - let total_staked_at_height = STAKED_TOTAL_AT_HEIGHT.get(store, &height); - if total_staked_at_height.is_none() { - let res = BALANCE.load(store)?; - Ok(Some(res)) + // Check if there is a total staked value recorded at the given height + if let Some(snapshot_value) = STAKED_TOTAL_AT_HEIGHT.get(store, &height) { + // If found, return the snapshot value + Ok(Some(snapshot_value)) } else { - let snapshot_value = STAKED_TOTAL_AT_HEIGHT.get(store, &height); - Ok(snapshot_value) + // If not found, fallback to the most recent total balance + let total_balance = TOTAL_BALANCE.load(store)?; + Ok(Some(total_balance)) } } } -pub static STAKED_BALANCES_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); -pub static STAKED_BALANCES_SNAPSHOT: Keymap<(u64, Addr), Uint128> = +pub const STAKED_BALANCES_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); +pub const STAKED_BALANCES_SNAPSHOT: Keymap<(u64, Addr), Uint128> = Keymap::new(b"staked_balances_snapshot"); -pub static USER_STAKED_AT_HEIGHT: Keymap> = Keymap::new(b"user_Staked_at_height"); +pub const USER_STAKED_AT_HEIGHT: Keymap> = Keymap::new(b"user_staked_at_height"); pub struct StakedBalancesStore {} + impl StakedBalancesStore { // Function to store a value at a specific block height pub fn save( @@ -68,58 +77,79 @@ impl StakedBalancesStore { key: Addr, value: Uint128, ) -> StdResult<()> { - let primary = STAKED_BALANCES_PRIMARY.get(store, &key.clone()); + let primary = STAKED_BALANCES_PRIMARY.get(store, &key); + if primary.is_none() { - STAKED_BALANCES_PRIMARY.insert(store, &key.clone(), &value)?; + // First time staking for this user + STAKED_BALANCES_PRIMARY.insert(store, &key, &value)?; STAKED_BALANCES_SNAPSHOT.insert( store, &(block_height, key.clone()), &Uint128::zero(), )?; - USER_STAKED_AT_HEIGHT.insert(store, &key.clone(), &vec![block_height])?; + USER_STAKED_AT_HEIGHT.insert(store, &key, &vec![block_height])?; } else { - let mut user_staked_height = USER_STAKED_AT_HEIGHT.get(store, &key.clone()).unwrap(); + // User has staked before, create a snapshot + let mut user_staked_height = USER_STAKED_AT_HEIGHT.get(store, &key).unwrap_or_default(); STAKED_BALANCES_SNAPSHOT.insert( store, &(block_height, key.clone()), &primary.unwrap(), )?; - STAKED_BALANCES_PRIMARY.insert(store, &key.clone(), &value)?; - user_staked_height.push(block_height); - USER_STAKED_AT_HEIGHT.insert(store, &key.clone(), &user_staked_height)?; + STAKED_BALANCES_PRIMARY.insert(store, &key, &value)?; + + // Ensure block height is not duplicated + if !user_staked_height.contains(&block_height) { + user_staked_height.push(block_height); + } + + USER_STAKED_AT_HEIGHT.insert(store, &key, &user_staked_height)?; } Ok(()) } + // Load the primary staked balance for the given user pub fn load(store: &dyn Storage, key: Addr) -> Uint128 { STAKED_BALANCES_PRIMARY.get(store, &key).unwrap_or_default() } + // Load staked balance at a specific height pub fn may_load_at_height( store: &dyn Storage, key: Addr, height: u64, ) -> StdResult> { + // Try to get the snapshot value at the given height let snapshot_key = (height, key.clone()); - let snapshot_value = STAKED_BALANCES_SNAPSHOT.get(store, &snapshot_key); + if snapshot_value.is_none() { + // No snapshot, fallback to the current primary balance + return Ok(STAKED_BALANCES_PRIMARY.get(store, &key)); + } + + // Get all heights at which the user staked + let user_staked_heights = USER_STAKED_AT_HEIGHT.get(store, &key).unwrap_or_default(); + + // If no staked heights exist, return primary value + if user_staked_heights.is_empty() { + return Ok(STAKED_BALANCES_PRIMARY.get(store, &key)); + } + + // Find the closest block height + let index = match user_staked_heights.binary_search(&height) { + Ok(i) => i, // Exact match + Err(i) => i.saturating_sub(1), // Closest lower height (prevent underflow) + }; + + // If we're querying for the most recent height, return primary + if index == user_staked_heights.len() - 1 { Ok(STAKED_BALANCES_PRIMARY.get(store, &key)) } else { - let x = USER_STAKED_AT_HEIGHT.get(store, &key).unwrap(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(STAKED_BALANCES_PRIMARY.get(store, &key)) - } else { - let snapshot_value = - STAKED_BALANCES_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize], key.clone())); - Ok(snapshot_value) - } + // Return snapshot at the closest lower height + let snapshot_height = user_staked_heights[index + 1]; + Ok(STAKED_BALANCES_SNAPSHOT.get(store, &(snapshot_height, key.clone()))) } } } diff --git a/contracts/voting/dao-voting-cw4/src/lib.rs b/contracts/voting/dao-voting-cw4/src/lib.rs index 2baf24d..f88d6a9 100644 --- a/contracts/voting/dao-voting-cw4/src/lib.rs +++ b/contracts/voting/dao-voting-cw4/src/lib.rs @@ -6,7 +6,7 @@ mod error; pub mod msg; pub mod state; -// #[cfg(test)] -// mod tests; +#[cfg(test)] +mod tests; pub use crate::error::ContractError; diff --git a/contracts/voting/dao-voting-cw4/src/tests.rs b/contracts/voting/dao-voting-cw4/src/tests.rs index 41d9885..7c3add6 100644 --- a/contracts/voting/dao-voting-cw4/src/tests.rs +++ b/contracts/voting/dao-voting-cw4/src/tests.rs @@ -1,18 +1,29 @@ -use cosmwasm_std::{from_binary, to_binary, Addr, ContractInfo, Empty, MessageInfo}; -use dao_interface::{state::AnyContractInfo, voting::InfoResponse}; +use cosmwasm_std::{ + from_binary, + testing::{mock_dependencies, mock_env, mock_info}, + to_binary, Addr, ContractInfo, CosmosMsg, Empty, MessageInfo, Uint128, WasmMsg, +}; +use dao_interface::{ + state::AnyContractInfo, + voting::{InfoResponse, TotalPowerAtHeightResponse, VotingPowerAtHeightResponse}, +}; use secret_cw2::ContractVersion; -use secret_multi_test::{App, Contract, ContractInstantiationInfo, ContractWrapper, Executor}; -use shade_protocol::utils::asset::RawContract; +use secret_multi_test::{ + next_block, App, Contract, ContractInstantiationInfo, ContractWrapper, Executor, +}; +use shade_protocol::{basic_staking::Auth, utils::asset::RawContract}; -use crate::msg::{GroupContract, InstantiateMsg, QueryMsg}; +use crate::{ + contract::{migrate, CONTRACT_NAME, CONTRACT_VERSION}, + msg::{GroupContract, InstantiateMsg, MigrateMsg, QueryMsg}, + ContractError, +}; const DAO_ADDR: &str = "dao"; const ADDR1: &str = "addr1"; const ADDR2: &str = "addr2"; const ADDR3: &str = "addr3"; -#[allow(dead_code)] const ADDR4: &str = "addr4"; -const OWNER: &str = "owner"; fn cw4_contract() -> Box> { let contract = ContractWrapper::new( @@ -22,14 +33,6 @@ fn cw4_contract() -> Box> { ); Box::new(contract) } -fn query_auth_contract() -> Box> { - let contract = ContractWrapper::new( - query_auth::contract::execute, - query_auth::contract::instantiate, - query_auth::contract::query, - ); - Box::new(contract) -} fn voting_contract() -> Box> { let contract = ContractWrapper::new( @@ -37,26 +40,18 @@ fn voting_contract() -> Box> { crate::contract::instantiate, crate::contract::query, ) - // .with_reply(crate::contract::reply) + .with_reply(crate::contract::reply) .with_migrate(crate::contract::migrate); Box::new(contract) } -#[allow(dead_code)] -fn instantiate_voting( - app: &mut App, - contract_instantiate_info: ContractInstantiationInfo, - msg: InstantiateMsg, -) -> ContractInfo { - app.instantiate_contract( - contract_instantiate_info, - Addr::unchecked(DAO_ADDR), - &msg, - &[], - "voting module", - None, - ) - .unwrap() +fn query_auth_contract() -> Box> { + let contract = ContractWrapper::new( + query_auth::contract::execute, + query_auth::contract::instantiate, + query_auth::contract::query, + ); + Box::new(contract) } fn instantiate_query_auth(app: &mut App) -> ContractInfo { @@ -71,7 +66,7 @@ fn instantiate_query_auth(app: &mut App) -> ContractInfo { app.instantiate_contract( query_auth_info, - Addr::unchecked(OWNER), + Addr::unchecked(DAO_ADDR), &msg, &[], "query_auth", @@ -80,7 +75,7 @@ fn instantiate_query_auth(app: &mut App) -> ContractInfo { .unwrap() } -fn _create_viewing_key(app: &mut App, contract_info: ContractInfo, info: MessageInfo) -> String { +fn create_viewing_key(app: &mut App, contract_info: ContractInfo, info: MessageInfo) -> String { let msg = shade_protocol::contract_interfaces::query_auth::ExecuteMsg::CreateViewingKey { entropy: "entropy".to_string(), padding: None, @@ -100,11 +95,25 @@ fn _create_viewing_key(app: &mut App, contract_info: ContractInfo, info: Message viewing_key } -fn _setup_test_case(app: &mut App) -> ContractInfo { - let cw4_instantiate_info = app.store_code(cw4_contract()); - let voting_instantiate_info = app.store_code(voting_contract()); +fn instantiate_voting( + app: &mut App, + voting_info: ContractInstantiationInfo, + msg: InstantiateMsg, +) -> ContractInfo { + app.instantiate_contract( + voting_info, + Addr::unchecked(DAO_ADDR), + &msg, + &[], + "voting module", + None, + ) + .unwrap() +} - let query_auth = instantiate_query_auth(app); +fn setup_test_case(app: &mut App) -> (ContractInfo, ContractInfo) { + let cw4_info = app.store_code(cw4_contract()); + let voting_info = app.store_code(voting_contract()); let members = vec![ cw4::Member { @@ -124,53 +133,184 @@ fn _setup_test_case(app: &mut App) -> ContractInfo { weight: 0, }, ]; - instantiate_voting( - app, - voting_instantiate_info, - InstantiateMsg { - group_contract: GroupContract::New { - cw4_group_code_id: cw4_instantiate_info.code_id, - cw4_group_code_hash: cw4_instantiate_info.code_hash, - initial_members: members, - query_auth: Some(RawContract { - address: query_auth.address.to_string(), - code_hash: query_auth.code_hash, - }), + + let query_auth = instantiate_query_auth(app); + + ( + instantiate_voting( + app, + voting_info, + InstantiateMsg { + group_contract: GroupContract::New { + cw4_group_code_id: cw4_info.code_id, + cw4_group_code_hash: cw4_info.code_hash, + initial_members: members, + query_auth: Some(RawContract::new( + &query_auth.address.clone().to_string(), + &query_auth.code_hash.clone().to_string(), + )), + }, + dao_code_hash: "todo!()".into(), }, - dao_code_hash: "dao_code_hash".to_string(), - }, + ), + query_auth, ) } #[test] fn test_instantiate() { let mut app = App::default(); + // Valid instantiate no panics + let _voting_contract_info = setup_test_case(&mut app); // Instantiate with no members, error - let voting_instantiate_info = app.store_code(voting_contract()); - let cw4_instantiate_info = app.store_code(cw4_contract()); - - let query_auth = instantiate_query_auth(&mut app); + let voting_info = app.store_code(voting_contract()); + let cw4_info = app.store_code(cw4_contract()); + let msg = InstantiateMsg { + group_contract: GroupContract::New { + cw4_group_code_id: cw4_info.code_id, + cw4_group_code_hash: cw4_info.code_hash.clone(), + initial_members: [].into(), + query_auth: None, + }, + dao_code_hash: "".to_string(), + }; + let _err = app + .instantiate_contract( + voting_info.clone(), + Addr::unchecked(DAO_ADDR), + &msg, + &[], + "voting module", + None, + ) + .unwrap_err(); + // Instantiate with members but no weight let msg = InstantiateMsg { group_contract: GroupContract::New { - cw4_group_code_id: cw4_instantiate_info.clone().code_id, - cw4_group_code_hash: cw4_instantiate_info.clone().code_hash, + cw4_group_code_id: cw4_info.code_id, + cw4_group_code_hash: cw4_info.code_hash.clone(), initial_members: vec![ cw4::Member { + addr: ADDR1.to_string(), + weight: 0, + }, + cw4::Member { + addr: ADDR2.to_string(), + weight: 0, + }, + cw4::Member { + addr: ADDR3.to_string(), + weight: 0, + }, + ], + query_auth: None, + }, + dao_code_hash: "".to_string(), + }; + let _err = app + .instantiate_contract( + voting_info, + Addr::unchecked(DAO_ADDR), + &msg, + &[], + "voting module", + None, + ) + .unwrap_err(); +} + +#[test] +pub fn test_instantiate_existing_contract() { + let mut app = App::default(); + + let voting_info = app.store_code(voting_contract()); + let cw4_info = app.store_code(cw4_contract()); + + let query_auth = instantiate_query_auth(&mut app); + + let viewing_key = create_viewing_key( + &mut app, + ContractInfo { + address: query_auth.address.clone(), + code_hash: query_auth.code_hash.clone(), + }, + mock_info(ADDR1, &[]), + ); + + // Fail with no members. + let cw4_contract_info = app + .instantiate_contract( + cw4_info.clone(), + Addr::unchecked(DAO_ADDR), + &cw4_group::msg::InstantiateMsg { + admin: Some(DAO_ADDR.to_string()), + members: vec![], + query_auth: RawContract::new( + &query_auth.address.clone().to_string(), + &query_auth.code_hash.clone().to_string(), + ), + }, + &[], + "cw4 group", + None, + ) + .unwrap(); + + let err: ContractError = app + .instantiate_contract( + voting_info.clone(), + Addr::unchecked(DAO_ADDR), + &InstantiateMsg { + group_contract: GroupContract::Existing { + address: cw4_contract_info.address.clone().to_string(), + code_hash: cw4_contract_info.code_hash.clone(), + }, + dao_code_hash: "todo!()".into(), + }, + &[], + "voting module", + None, + ) + .unwrap_err() + .downcast() + .unwrap(); + + assert_eq!(err, ContractError::NoMembers {}); + + let cw4_contract_info = app + .instantiate_contract( + cw4_info, + Addr::unchecked(DAO_ADDR), + &cw4_group::msg::InstantiateMsg { + admin: Some(DAO_ADDR.to_string()), + members: vec![cw4::Member { addr: ADDR1.to_string(), weight: 1, }], - query_auth: Some(RawContract { - address: query_auth.address.to_string(), - code_hash: query_auth.code_hash.clone(), - }), + query_auth: RawContract::new( + &query_auth.address.clone().to_string(), + &query_auth.code_hash.clone().to_string(), + ), + }, + &[], + "cw4 group", + None, + ) + .unwrap(); + + // Instantiate with existing contract + let msg = InstantiateMsg { + group_contract: GroupContract::Existing { + address: cw4_contract_info.address.clone().to_string(), + code_hash: cw4_contract_info.code_hash.clone(), }, - dao_code_hash: "dao_code_Hash".to_string(), + dao_code_hash: "".into(), }; let _err = app .instantiate_contract( - voting_instantiate_info.clone(), + voting_info, Addr::unchecked(DAO_ADDR), &msg, &[], @@ -179,145 +319,675 @@ fn test_instantiate() { ) .unwrap(); - // // Instantiate with members but no weight - // let msg = InstantiateMsg { - // group_contract: GroupContract::New { - // cw4_group_code_id: cw4_instantiate_info.clone().code_id, - // cw4_group_code_hash: cw4_instantiate_info.clone().code_hash, - // initial_members: vec![ - // cw4::Member { - // addr: ADDR1.to_string(), - // weight: 0, - // }, - // cw4::Member { - // addr: ADDR2.to_string(), - // weight: 0, - // }, - // cw4::Member { - // addr: ADDR3.to_string(), - // weight: 0, - // }, - // ], - // query_auth: Some(RawContract { - // address: query_auth.address.to_string(), - // code_hash: query_auth.code_hash, - // }), - // }, - - // dao_code_hash: "dao_code_hash".to_string(), - // }; - // let _err = app - // .instantiate_contract( - // voting_instantiate_info, - // Addr::unchecked(DAO_ADDR), - // &msg, - // &[], - // "voting module", - // None, - // ) - // .unwrap_err(); + // Update ADDR1's weight to 2 + let msg = cw4_group::msg::ExecuteMsg::UpdateMembers { + remove: vec![], + add: vec![cw4::Member { + addr: ADDR1.to_string(), + weight: 2, + }], + }; + + app.execute_contract( + Addr::unchecked(DAO_ADDR), + &cw4_contract_info.clone(), + &msg, + &[], + ) + .unwrap(); + + // Same should be true about the groups contract. + let cw4_power: cw4::MemberResponse = app + .wrap() + .query_wasm_smart( + cw4_contract_info.code_hash.clone(), + cw4_contract_info.address.clone(), + &cw4::Cw4QueryMsg::Member { + at_height: None, + auth: Auth::ViewingKey { + key: viewing_key.into(), + address: ADDR1.into(), + }, + }, + ) + .unwrap(); + assert_eq!(cw4_power.weight.unwrap(), 2); } -// #[test] -// fn test_contract_info() { -// let mut app = App::default(); +#[test] +fn test_contract_info() { + let mut app = App::default(); + let (voting_contract_info, _) = setup_test_case(&mut app); -// let voting_instantiate_info = app.store_code(voting_contract()); -// let cw4_instantiate_info = app.store_code(cw4_contract()); -// let query_auth = instantiate_query_auth(&mut app); + let info: InfoResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::Info {}, + ) + .unwrap(); + assert_eq!( + info, + InfoResponse { + info: ContractVersion { + contract: "crates.io:dao-voting-cw4".to_string(), + version: env!("CARGO_PKG_VERSION").to_string() + } + } + ); -// let cw4_info_with_member = app -// .instantiate_contract( -// cw4_instantiate_info.clone(), -// Addr::unchecked(DAO_ADDR), -// &cw4_group::msg::InstantiateMsg { -// admin: Some(DAO_ADDR.to_string()), -// members: vec![ -// cw4::Member { -// addr: ADDR1.to_string(), -// weight: 0, -// }, -// cw4::Member { -// addr: ADDR2.to_string(), -// weight: 0, -// }, -// cw4::Member { -// addr: ADDR3.to_string(), -// weight: 0, -// }, -// ], -// query_auth: RawContract { -// address: query_auth.clone().address.to_string(), -// code_hash: query_auth.clone().code_hash, -// }, -// voting_code_hash: cw4_instantiate_info.code_hash.clone(), -// }, -// &[], -// "cw4 group", -// None, -// ) -// .unwrap(); + // Ensure group contract is set + let _group_contract: AnyContractInfo = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::GroupContract {}, + ) + .unwrap(); + + let dao_contract: AnyContractInfo = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash, + voting_contract_info.address, + &QueryMsg::Dao {}, + ) + .unwrap(); + assert_eq!( + dao_contract, + AnyContractInfo { + addr: Addr::unchecked(DAO_ADDR), + code_hash: "todo!()".into(), + } + ); +} + +#[test] +fn test_power_at_height() { + let mut app = App::default(); + let (voting_contract_info, query_auth) = setup_test_case(&mut app); + + let viewing_key = create_viewing_key( + &mut app, + ContractInfo { + address: query_auth.address.clone(), + code_hash: query_auth.code_hash.clone(), + }, + mock_info(ADDR1, &[]), + ); + + let viewing_key_addr2 = create_viewing_key( + &mut app, + ContractInfo { + address: query_auth.address.clone(), + code_hash: query_auth.code_hash.clone(), + }, + mock_info(ADDR2, &[]), + ); + + app.update_block(next_block); + + let cw4_contract_info: AnyContractInfo = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::GroupContract {}, + ) + .unwrap(); + + let addr1_voting_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key.clone(), + address: ADDR1.to_string(), + }, + height: None, + }, + ) + .unwrap(); + assert_eq!(addr1_voting_power.power, Uint128::new(1u128)); + assert_eq!(addr1_voting_power.height, app.block_info().height); + + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { height: None }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(3u128)); + assert_eq!(total_voting_power.height, app.block_info().height); -// // Instantiate with existing contract + // Update ADDR1's weight to 2 + let msg = cw4_group::msg::ExecuteMsg::UpdateMembers { + remove: vec![], + add: vec![cw4::Member { + addr: ADDR1.to_string(), + weight: 2, + }], + }; + + // Should still be one as voting power should not update until + // the following block. + let addr1_voting_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key.clone(), + address: ADDR1.to_string(), + }, + height: None, + }, + ) + .unwrap(); + assert_eq!(addr1_voting_power.power, Uint128::new(1u128)); + + // Same should be true about the groups contract. + let cw4_power: cw4::MemberResponse = app + .wrap() + .query_wasm_smart( + cw4_contract_info.code_hash.clone(), + cw4_contract_info.addr.clone(), + &cw4::Cw4QueryMsg::Member { + auth: Auth::ViewingKey { + key: viewing_key.clone(), + address: ADDR1.to_string(), + }, + at_height: None, + }, + ) + .unwrap(); + assert_eq!(cw4_power.weight.unwrap(), 1); + + app.execute_contract( + Addr::unchecked(DAO_ADDR), + &ContractInfo { + address: cw4_contract_info.addr.clone(), + code_hash: cw4_contract_info.code_hash.clone(), + }, + &msg, + &[], + ) + .unwrap(); + app.update_block(next_block); + + // Should now be 2 + let addr1_voting_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key.clone(), + address: ADDR1.to_string(), + }, + height: None, + }, + ) + .unwrap(); + assert_eq!(addr1_voting_power.power, Uint128::new(2u128)); + assert_eq!(addr1_voting_power.height, app.block_info().height); + + // Check we can still get the 1 weight he had last block + let addr1_voting_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key.clone(), + address: ADDR1.to_string(), + }, + height: Some(app.block_info().height - 1), + }, + ) + .unwrap(); + assert_eq!(addr1_voting_power.power, Uint128::new(1u128)); + assert_eq!(addr1_voting_power.height, app.block_info().height - 1); + + // Check total power is now 4 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { height: None }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(4u128)); + assert_eq!(total_voting_power.height, app.block_info().height); + + // Check total power for last block is 3 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { + height: Some(app.block_info().height - 1), + }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(3u128)); + assert_eq!(total_voting_power.height, app.block_info().height - 1); + + // Update ADDR1's weight back to 1 + let msg = cw4_group::msg::ExecuteMsg::UpdateMembers { + remove: vec![], + add: vec![cw4::Member { + addr: ADDR1.to_string(), + weight: 1, + }], + }; + + app.execute_contract( + Addr::unchecked(DAO_ADDR), + &ContractInfo { + address: cw4_contract_info.addr.clone(), + code_hash: cw4_contract_info.code_hash.clone(), + }, + &msg, + &[], + ) + .unwrap(); + app.update_block(next_block); + + // Should now be 1 again + let addr1_voting_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key.clone(), + address: ADDR1.to_string(), + }, + height: None, + }, + ) + .unwrap(); + assert_eq!(addr1_voting_power.power, Uint128::new(1u128)); + assert_eq!(addr1_voting_power.height, app.block_info().height); + + // Check total power for current block is now 3 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { height: None }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(3u128)); + assert_eq!(total_voting_power.height, app.block_info().height); + + // Check total power for last block is 4 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { + height: Some(app.block_info().height - 1), + }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(4u128)); + assert_eq!(total_voting_power.height, app.block_info().height - 1); + + // Remove address 2 completely + let msg = cw4_group::msg::ExecuteMsg::UpdateMembers { + remove: vec![ADDR2.to_string()], + add: vec![], + }; + + app.execute_contract( + Addr::unchecked(DAO_ADDR), + &ContractInfo { + address: cw4_contract_info.addr.clone(), + code_hash: cw4_contract_info.code_hash.clone(), + }, + &msg, + &[], + ) + .unwrap(); + app.update_block(next_block); + + // ADDR2 power is now 0 + let addr2_voting_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key_addr2.clone(), + address: ADDR2.to_string(), + }, + height: None, + }, + ) + .unwrap(); + assert_eq!(addr2_voting_power.power, Uint128::zero()); + assert_eq!(addr2_voting_power.height, app.block_info().height); + + // Check total power for current block is now 2 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { height: None }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(2u128)); + assert_eq!(total_voting_power.height, app.block_info().height); + + // Check total power for last block is 3 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { + height: Some(app.block_info().height - 1), + }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(3u128)); + assert_eq!(total_voting_power.height, app.block_info().height - 1); + + // Readd ADDR2 with 10 power + let msg = cw4_group::msg::ExecuteMsg::UpdateMembers { + remove: vec![], + add: vec![cw4::Member { + addr: ADDR2.to_string(), + weight: 10, + }], + }; + + app.execute_contract( + Addr::unchecked(DAO_ADDR), + &ContractInfo { + address: cw4_contract_info.addr.clone(), + code_hash: cw4_contract_info.code_hash.clone(), + }, + &msg, + &[], + ) + .unwrap(); + app.update_block(next_block); + + // ADDR2 power is now 10 + let addr2_voting_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key_addr2.clone(), + address: ADDR2.to_string(), + }, + height: None, + }, + ) + .unwrap(); + assert_eq!(addr2_voting_power.power, Uint128::new(10u128)); + assert_eq!(addr2_voting_power.height, app.block_info().height); + + // Check total power for current block is now 12 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::TotalPowerAtHeight { height: None }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(12u128)); + assert_eq!(total_voting_power.height, app.block_info().height); + + // Check total power for last block is 2 + let total_voting_power: TotalPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash, + voting_contract_info.address, + &QueryMsg::TotalPowerAtHeight { + height: Some(app.block_info().height - 1), + }, + ) + .unwrap(); + assert_eq!(total_voting_power.power, Uint128::new(2u128)); + assert_eq!(total_voting_power.height, app.block_info().height - 1); +} + +#[test] +fn test_migrate() { + let mut app = App::default(); + + let initial_members = vec![ + cw4::Member { + addr: ADDR1.to_string(), + weight: 1, + }, + cw4::Member { + addr: ADDR2.to_string(), + weight: 1, + }, + cw4::Member { + addr: ADDR3.to_string(), + weight: 1, + }, + ]; + + let query_auth = instantiate_query_auth(&mut app); + + let viewing_key = create_viewing_key( + &mut app, + ContractInfo { + address: query_auth.address.clone(), + code_hash: query_auth.code_hash.clone(), + }, + mock_info(ADDR1, &[]), + ); + + // Instantiate with no members, error + let voting_info = app.store_code(voting_contract()); + let cw4_info = app.store_code(cw4_contract()); + let msg = InstantiateMsg { + group_contract: GroupContract::New { + cw4_group_code_id: cw4_info.code_id, + cw4_group_code_hash: cw4_info.code_hash, + initial_members, + query_auth: Some(RawContract::new( + &query_auth.address.to_string(), + &query_auth.code_hash.to_string(), + )), + }, + dao_code_hash: "todo!()".into(), + }; + let voting_contract_info = app + .instantiate_contract( + voting_info.clone(), + Addr::unchecked(DAO_ADDR), + &msg, + &[], + "voting module", + Some(DAO_ADDR.to_string()), + ) + .unwrap(); + + let power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash.clone(), + voting_contract_info.address.clone(), + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key.clone(), + address: ADDR1.into(), + }, + height: None, + }, + ) + .unwrap(); + + app.execute( + Addr::unchecked(DAO_ADDR), + CosmosMsg::Wasm(WasmMsg::Migrate { + contract_addr: voting_contract_info.address.clone().to_string(), + code_id: voting_info.code_id, + code_hash: voting_contract_info.code_hash.clone(), + msg: to_binary(&MigrateMsg {}).unwrap(), + }), + ) + .unwrap(); + + let new_power: VotingPowerAtHeightResponse = app + .wrap() + .query_wasm_smart( + voting_contract_info.code_hash, + voting_contract_info.address, + &QueryMsg::VotingPowerAtHeight { + auth: Auth::ViewingKey { + key: viewing_key, + address: ADDR1.into(), + }, + height: None, + }, + ) + .unwrap(); + + assert_eq!(new_power, power) +} + +// #[test] +// fn test_duplicate_member() { +// let mut app = App::default(); +// let _voting_addr = setup_test_case(&mut app); +// let voting_info = app.store_code(voting_contract()); +// let cw4_info = app.store_code(cw4_contract()); +// // Instantiate with members but have a duplicate +// // Total weight is actually 69 but ADDR3 appears twice. // let msg = InstantiateMsg { -// group_contract: GroupContract::Existing { -// address: cw4_info_with_member.clone().address.to_string(), -// code_hash: cw4_info_with_member.clone().code_hash, +// group_contract: GroupContract::New { +// cw4_group_code_id: cw4_info, +// initial_members: vec![ +// cw4::Member { +// addr: ADDR3.to_string(), // same address above +// weight: 19, +// }, +// cw4::Member { +// addr: ADDR1.to_string(), +// weight: 25, +// }, +// cw4::Member { +// addr: ADDR2.to_string(), +// weight: 25, +// }, +// cw4::Member { +// addr: ADDR3.to_string(), +// weight: 19, +// }, +// ], // }, -// dao_code_hash: "dao_code_hash".to_string(), // }; -// let voting_info = app +// // Previous versions voting power was 100, due to no dedup. +// // Now we error +// // Bug busted : ) +// let _voting_addr = app // .instantiate_contract( -// voting_instantiate_info.clone(), +// voting_info, // Addr::unchecked(DAO_ADDR), // &msg, // &[], // "voting module", // None, // ) +// .unwrap_err(); +// } + +// #[test] +// fn test_zero_voting_power() { +// let mut app = App::default(); +// let voting_contract_info = setup_test_case(&mut app); +// app.update_block(next_block); + +// let cw4_contract_info: Addr = app +// .wrap() +// .query_wasm_smart(voting_contract_info.clone(), &QueryMsg::GroupContract {}) // .unwrap(); -// let info: InfoResponse = app +// // check that ADDR4 weight is 0 +// let addr4_voting_power: VotingPowerAtHeightResponse = app // .wrap() // .query_wasm_smart( -// voting_info.clone().code_hash, -// voting_info.clone().address.to_string(), -// &QueryMsg::Info {}, +// voting_contract_info.clone(), +// &QueryMsg::VotingPowerAtHeight { +// address: ADDR4.to_string(), +// height: None, +// }, // ) // .unwrap(); -// assert_eq!( -// info, -// InfoResponse { -// info: ContractVersion { -// contract: "crates.io:dao-voting-cw4".to_string(), -// version: env!("CARGO_PKG_VERSION").to_string() -// } -// } -// ); - -// // Ensure group contract is set -// let _group_contract: AnyContractInfo = app +// assert_eq!(addr4_voting_power.power, Uint128::new(0)); +// assert_eq!(addr4_voting_power.height, app.block_info().height); + +// // Update ADDR1's weight to 0 +// let msg = cw4_group::msg::ExecuteMsg::UpdateMembers { +// remove: vec![], +// add: vec![cw4::Member { +// addr: ADDR1.to_string(), +// weight: 0, +// }], +// }; +// app.execute_contract(Addr::unchecked(DAO_ADDR), cw4_contract_info, &msg, &[]) +// .unwrap(); + +// // Check ADDR1's power is now 0 +// let addr1_voting_power: VotingPowerAtHeightResponse = app // .wrap() // .query_wasm_smart( -// voting_info.clone().code_hash, -// voting_info.clone().address.to_string(), -// &QueryMsg::GroupContract {}, +// voting_contract_info.clone(), +// &QueryMsg::VotingPowerAtHeight { +// address: ADDR1.to_string(), +// height: None, +// }, // ) // .unwrap(); +// assert_eq!(addr1_voting_power.power, Uint128::new(0u128)); +// assert_eq!(addr1_voting_power.height, app.block_info().height); -// let dao_contract: AnyContractInfo = app +// // Check total power is now 2 +// let total_voting_power: TotalPowerAtHeightResponse = app // .wrap() -// .query_wasm_smart( -// voting_info.code_hash, -// voting_info.address.to_string(), -// &QueryMsg::Dao {}, -// ) +// .query_wasm_smart(voting_contract_info, &QueryMsg::TotalPowerAtHeight { height: None }) // .unwrap(); -// assert_eq!( -// dao_contract, -// AnyContractInfo { -// addr: Addr::unchecked(DAO_ADDR), -// code_hash: "dao_code_hash".to_string() -// } -// ); +// assert_eq!(total_voting_power.power, Uint128::new(2u128)); +// assert_eq!(total_voting_power.height, app.block_info().height); +// } + +// #[test] +// pub fn test_migrate_update_version() { +// let mut deps = mock_dependencies(); +// cw2::set_contract_version(&mut deps.storage, "my-contract", "1.0.0").unwrap(); +// migrate(deps.as_mut(), mock_env(), MigrateMsg {}).unwrap(); +// let version = cw2::get_contract_version(&deps.storage).unwrap(); +// assert_eq!(version.version, CONTRACT_VERSION); +// assert_eq!(version.contract, CONTRACT_NAME); // } diff --git a/contracts/voting/dao-voting-snip721-staked/src/state.rs b/contracts/voting/dao-voting-snip721-staked/src/state.rs index 108bdd4..a32f0d1 100644 --- a/contracts/voting/dao-voting-snip721-staked/src/state.rs +++ b/contracts/voting/dao-voting-snip721-staked/src/state.rs @@ -29,17 +29,17 @@ pub const INITIAL_NFTS: Item> = Item::new("initial_nfts"); /// The set of NFTs currently staked by each address. The existence of /// an `(address, token_id)` pair implies that `address` has staked /// `token_id`. -pub static STAKED_NFTS_PER_OWNER: Keymap<(Addr, String), Empty, Json> = Keymap::new(b"snpw"); +pub const STAKED_NFTS_PER_OWNER: Keymap<(Addr, String), Empty, Json> = Keymap::new(b"snpw"); -/// The number of NFTs staked by an address as a function of block -/// height. -pub static NFT_BALANCES_PRIMARY: Keymap = Keymap::new(b"nft_balances_primary"); -pub static NFT_BALANCES_SNAPSHOT: Keymap<(u64, Addr), Uint128, Json> = +/// The number of NFTs staked by an address as a function of block height. +pub const NFT_BALANCES_PRIMARY: Keymap = Keymap::new(b"nft_balances_primary"); +pub const NFT_BALANCES_SNAPSHOT: Keymap<(u64, Addr), Uint128> = Keymap::new(b"nft_balances_snapshot"); -pub static USER_STAKED_NFT_AT_HEIGHT: Keymap, Json> = - Keymap::new(b"user_Staked_Nft_at_height"); +pub const USER_STAKED_NFT_AT_HEIGHT: Keymap> = + Keymap::new(b"user_staked_nft_at_height"); pub struct NftBalancesStore {} + impl NftBalancesStore { // Function to store a value at a specific block height pub fn save( @@ -48,27 +48,31 @@ impl NftBalancesStore { key: Addr, value: Uint128, ) -> StdResult<()> { - let primary = NFT_BALANCES_PRIMARY.get(store, &key.clone()); + let primary = NFT_BALANCES_PRIMARY.get(store, &key); + if primary.is_none() { - NFT_BALANCES_PRIMARY.insert(store, &key.clone(), &value)?; + // If no previous balance, store the initial balance + NFT_BALANCES_PRIMARY.insert(store, &key, &value)?; NFT_BALANCES_SNAPSHOT.insert(store, &(block_height, key.clone()), &Uint128::zero())?; - USER_STAKED_NFT_AT_HEIGHT.insert(store, &key.clone(), &vec![block_height])?; + USER_STAKED_NFT_AT_HEIGHT.insert(store, &key, &vec![block_height])?; } else { - let mut user_staked_height = - USER_STAKED_NFT_AT_HEIGHT.get(store, &key.clone()).unwrap(); + // If a balance exists, create a snapshot and update the balance + let mut user_staked_height = USER_STAKED_NFT_AT_HEIGHT.get(store, &key).unwrap(); NFT_BALANCES_SNAPSHOT.insert(store, &(block_height, key.clone()), &primary.unwrap())?; - NFT_BALANCES_PRIMARY.insert(store, &key.clone(), &value)?; + NFT_BALANCES_PRIMARY.insert(store, &key, &value)?; user_staked_height.push(block_height); - USER_STAKED_NFT_AT_HEIGHT.insert(store, &key.clone(), &user_staked_height)?; + USER_STAKED_NFT_AT_HEIGHT.insert(store, &key, &user_staked_height)?; } Ok(()) } + // Load the current NFT balance for a specific address pub fn load(store: &dyn Storage, key: Addr) -> Uint128 { NFT_BALANCES_PRIMARY.get(store, &key).unwrap_or_default() } + // Load the NFT balance at a specific block height, falling back to the most recent balance if not found pub fn may_load_at_height( store: &dyn Storage, key: Addr, @@ -76,38 +80,47 @@ impl NftBalancesStore { ) -> StdResult> { let snapshot_key = (height, key.clone()); - let snapshot_value = NFT_BALANCES_SNAPSHOT.get(store, &snapshot_key); - if snapshot_value.is_none() { - Ok(NFT_BALANCES_PRIMARY.get(store, &key)) + // Try loading the snapshot balance at the given block height + if let Some(snapshot_value) = NFT_BALANCES_SNAPSHOT.get(store, &snapshot_key) { + Ok(Some(snapshot_value)) } else { - let x = USER_STAKED_NFT_AT_HEIGHT.get(store, &key).unwrap(); - let id = match x.binary_search(&height) { + // Fall back to the most recent snapshot before the height + let user_staked_heights = USER_STAKED_NFT_AT_HEIGHT + .get(store, &key) + .unwrap_or_default(); + let id = match user_staked_heights.binary_search(&height) { Ok(index) => Some(index), Err(_) => None, }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(NFT_BALANCES_PRIMARY.get(store, &key)) + + // If it's the last height, return the current primary balance + if let Some(index) = id { + if index == user_staked_heights.len() - 1 { + return Ok(NFT_BALANCES_PRIMARY.get(store, &key)); + } + // Otherwise, return the balance at the next height + let next_height = user_staked_heights[index + 1]; + Ok(NFT_BALANCES_SNAPSHOT.get(store, &(next_height, key.clone()))) } else { - let snapshot_value = - NFT_BALANCES_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize], key.clone())); - Ok(snapshot_value) + Ok(NFT_BALANCES_PRIMARY.get(store, &key)) } } } } -/// The number of NFTs staked with this contract as a function of -/// block height. +/// The number of NFTs staked with this contract as a function of block height. pub const TOTAL_STAKED_NFTS_PRIMARY: Item = Item::new("tsnP"); -pub static TOTAL_STAKED_NFTS_SNAPSHOT: Keymap = Keymap::new(b"tsns"); +pub const TOTAL_STAKED_NFTS_SNAPSHOT: Keymap = Keymap::new(b"tsns"); pub const TOTAL_STAKED_NFTS_AT_HEIGHTS: Item> = Item::new("tsnah"); pub struct StakedNftsTotalStore {} + impl StakedNftsTotalStore { // Function to store a value at a specific block height pub fn save(store: &mut dyn Storage, block_height: u64, value: Uint128) -> StdResult<()> { let primary = TOTAL_STAKED_NFTS_PRIMARY.load(store).unwrap_or_default(); + + // Insert or update the primary total staked NFTs if primary.is_zero() { TOTAL_STAKED_NFTS_PRIMARY.save(store, &value)?; TOTAL_STAKED_NFTS_SNAPSHOT.insert(store, &block_height, &Uint128::zero())?; @@ -123,29 +136,30 @@ impl StakedNftsTotalStore { Ok(()) } + // Load the current total staked NFTs pub fn load(store: &dyn Storage) -> Uint128 { TOTAL_STAKED_NFTS_PRIMARY.load(store).unwrap_or_default() } + // Load the total staked NFTs at a specific block height, falling back to the most recent total if not found pub fn may_load_at_height(store: &dyn Storage, height: u64) -> StdResult> { - let snapshot_key = height; - - let snapshot_value = TOTAL_STAKED_NFTS_SNAPSHOT.get(store, &snapshot_key); - if snapshot_value.is_none() { - Ok(Some(TOTAL_STAKED_NFTS_PRIMARY.load(store)?)) + // Check for a snapshot at the given block height + if let Some(snapshot_value) = TOTAL_STAKED_NFTS_SNAPSHOT.get(store, &height) { + Ok(Some(snapshot_value)) } else { - let x = TOTAL_STAKED_NFTS_AT_HEIGHTS.load(store).unwrap(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(Some(TOTAL_STAKED_NFTS_PRIMARY.load(store)?)) - } else { - let snapshot_value = - TOTAL_STAKED_NFTS_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize])); - Ok(snapshot_value) + // Fall back to the current primary total staked NFTs + let total_current = TOTAL_STAKED_NFTS_PRIMARY.load(store)?; + let heights = TOTAL_STAKED_NFTS_AT_HEIGHTS.load(store).unwrap_or_default(); + + // Determine the position in the heights array + match heights.binary_search(&height) { + Ok(index) if index == heights.len() - 1 => Ok(Some(total_current)), + Ok(index) => { + // Return the snapshot value at the next height + let next_height = heights[index + 1]; + Ok(TOTAL_STAKED_NFTS_SNAPSHOT.get(store, &next_height)) + } + Err(_) => Ok(Some(total_current)), // If height not found, return current total } } } diff --git a/contracts/voting/dao-voting-token-staked/src/state.rs b/contracts/voting/dao-voting-token-staked/src/state.rs index 215fd31..42a41ca 100644 --- a/contracts/voting/dao-voting-token-staked/src/state.rs +++ b/contracts/voting/dao-voting-token-staked/src/state.rs @@ -27,12 +27,13 @@ pub const DAO: Item = Item::new("dao"); pub const DENOM: Item = Item::new("denom"); /// Keeps track of staked balances by address over time -pub static STAKED_BALANCES_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); -pub static STAKED_BALANCES_SNAPSHOT: Keymap<(u64, Addr), Uint128> = +pub const STAKED_BALANCES_PRIMARY: Keymap = Keymap::new(b"staked_balances_primary"); +pub const STAKED_BALANCES_SNAPSHOT: Keymap<(u64, Addr), Uint128> = Keymap::new(b"staked_balances_snapshot"); -pub static USER_STAKED_AT_HEIGHT: Keymap> = Keymap::new(b"user_Staked_at_height"); +pub const USER_STAKED_AT_HEIGHT: Keymap> = Keymap::new(b"user_staked_at_height"); pub struct StakedBalancesStore {} + impl StakedBalancesStore { // Function to store a value at a specific block height pub fn save( @@ -41,58 +42,65 @@ impl StakedBalancesStore { key: Addr, value: Uint128, ) -> StdResult<()> { - let primary = STAKED_BALANCES_PRIMARY.get(store, &key.clone()); - if primary.is_none() { - STAKED_BALANCES_PRIMARY.insert(store, &key.clone(), &value)?; + // Get the current primary balance for the address + let primary_balance = STAKED_BALANCES_PRIMARY.get(store, &key).unwrap_or_default(); + + // If there's no primary balance, initialize it + if primary_balance.is_zero() { + STAKED_BALANCES_PRIMARY.insert(store, &key, &value)?; STAKED_BALANCES_SNAPSHOT.insert( store, &(block_height, key.clone()), &Uint128::zero(), )?; - USER_STAKED_AT_HEIGHT.insert(store, &key.clone(), &vec![block_height])?; + USER_STAKED_AT_HEIGHT.insert(store, &key, &vec![block_height])?; } else { - let mut user_staked_height = USER_STAKED_AT_HEIGHT.get(store, &key.clone()).unwrap(); + // Update the existing balance STAKED_BALANCES_SNAPSHOT.insert( store, &(block_height, key.clone()), - &primary.unwrap(), + &primary_balance, )?; - STAKED_BALANCES_PRIMARY.insert(store, &key.clone(), &value)?; + STAKED_BALANCES_PRIMARY.insert(store, &key, &value)?; + + // Update the list of staked heights for the user + let mut user_staked_height = USER_STAKED_AT_HEIGHT.get(store, &key).unwrap(); user_staked_height.push(block_height); - USER_STAKED_AT_HEIGHT.insert(store, &key.clone(), &user_staked_height)?; + USER_STAKED_AT_HEIGHT.insert(store, &key, &user_staked_height)?; } Ok(()) } + // Load the current staked balance for an address pub fn load(store: &dyn Storage, key: Addr) -> Uint128 { STAKED_BALANCES_PRIMARY.get(store, &key).unwrap_or_default() } + // Load the staked balance at a specific block height, falling back to the current balance if not found pub fn may_load_at_height( store: &dyn Storage, key: Addr, height: u64, ) -> StdResult> { + // Check for a snapshot at the given height let snapshot_key = (height, key.clone()); + if let Some(snapshot_value) = STAKED_BALANCES_SNAPSHOT.get(store, &snapshot_key) { + return Ok(Some(snapshot_value)); + } - let snapshot_value = STAKED_BALANCES_SNAPSHOT.get(store, &snapshot_key); - if snapshot_value.is_none() { - Ok(STAKED_BALANCES_PRIMARY.get(store, &key)) - } else { - let x = USER_STAKED_AT_HEIGHT.get(store, &key).unwrap(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(STAKED_BALANCES_PRIMARY.get(store, &key)) - } else { - let snapshot_value = - STAKED_BALANCES_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize], key.clone())); - Ok(snapshot_value) + // If no snapshot exists, check the primary balance + let primary_balance = STAKED_BALANCES_PRIMARY.get(store, &key).unwrap_or_default(); + let staked_heights = USER_STAKED_AT_HEIGHT.get(store, &key).unwrap_or_default(); + + // Check the index of the height in the staked heights + match staked_heights.binary_search(&height) { + Ok(index) if index == staked_heights.len() - 1 => Ok(Some(primary_balance)), // Last index + Ok(index) => { + let next_height = staked_heights[index + 1]; + Ok(STAKED_BALANCES_SNAPSHOT.get(store, &(next_height, key))) } + Err(_) => Ok(Some(primary_balance)), // Height not found, return current balance } } } @@ -100,51 +108,60 @@ impl StakedBalancesStore { /// Keeps track of staked total over time pub const STAKED_TOTAL_PRIMARY: Item = Item::new("staked_balances_primary"); pub static STAKED_TOTAL_SNAPSHOT: Keymap = Keymap::new(b"staked_balances_snapshot"); -pub const STAKED_TOTAL_AT_HEIGHTS: Item> = Item::new("user_Staked_at_height"); +pub const STAKED_TOTAL_AT_HEIGHTS: Item> = Item::new("user_staked_at_height"); pub struct TotalStakedStore {} + impl TotalStakedStore { // Function to store a value at a specific block height pub fn save(store: &mut dyn Storage, block_height: u64, value: Uint128) -> StdResult<()> { - let primary = STAKED_TOTAL_PRIMARY.load(store).unwrap_or_default(); - if primary.is_zero() { + // Load the current total staked amount + let primary_total = STAKED_TOTAL_PRIMARY.load(store).unwrap_or_default(); + + // If there is no staked total, initialize it + if primary_total.is_zero() { STAKED_TOTAL_PRIMARY.save(store, &value)?; STAKED_TOTAL_SNAPSHOT.insert(store, &block_height, &Uint128::zero())?; STAKED_TOTAL_AT_HEIGHTS.save(store, &vec![block_height])?; } else { - let mut user_staked_height = STAKED_TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); - STAKED_TOTAL_SNAPSHOT.insert(store, &block_height, &primary)?; + // Update existing totals + let mut user_staked_heights = STAKED_TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); + STAKED_TOTAL_SNAPSHOT.insert(store, &block_height, &primary_total)?; STAKED_TOTAL_PRIMARY.save(store, &value)?; - user_staked_height.push(block_height); - STAKED_TOTAL_AT_HEIGHTS.save(store, &user_staked_height)?; + user_staked_heights.push(block_height); + STAKED_TOTAL_AT_HEIGHTS.save(store, &user_staked_heights)?; } Ok(()) } + // Load the current total staked amount pub fn load(store: &dyn Storage) -> Uint128 { STAKED_TOTAL_PRIMARY.load(store).unwrap_or_default() } + // Load the staked total at a specific block height, falling back to the current total if not found pub fn may_load_at_height(store: &dyn Storage, height: u64) -> StdResult> { - let snapshot_key = height; + // Check for a snapshot at the given height + let snapshot_value = STAKED_TOTAL_SNAPSHOT.get(store, &height); - let snapshot_value = STAKED_TOTAL_SNAPSHOT.get(store, &snapshot_key); - if snapshot_value.is_none() { - Ok(Some(STAKED_TOTAL_PRIMARY.load(store).unwrap_or_default())) - } else { - let x = STAKED_TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); - let id = match x.binary_search(&height) { - Ok(index) => Some(index), - Err(_) => None, - }; - // return Ok(Some(Uint128::new(x.len() as u128))); - if id.unwrap() == (x.len() - 1) { - Ok(Some(STAKED_TOTAL_PRIMARY.load(store).unwrap_or_default())) - } else { - let snapshot_value = STAKED_TOTAL_SNAPSHOT.get(store, &(x[id.unwrap() + 1_usize])); - Ok(snapshot_value) + // If a snapshot exists, return it + if let Some(value) = snapshot_value { + return Ok(Some(value)); + } + + // If no snapshot exists, check the primary total + let primary_total = STAKED_TOTAL_PRIMARY.load(store).unwrap_or_default(); + let staked_heights = STAKED_TOTAL_AT_HEIGHTS.load(store).unwrap_or_default(); + + // Check the index of the height in the staked heights + match staked_heights.binary_search(&height) { + Ok(index) if index == staked_heights.len() - 1 => Ok(Some(primary_total)), // Last index + Ok(index) => { + let next_height = staked_heights[index + 1]; + Ok(STAKED_TOTAL_SNAPSHOT.get(store, &next_height)) } + Err(_) => Ok(Some(primary_total)), // Height not found, return current total } } }