From 191f9c8ddb67479b5967392d0a8d668ef5c2fc18 Mon Sep 17 00:00:00 2001 From: Kayanski Date: Tue, 26 Nov 2024 14:31:09 +0000 Subject: [PATCH] Better account id storage serialization --- .../src/objects/account/account_id.rs | 105 +++++++++++++----- .../src/objects/account/account_trace.rs | 59 +++++----- .../src/objects/entry/asset_entry.rs | 2 +- .../src/objects/entry/channel_entry.rs | 2 +- .../src/objects/entry/contract_entry.rs | 2 +- .../src/objects/entry/dex_asset_pairing.rs | 2 +- .../abstract-std/src/objects/module.rs | 2 +- 7 files changed, 117 insertions(+), 57 deletions(-) diff --git a/framework/packages/abstract-std/src/objects/account/account_id.rs b/framework/packages/abstract-std/src/objects/account/account_id.rs index 7946a3f36..0611ff5f6 100644 --- a/framework/packages/abstract-std/src/objects/account/account_id.rs +++ b/framework/packages/abstract-std/src/objects/account/account_id.rs @@ -1,7 +1,8 @@ use std::{fmt::Display, str::FromStr}; -use cosmwasm_std::{StdError, StdResult}; +use cosmwasm_std::StdResult; use cw_storage_plus::{Key, KeyDeserialize, Prefixer, PrimaryKey}; +use deser::split_first_key; use super::{account_trace::AccountTrace, AccountSequence}; use crate::{objects::TruncatedChainId, AbstractError}; @@ -160,48 +161,72 @@ impl<'a> Prefixer<'a> for AccountId { impl KeyDeserialize for &AccountId { type Output = AccountId; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = AccountId::KEY_ELEMS; #[inline(always)] - fn from_vec(mut value: Vec) -> StdResult { - let mut tu = value.split_off(2); - let t_len = parse_length(&value)?; - let u = tu.split_off(t_len); + fn from_vec(value: Vec) -> StdResult { + let (trace, seq) = split_first_key(AccountTrace::KEY_ELEMS, value.as_ref())?; + + println!("{:x?} - {:?}", trace, seq); Ok(AccountId { - seq: AccountSequence::from_vec(u)?, - trace: AccountTrace::from_string(String::from_vec(tu)?), + seq: AccountSequence::from_vec(seq.to_vec())?, + trace: AccountTrace::from_vec(trace)?, }) } } impl KeyDeserialize for AccountId { type Output = AccountId; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = AccountTrace::KEY_ELEMS + u32::KEY_ELEMS; #[inline(always)] - fn from_vec(mut value: Vec) -> StdResult { - let mut tu = value.split_off(2); - let t_len = parse_length(&value)?; - let u = tu.split_off(t_len); - - Ok(AccountId { - seq: AccountSequence::from_vec(u)?, - trace: AccountTrace::from_string(String::from_vec(tu)?), - }) + fn from_vec(value: Vec) -> StdResult { + <&AccountId>::from_vec(value) } } -#[inline(always)] -fn parse_length(value: &[u8]) -> StdResult { - Ok(u16::from_be_bytes( - value - .try_into() - .map_err(|_| StdError::generic_err("Could not read 2 byte length"))?, - ) - .into()) -} +/// This was copied from cosmwasm-std +/// +/// https://github.com/CosmWasm/cw-storage-plus/blob/f65cd4000a0dc1c009f3f99e23f9e10a1c256a68/src/de.rs#L173 +pub(crate) mod deser { + use cosmwasm_std::{StdError, StdResult}; + + /// Splits the first key from the value based on the provided number of key elements. + /// The return value is ordered as (first_key, remainder). + /// + pub fn split_first_key(key_elems: u16, value: &[u8]) -> StdResult<(Vec, &[u8])> { + let mut index = 0; + let mut first_key = Vec::new(); + + // Iterate over the sub keys + for i in 0..key_elems { + let len_slice = &value[index..index + 2]; + index += 2; + let is_last_key = i == key_elems - 1; + + if !is_last_key { + first_key.extend_from_slice(len_slice); + } + + let subkey_len = parse_length(len_slice)?; + first_key.extend_from_slice(&value[index..index + subkey_len]); + index += subkey_len; + } + let remainder = &value[index..]; + Ok((first_key, remainder)) + } + + fn parse_length(value: &[u8]) -> StdResult { + Ok(u16::from_be_bytes( + value + .try_into() + .map_err(|_| StdError::generic_err("Could not read 2 byte length"))?, + ) + .into()) + } +} //-------------------------------------------------------------------------------------------------- // Tests //-------------------------------------------------------------------------------------------------- @@ -226,6 +251,13 @@ mod test { } } + fn mock_local_key() -> AccountId { + AccountId { + seq: 54, + trace: AccountTrace::Remote(vec![]), + } + } + fn mock_keys() -> (AccountId, AccountId, AccountId) { ( AccountId { @@ -268,6 +300,25 @@ mod test { assert_eq!(items[0], (key, 42069)); } + #[coverage_helper::test] + fn storage_key_local_works() { + let mut deps = mock_dependencies(); + let key = mock_local_key(); + let map: Map<&AccountId, u64> = Map::new("map"); + + map.save(deps.as_mut().storage, &key, &42069).unwrap(); + + assert_eq!(map.load(deps.as_ref().storage, &key).unwrap(), 42069); + + let items = map + .range(deps.as_ref().storage, None, None, Order::Ascending) + .map(|item| item.unwrap()) + .collect::>(); + + assert_eq!(items.len(), 1); + assert_eq!(items[0], (key, 42069)); + } + #[coverage_helper::test] fn composite_key_works() { let mut deps = mock_dependencies(); diff --git a/framework/packages/abstract-std/src/objects/account/account_trace.rs b/framework/packages/abstract-std/src/objects/account/account_trace.rs index 53becf9a8..f4574b3a6 100644 --- a/framework/packages/abstract-std/src/objects/account/account_trace.rs +++ b/framework/packages/abstract-std/src/objects/account/account_trace.rs @@ -1,11 +1,12 @@ use std::fmt::Display; +use super::account_id::deser::split_first_key; use cosmwasm_std::{ensure, Env, StdError, StdResult}; use cw_storage_plus::{Key, KeyDeserialize, Prefixer, PrimaryKey}; use crate::{constants::CHAIN_DELIMITER, objects::TruncatedChainId, AbstractError}; -pub const MAX_TRACE_LENGTH: usize = 6; +pub const MAX_TRACE_LENGTH: u16 = 6; pub(crate) const LOCAL: &str = "local"; /// The identifier of chain that triggered the account creation @@ -16,14 +17,30 @@ pub enum AccountTrace { Remote(Vec), } +pub const ACCOUNT_TRACE_KEY_PLACEHOLDER: &str = "place-holder-key"; + impl KeyDeserialize for &AccountTrace { type Output = AccountTrace; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = AccountTrace::KEY_ELEMS; #[inline(always)] fn from_vec(value: Vec) -> StdResult { - let value = value.into_iter().filter(|b| *b > 32).collect(); - Ok(AccountTrace::from_string(String::from_vec(value)?)) + let mut trace = vec![]; + // We parse the whole data for the MAX_TRACE_LENGTH keys + let mut value = value.as_ref(); + for i in 0..MAX_TRACE_LENGTH - 1 { + let (t, remainder) = split_first_key(1, value)?; + value = remainder; + let chain = String::from_utf8(t)?; + if i == 0 && chain == "local" { + return Ok(AccountTrace::Local); + } + if chain != ACCOUNT_TRACE_KEY_PLACEHOLDER { + trace.push(TruncatedChainId::from_string(chain).unwrap()) + } + } + + Ok(AccountTrace::Remote(trace)) } } @@ -34,35 +51,27 @@ impl<'a> PrimaryKey<'a> for AccountTrace { type SuperSuffix = Self; fn key(&self) -> Vec { - match self { + let mut serialization_result = match self { AccountTrace::Local => LOCAL.key(), - AccountTrace::Remote(chain_name) => { - let len = chain_name.len(); - chain_name - .iter() - .rev() - .enumerate() - .flat_map(|(s, c)| { - if s == len - 1 { - vec![c.str_ref().key()] - } else { - vec![c.str_ref().key(), CHAIN_DELIMITER.key()] - } - }) - .flatten() - .collect::>() - } + AccountTrace::Remote(chain_name) => chain_name + .iter() + .flat_map(|c| c.str_ref().key()) + .collect::>(), + }; + for _ in serialization_result.len()..(MAX_TRACE_LENGTH as usize) { + serialization_result.extend(ACCOUNT_TRACE_KEY_PLACEHOLDER.key()); } + serialization_result } } impl KeyDeserialize for AccountTrace { type Output = AccountTrace; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = MAX_TRACE_LENGTH; #[inline(always)] fn from_vec(value: Vec) -> StdResult { - Ok(AccountTrace::from_string(String::from_vec(value)?)) + <&AccountTrace>::from_vec(value) } } @@ -80,7 +89,7 @@ impl AccountTrace { AccountTrace::Remote(chain_trace) => { // Ensure the trace length is limited ensure!( - chain_trace.len() <= MAX_TRACE_LENGTH, + chain_trace.len() <= MAX_TRACE_LENGTH as usize, AbstractError::FormattingError { object: "chain-seq".into(), expected: format!("between 1 and {MAX_TRACE_LENGTH}"), @@ -370,7 +379,7 @@ mod test { .map(|item| item.unwrap()) .collect::>(); - assert_eq!(items.len(), 3); + assert_eq!(items.len(), 2); assert_eq!(items[0], (Addr::unchecked("jake"), 69420)); assert_eq!(items[1], (Addr::unchecked("larry"), 42069)); } diff --git a/framework/packages/abstract-std/src/objects/entry/asset_entry.rs b/framework/packages/abstract-std/src/objects/entry/asset_entry.rs index 513daee0d..ce538a43a 100644 --- a/framework/packages/abstract-std/src/objects/entry/asset_entry.rs +++ b/framework/packages/abstract-std/src/objects/entry/asset_entry.rs @@ -111,7 +111,7 @@ impl KeyDeserialize for AssetEntry { impl KeyDeserialize for &AssetEntry { type Output = AssetEntry; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = AssetEntry::KEY_ELEMS; #[inline(always)] fn from_vec(value: Vec) -> StdResult { diff --git a/framework/packages/abstract-std/src/objects/entry/channel_entry.rs b/framework/packages/abstract-std/src/objects/entry/channel_entry.rs index a75272b51..4a03fc1ff 100644 --- a/framework/packages/abstract-std/src/objects/entry/channel_entry.rs +++ b/framework/packages/abstract-std/src/objects/entry/channel_entry.rs @@ -83,7 +83,7 @@ impl<'a> Prefixer<'a> for &ChannelEntry { impl KeyDeserialize for &ChannelEntry { type Output = ChannelEntry; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = 2; #[inline(always)] fn from_vec(mut value: Vec) -> StdResult { diff --git a/framework/packages/abstract-std/src/objects/entry/contract_entry.rs b/framework/packages/abstract-std/src/objects/entry/contract_entry.rs index efd1723f6..c7ea692b8 100644 --- a/framework/packages/abstract-std/src/objects/entry/contract_entry.rs +++ b/framework/packages/abstract-std/src/objects/entry/contract_entry.rs @@ -107,7 +107,7 @@ impl<'a> Prefixer<'a> for &ContractEntry { impl KeyDeserialize for &ContractEntry { type Output = ContractEntry; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = 2; #[inline(always)] fn from_vec(mut value: Vec) -> StdResult { diff --git a/framework/packages/abstract-std/src/objects/entry/dex_asset_pairing.rs b/framework/packages/abstract-std/src/objects/entry/dex_asset_pairing.rs index 4f370bb1a..a3280572f 100644 --- a/framework/packages/abstract-std/src/objects/entry/dex_asset_pairing.rs +++ b/framework/packages/abstract-std/src/objects/entry/dex_asset_pairing.rs @@ -66,7 +66,7 @@ impl<'a> Prefixer<'a> for &DexAssetPairing { impl KeyDeserialize for &DexAssetPairing { type Output = DexAssetPairing; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = 3; #[inline(always)] fn from_vec(value: Vec) -> StdResult { diff --git a/framework/packages/abstract-std/src/objects/module.rs b/framework/packages/abstract-std/src/objects/module.rs index c5f22b813..1dea3f02d 100644 --- a/framework/packages/abstract-std/src/objects/module.rs +++ b/framework/packages/abstract-std/src/objects/module.rs @@ -175,7 +175,7 @@ impl<'a> Prefixer<'a> for &ModuleInfo { impl KeyDeserialize for &ModuleInfo { type Output = ModuleInfo; - const KEY_ELEMS: u16 = 1; + const KEY_ELEMS: u16 = Namespace::KEY_ELEMS + String::KEY_ELEMS + ModuleVersion::KEY_ELEMS; #[inline(always)] fn from_vec(mut value: Vec) -> StdResult {