use cosmwasm_std::{storage_keys::namespace_with_key, Addr};
use crate::de::KeyDeserialize;
use crate::int_key::IntKey;
#[derive(Debug)]
pub enum Key<'a> {
Ref(&'a [u8]),
Val8([u8; 1]),
Val16([u8; 2]),
Val32([u8; 4]),
Val64([u8; 8]),
Val128([u8; 16]),
}
impl<'a> AsRef<[u8]> for Key<'a> {
fn as_ref(&self) -> &[u8] {
match self {
Key::Ref(r) => r,
Key::Val8(v) => v,
Key::Val16(v) => v,
Key::Val32(v) => v,
Key::Val64(v) => v,
Key::Val128(v) => v,
}
}
}
impl<'a> PartialEq<&[u8]> for Key<'a> {
fn eq(&self, other: &&[u8]) -> bool {
self.as_ref() == *other
}
}
pub trait PrimaryKey<'a>: Clone {
type Prefix: Prefixer<'a>;
type SubPrefix: Prefixer<'a>;
type Suffix: KeyDeserialize;
type SuperSuffix: KeyDeserialize;
fn key(&self) -> Vec<Key>;
fn joined_key(&self) -> Vec<u8> {
let keys = self.key();
let l = keys.len();
namespace_with_key(
&keys[0..l - 1].iter().map(Key::as_ref).collect::<Vec<_>>(),
keys[l - 1].as_ref(),
)
}
fn joined_extra_key(&self, key: &[u8]) -> Vec<u8> {
let keys = self.key();
namespace_with_key(&keys.iter().map(Key::as_ref).collect::<Vec<_>>(), key)
}
}
impl<'a> PrimaryKey<'a> for () {
type Prefix = Self;
type SubPrefix = Self;
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![]
}
}
impl<'a> PrimaryKey<'a> for &'a [u8] {
type Prefix = ();
type SubPrefix = ();
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![Key::Ref(self)]
}
}
impl<'a, const N: usize> PrimaryKey<'a> for [u8; N] {
type Prefix = ();
type SubPrefix = ();
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![Key::Ref(self.as_slice())]
}
}
impl<'a> PrimaryKey<'a> for &'a str {
type Prefix = ();
type SubPrefix = ();
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![Key::Ref(self.as_bytes())]
}
}
impl<'a, T: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize, U: PrimaryKey<'a> + KeyDeserialize>
PrimaryKey<'a> for (T, U)
{
type Prefix = T;
type SubPrefix = ();
type Suffix = U;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
let mut keys = self.0.key();
keys.extend(self.1.key());
keys
}
}
impl<'a, T> PrimaryKey<'a> for &'a T
where
T: PrimaryKey<'a>,
{
type Prefix = <T as PrimaryKey<'a>>::Prefix;
type SubPrefix = <T as PrimaryKey<'a>>::SubPrefix;
type Suffix = T::Suffix;
type SuperSuffix = T::SuperSuffix;
fn key(&self) -> Vec<Key> {
<T as PrimaryKey<'a>>::key(self)
}
}
impl<
'a,
T: PrimaryKey<'a> + Prefixer<'a>,
U: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
V: PrimaryKey<'a> + KeyDeserialize,
> PrimaryKey<'a> for (T, U, V)
{
type Prefix = (T, U);
type SubPrefix = T;
type Suffix = V;
type SuperSuffix = (U, V);
fn key(&self) -> Vec<Key> {
let mut keys = self.0.key();
keys.extend(self.1.key());
keys.extend(self.2.key());
keys
}
}
pub trait Prefixer<'a> {
fn prefix(&self) -> Vec<Key>;
fn joined_prefix(&self) -> Vec<u8> {
let prefixes = self.prefix();
namespace_with_key(&prefixes.iter().map(Key::as_ref).collect::<Vec<_>>(), &[])
}
}
impl<'a> Prefixer<'a> for () {
fn prefix(&self) -> Vec<Key> {
vec![]
}
}
impl<'a> Prefixer<'a> for &'a [u8] {
fn prefix(&self) -> Vec<Key> {
vec![Key::Ref(self)]
}
}
impl<'a, T: Prefixer<'a>, U: Prefixer<'a>> Prefixer<'a> for (T, U) {
fn prefix(&self) -> Vec<Key> {
let mut res = self.0.prefix();
res.extend(self.1.prefix());
res
}
}
impl<'a, T: Prefixer<'a>, U: Prefixer<'a>, V: Prefixer<'a>> Prefixer<'a> for (T, U, V) {
fn prefix(&self) -> Vec<Key> {
let mut res = self.0.prefix();
res.extend(self.1.prefix());
res.extend(self.2.prefix());
res
}
}
impl<'a, T> Prefixer<'a> for &'a T
where
T: Prefixer<'a>,
{
fn prefix(&self) -> Vec<Key> {
<T as Prefixer<'a>>::prefix(self)
}
}
impl<'a> Prefixer<'a> for &'a str {
fn prefix(&self) -> Vec<Key> {
vec![Key::Ref(self.as_bytes())]
}
}
impl<'a> PrimaryKey<'a> for Vec<u8> {
type Prefix = ();
type SubPrefix = ();
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![Key::Ref(self)]
}
}
impl<'a> Prefixer<'a> for Vec<u8> {
fn prefix(&self) -> Vec<Key> {
vec![Key::Ref(self.as_ref())]
}
}
impl<'a> PrimaryKey<'a> for String {
type Prefix = ();
type SubPrefix = ();
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![Key::Ref(self.as_bytes())]
}
}
impl<'a> Prefixer<'a> for String {
fn prefix(&self) -> Vec<Key> {
vec![Key::Ref(self.as_bytes())]
}
}
impl<'a> PrimaryKey<'a> for Addr {
type Prefix = ();
type SubPrefix = ();
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![Key::Ref(self.as_bytes())]
}
}
impl<'a> Prefixer<'a> for Addr {
fn prefix(&self) -> Vec<Key> {
vec![Key::Ref(self.as_bytes())]
}
}
macro_rules! integer_key {
(for $($t:ty, $v:tt),+) => {
$(impl<'a> PrimaryKey<'a> for $t {
type Prefix = ();
type SubPrefix = ();
type Suffix = Self;
type SuperSuffix = Self;
fn key(&self) -> Vec<Key> {
vec![Key::$v(self.to_cw_bytes())]
}
})*
}
}
integer_key!(for i8, Val8, u8, Val8, i16, Val16, u16, Val16, i32, Val32, u32, Val32, i64, Val64, u64, Val64, i128, Val128, u128, Val128);
macro_rules! integer_prefix {
(for $($t:ty, $v:tt),+) => {
$(impl<'a> Prefixer<'a> for $t {
fn prefix(&self) -> Vec<Key> {
vec![Key::$v(self.to_cw_bytes())]
}
})*
}
}
integer_prefix!(for i8, Val8, u8, Val8, i16, Val16, u16, Val16, i32, Val32, u32, Val32, i64, Val64, u64, Val64, i128, Val128, u128, Val128);
#[cfg(test)]
mod test {
use cosmwasm_std::{Uint256, Uint512};
use super::*;
#[test]
fn naked_8key_works() {
let k: u8 = 42u8;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(42u8.to_cw_bytes(), path[0].as_ref());
let k: i8 = 42i8;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(42i8.to_cw_bytes(), path[0].as_ref());
}
#[test]
fn naked_16key_works() {
let k: u16 = 4242u16;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242u16.to_cw_bytes(), path[0].as_ref());
let k: i16 = 4242i16;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242i16.to_cw_bytes(), path[0].as_ref());
}
#[test]
fn naked_32key_works() {
let k: u32 = 4242u32;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242u32.to_cw_bytes(), path[0].as_ref());
let k: i32 = 4242i32;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242i32.to_cw_bytes(), path[0].as_ref());
}
#[test]
fn naked_64key_works() {
let k: u64 = 4242u64;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242u64.to_cw_bytes(), path[0].as_ref());
let k: i64 = 4242i64;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242i64.to_cw_bytes(), path[0].as_ref());
}
#[test]
fn naked_128key_works() {
let k: u128 = 4242u128;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242u128.to_cw_bytes(), path[0].as_ref());
let k: i128 = 4242i128;
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(4242i128.to_cw_bytes(), path[0].as_ref());
}
#[test]
fn str_key_works() {
type K<'a> = &'a str;
let k: K = "hello";
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(b"hello", path[0].as_ref());
let joined = k.joined_key();
assert_eq!(joined, b"hello")
}
#[test]
fn string_key_works() {
type K = String;
let k: K = "hello".to_string();
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(b"hello", path[0].as_ref());
let joined = k.joined_key();
assert_eq!(joined, b"hello")
}
#[test]
fn fixed_size_bytes_key_works() {
type K = [u8; 32];
let k: K = Uint256::MAX.to_be_bytes();
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(k, path[0].as_ref());
let joined = k.joined_key();
assert_eq!(joined, k);
type K2<'a> = &'a [u8; 64];
let k: K2 = &Uint512::MAX.to_be_bytes();
let path = k.key();
assert_eq!(1, path.len());
assert_eq!(k, path[0].as_ref());
let joined = k.joined_key();
assert_eq!(joined, k);
}
#[test]
fn nested_str_key_works() {
type K<'a> = (&'a str, &'a [u8]);
let k: K = ("hello", b"world");
let path = k.key();
assert_eq!(2, path.len());
assert_eq!(b"hello", path[0].as_ref());
assert_eq!(b"world", path[1].as_ref());
}
#[test]
fn composite_byte_key() {
let k: (&[u8], &[u8]) = ("foo".as_bytes(), b"bar");
let path = k.key();
assert_eq!(2, path.len());
assert_eq!(path, vec!["foo".as_bytes(), b"bar"],);
}
#[test]
fn naked_composite_int_key() {
let k: (u32, u64) = (123, 87654);
let path = k.key();
assert_eq!(2, path.len());
assert_eq!(4, path[0].as_ref().len());
assert_eq!(8, path[1].as_ref().len());
assert_eq!(path[0].as_ref(), 123u32.to_cw_bytes());
assert_eq!(path[1].as_ref(), 87654u64.to_cw_bytes());
}
#[test]
fn nested_composite_keys() {
let first: &[u8] = b"foo";
let k: ((&[u8], &[u8]), &[u8]) = ((first, b"bar"), b"zoom");
let path = k.key();
assert_eq!(3, path.len());
assert_eq!(path, vec![first, b"bar", b"zoom"]);
let dir = k.0.prefix();
assert_eq!(2, dir.len());
assert_eq!(dir, vec![first, b"bar"]);
}
#[test]
fn naked_8bit_prefixes() {
let pair: (u8, &[u8]) = (123, b"random");
let one: Vec<u8> = vec![123];
let two: Vec<u8> = b"random".to_vec();
assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
let pair: (i8, &[u8]) = (123, b"random");
let one: Vec<u8> = vec![123 ^ 0x80];
let two: Vec<u8> = b"random".to_vec();
assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
}
#[test]
fn naked_16bit_prefixes() {
let pair: (u16, &[u8]) = (12345, b"random");
let one: Vec<u8> = vec![48, 57];
let two: Vec<u8> = b"random".to_vec();
assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
let pair: (i16, &[u8]) = (12345, b"random");
let one: Vec<u8> = vec![48 ^ 0x80, 57];
let two: Vec<u8> = b"random".to_vec();
assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
}
#[test]
fn naked_64bit_prefixes() {
let pair: (u64, &[u8]) = (12345, b"random");
let one: Vec<u8> = vec![0, 0, 0, 0, 0, 0, 48, 57];
let two: Vec<u8> = b"random".to_vec();
assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
let pair: (i64, &[u8]) = (12345, b"random");
#[allow(clippy::identity_op)]
let one: Vec<u8> = vec![0 ^ 0x80, 0, 0, 0, 0, 0, 48, 57];
let two: Vec<u8> = b"random".to_vec();
assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
}
#[test]
fn naked_proper_prefixes() {
let pair: (u32, &[u8]) = (12345, b"random");
let one: Vec<u8> = vec![0, 0, 48, 57];
let two: Vec<u8> = b"random".to_vec();
assert_eq!(pair.prefix(), vec![one.as_slice(), two.as_slice()]);
let triple: (&str, u32, &[u8]) = ("begin", 12345, b"end");
let one: Vec<u8> = b"begin".to_vec();
let two: Vec<u8> = vec![0, 0, 48, 57];
let three: Vec<u8> = b"end".to_vec();
assert_eq!(
triple.prefix(),
vec![one.as_slice(), two.as_slice(), three.as_slice()]
);
let owned_triple: (String, u32, Vec<u8>) = ("begin".to_string(), 12345, b"end".to_vec());
assert_eq!(
owned_triple.prefix(),
vec![one.as_slice(), two.as_slice(), three.as_slice()]
);
}
}