From 83f456e92e56f25c0b25a95cb1b87f47505ea737 Mon Sep 17 00:00:00 2001 From: Simon Warta Date: Tue, 2 May 2023 17:45:33 +0200 Subject: [PATCH] Create and use namespace_with_key --- .../std/src/storage_keys/length_prefixed.rs | 82 +++++++++++++------ packages/std/src/storage_keys/mod.rs | 4 +- packages/storage/src/namespace_helpers.rs | 21 ++--- 3 files changed, 69 insertions(+), 38 deletions(-) diff --git a/packages/std/src/storage_keys/length_prefixed.rs b/packages/std/src/storage_keys/length_prefixed.rs index 721ea6ce59..9e95f17f04 100644 --- a/packages/std/src/storage_keys/length_prefixed.rs +++ b/packages/std/src/storage_keys/length_prefixed.rs @@ -6,44 +6,59 @@ /// Calculates the raw key prefix for a given namespace as documented /// in https://github.com/webmaster128/key-namespacing#length-prefixed-keys -pub fn to_length_prefixed(namespace: &[u8]) -> Vec { - let mut out = Vec::with_capacity(namespace.len() + 2); - out.extend_from_slice(&encode_length(namespace)); - out.extend_from_slice(namespace); +pub fn to_length_prefixed(namespace_component: &[u8]) -> Vec { + let mut out = Vec::with_capacity(namespace_component.len() + 2); + out.extend_from_slice(&encode_length(namespace_component)); + out.extend_from_slice(namespace_component); out } /// Calculates the raw key prefix for a given nested namespace /// as documented in https://github.com/webmaster128/key-namespacing#nesting -pub fn to_length_prefixed_nested(namespaces: &[&[u8]]) -> Vec { +pub fn to_length_prefixed_nested(namespace: &[&[u8]]) -> Vec { let mut size = 0; - for &namespace in namespaces { - size += namespace.len() + 2; + for component in namespace { + size += component.len() + 2; } let mut out = Vec::with_capacity(size); - for &namespace in namespaces { - out.extend_from_slice(&encode_length(namespace)); - out.extend_from_slice(namespace); + for component in namespace { + out.extend_from_slice(&encode_length(component)); + out.extend_from_slice(component); } out } -/// Encodes the length of a given namespace as a 2 byte big endian encoded integer -fn encode_length(namespace: &[u8]) -> [u8; 2] { - if namespace.len() > 0xFFFF { - panic!("only supports namespaces up to length 0xFFFF") +/// Encodes the length of a given namespace component +/// as a 2 byte big endian encoded integer +fn encode_length(namespace_component: &[u8]) -> [u8; 2] { + if namespace_component.len() > 0xFFFF { + panic!("only supports namespace components up to length 0xFFFF") } - let length_bytes = (namespace.len() as u32).to_be_bytes(); + let length_bytes = (namespace_component.len() as u32).to_be_bytes(); [length_bytes[2], length_bytes[3]] } -#[inline] -#[allow(unused)] -fn concat(namespace: &[u8], key: &[u8]) -> Vec { - let mut k = namespace.to_vec(); - k.extend_from_slice(key); - k +/// Encodes a namespace + key to a raw storage key. +/// +/// This is equivalent concat(to_length_prefixed_nested(namespace), key) +/// but more efficient when the namespace serialization is not persisted because +/// here we only need one vector allocation. +pub fn namespace_with_key(namespace: &[&[u8]], key: &[u8]) -> Vec { + // As documented in docs/STORAGE_KEYS.md, we know the final size of the key, + // which allows us to avoid reallocations of vectors. + let mut size = key.len(); + for component in namespace { + size += 2 /* encoded component length */ + component.len() /* the actual component data */; + } + + let mut out = Vec::with_capacity(size); + for component in namespace { + out.extend_from_slice(&encode_length(component)); + out.extend_from_slice(component); + } + out.extend_from_slice(key); + out } #[cfg(test)] @@ -77,7 +92,7 @@ mod tests { } #[test] - #[should_panic(expected = "only supports namespaces up to length 0xFFFF")] + #[should_panic(expected = "only supports namespace components up to length 0xFFFF")] fn to_length_prefixed_panics_for_too_long_prefix() { let limit = 0xFFFF; let long_namespace = vec![0; limit + 1]; @@ -186,8 +201,29 @@ mod tests { } #[test] - #[should_panic(expected = "only supports namespaces up to length 0xFFFF")] + #[should_panic(expected = "only supports namespace components up to length 0xFFFF")] fn encode_length_panics_for_large_values() { encode_length(&vec![1; 65536]); } + + #[test] + fn namespace_with_key_works() { + // Empty namespace + let enc = namespace_with_key(&[], b"foo"); + assert_eq!(enc, b"foo"); + let enc = namespace_with_key(&[], b""); + assert_eq!(enc, b""); + + // One component namespace + let enc = namespace_with_key(&[b"bar"], b"foo"); + assert_eq!(enc, b"\x00\x03barfoo"); + let enc = namespace_with_key(&[b"bar"], b""); + assert_eq!(enc, b"\x00\x03bar"); + + // Multi component namespace + let enc = namespace_with_key(&[b"bar", b"cool"], b"foo"); + assert_eq!(enc, b"\x00\x03bar\x00\x04coolfoo"); + let enc = namespace_with_key(&[b"bar", b"cool"], b""); + assert_eq!(enc, b"\x00\x03bar\x00\x04cool"); + } } diff --git a/packages/std/src/storage_keys/mod.rs b/packages/std/src/storage_keys/mod.rs index a9410c3f2a..630feaa18f 100644 --- a/packages/std/src/storage_keys/mod.rs +++ b/packages/std/src/storage_keys/mod.rs @@ -1,3 +1,5 @@ mod length_prefixed; -pub use length_prefixed::{to_length_prefixed, to_length_prefixed_nested}; +// Please note that the entire storage_keys module is public. So be careful +// when adding elements here. +pub use length_prefixed::{namespace_with_key, to_length_prefixed, to_length_prefixed_nested}; diff --git a/packages/storage/src/namespace_helpers.rs b/packages/storage/src/namespace_helpers.rs index 6a24829856..4b3aa9935c 100644 --- a/packages/storage/src/namespace_helpers.rs +++ b/packages/storage/src/namespace_helpers.rs @@ -1,4 +1,4 @@ -use cosmwasm_std::Storage; +use cosmwasm_std::{storage_keys::namespace_with_key, Storage}; #[cfg(feature = "iterator")] use cosmwasm_std::{Order, Record}; @@ -7,7 +7,7 @@ pub(crate) fn get_with_prefix( namespace: &[u8], key: &[u8], ) -> Option> { - storage.get(&concat(namespace, key)) + storage.get(&namespace_with_key(&[namespace], key)) } pub(crate) fn set_with_prefix( @@ -16,18 +16,11 @@ pub(crate) fn set_with_prefix( key: &[u8], value: &[u8], ) { - storage.set(&concat(namespace, key), value); + storage.set(&namespace_with_key(&[namespace], key), value); } pub(crate) fn remove_with_prefix(storage: &mut dyn Storage, namespace: &[u8], key: &[u8]) { - storage.remove(&concat(namespace, key)); -} - -#[inline] -fn concat(namespace: &[u8], key: &[u8]) -> Vec { - let mut k = namespace.to_vec(); - k.extend_from_slice(key); - k + storage.remove(&namespace_with_key(&[namespace], key)); } #[cfg(feature = "iterator")] @@ -40,11 +33,11 @@ pub(crate) fn range_with_prefix<'a>( ) -> Box + 'a> { // prepare start, end with prefix let start = match start { - Some(s) => concat(namespace, s), + Some(s) => namespace_with_key(&[namespace], s), None => namespace.to_vec(), }; let end = match end { - Some(e) => concat(namespace, e), + Some(e) => namespace_with_key(&[namespace], e), // end is updating last byte by one None => namespace_upper_bound(namespace), }; @@ -131,7 +124,7 @@ mod tests { // foo comes first let mut iter = storage.range(None, None, Order::Ascending); let first = iter.next().unwrap(); - let expected_key = concat(&prefix, b"bar"); + let expected_key = namespace_with_key(&[&prefix], b"bar"); assert_eq!(first, (expected_key, b"none".to_vec())); }