oasis_contract_sdk_storage/
map.rs

1//! A map backed by contract storage.
2use std::{convert::TryInto, marker::PhantomData};
3
4use oasis_contract_sdk::{
5    storage::{ConfidentialStore, PublicStore},
6    types::address::Address,
7};
8
9use crate::cell::{ConfidentialCell, PublicCell};
10
11macro_rules! declare_map {
12    ($name:ident, $cell:ident, $store:ident) => {
13        /// A map backed by contract storage.
14        pub struct $name<'key, K, V> {
15            /// Unique map identifier.
16            key: &'key [u8],
17
18            _key: PhantomData<K>,
19            _value: PhantomData<V>,
20        }
21
22        impl<'key, K, V> $name<'key, K, V> {
23            /// Create a new map instance.
24            pub const fn new(key: &'key [u8]) -> Self {
25                Self {
26                    key,
27                    _key: PhantomData,
28                    _value: PhantomData,
29                }
30            }
31        }
32
33        impl<'key, K, V> $name<'key, K, V>
34        where
35            K: MapKey,
36            V: cbor::Encode + cbor::Decode,
37        {
38            fn key(&self, key: K) -> Vec<u8> {
39                let raw_key = key.key();
40                encode_length_prefixed_path(
41                    self.key,
42                    &raw_key[..raw_key.len() - 1],
43                    raw_key[raw_key.len() - 1],
44                )
45            }
46
47            /// Lookup a given key.
48            pub fn get(&self, store: &dyn $store, key: K) -> Option<V> {
49                $cell::new(&self.key(key)).get(store)
50            }
51
52            /// Insert a given key/value pair.
53            pub fn insert(&self, store: &mut dyn $store, key: K, value: V) {
54                $cell::new(&self.key(key)).set(store, value);
55            }
56
57            /// Remove a given key.
58            pub fn remove(&self, store: &mut dyn $store, key: K) {
59                $cell::<V>::new(&self.key(key)).clear(store);
60            }
61        }
62    };
63}
64
65declare_map!(PublicMap, PublicCell, PublicStore);
66declare_map!(ConfidentialMap, ConfidentialCell, ConfidentialStore);
67
68/// A trait for types which can be used as map keys.
69pub trait MapKey {
70    /// Return the composite key.
71    fn key(&self) -> Vec<&[u8]>;
72}
73
74impl<const N: usize> MapKey for [u8; N] {
75    fn key(&self) -> Vec<&[u8]> {
76        vec![self]
77    }
78}
79
80impl MapKey for &[u8] {
81    fn key(&self) -> Vec<&[u8]> {
82        vec![self]
83    }
84}
85
86impl MapKey for &str {
87    fn key(&self) -> Vec<&[u8]> {
88        vec![self.as_bytes()]
89    }
90}
91
92impl MapKey for Vec<u8> {
93    fn key(&self) -> Vec<&[u8]> {
94        vec![self]
95    }
96}
97
98impl MapKey for String {
99    fn key(&self) -> Vec<&[u8]> {
100        vec![self.as_bytes()]
101    }
102}
103
104impl<T, U> MapKey for (T, U)
105where
106    T: MapKey,
107    U: MapKey,
108{
109    fn key(&self) -> Vec<&[u8]> {
110        let mut key = self.0.key();
111        key.extend(self.1.key());
112        key
113    }
114}
115
116impl<T, U, V> MapKey for (T, U, V)
117where
118    T: MapKey,
119    U: MapKey,
120    V: MapKey,
121{
122    fn key(&self) -> Vec<&[u8]> {
123        let mut key = self.0.key();
124        key.extend(self.1.key());
125        key.extend(self.2.key());
126        key
127    }
128}
129
130impl MapKey for Address {
131    fn key(&self) -> Vec<&[u8]> {
132        vec![self.as_ref()]
133    }
134}
135
136/// A trait representing an integer that can be encoded into big-endian bytes.
137pub trait Integer {
138    /// Type of the encoded representation.
139    type Encoded: AsRef<[u8]>;
140
141    /// Return the memory representation of this integer as a byte array in big-endian byte order.
142    fn to_be_bytes(self) -> Self::Encoded;
143}
144
145macro_rules! impl_integer_for_primitive {
146    ($ty:ty) => {
147        impl Integer for $ty {
148            type Encoded = [u8; std::mem::size_of::<$ty>()];
149
150            fn to_be_bytes(self) -> Self::Encoded {
151                <$ty>::to_be_bytes(self)
152            }
153        }
154    };
155}
156
157impl_integer_for_primitive!(u8);
158impl_integer_for_primitive!(u16);
159impl_integer_for_primitive!(u32);
160impl_integer_for_primitive!(u64);
161impl_integer_for_primitive!(u128);
162
163impl_integer_for_primitive!(i8);
164impl_integer_for_primitive!(i16);
165impl_integer_for_primitive!(i32);
166impl_integer_for_primitive!(i64);
167impl_integer_for_primitive!(i128);
168
169/// An integer in big-endian representation.
170pub struct Int<I: Integer> {
171    encoded: I::Encoded,
172    _type: PhantomData<I>,
173}
174
175impl<I: Integer> Int<I> {
176    /// Create a new integer in big-endian representation.
177    pub fn new(v: I) -> Self {
178        Self {
179            encoded: v.to_be_bytes(),
180            _type: PhantomData,
181        }
182    }
183}
184
185impl<I: Integer> From<I> for Int<I> {
186    fn from(v: I) -> Self {
187        Self::new(v)
188    }
189}
190
191impl<I: Integer> MapKey for Int<I> {
192    fn key(&self) -> Vec<&[u8]> {
193        vec![self.encoded.as_ref()]
194    }
195}
196
197/// Encodes the given components as a length-prefixed path.
198fn encode_length_prefixed_path(front: &[u8], middle: &[&[u8]], back: &[u8]) -> Vec<u8> {
199    let size = middle.iter().fold(0, |acc, k| acc + k.len() + 1);
200    let mut output = Vec::with_capacity(front.len() + size + back.len());
201
202    // Front is not length-prefixed.
203    output.extend_from_slice(front);
204    // Middle keys are length-prefixed.
205    for key in middle {
206        output.extend_from_slice(&encode_length(key));
207        output.extend_from_slice(key);
208    }
209    // Back is not length-prefixed.
210    output.extend_from_slice(back);
211
212    output
213}
214
215/// Encode the length of a storage key.
216///
217/// # Panics
218///
219/// This function will panic if the key length is greater than 255 bytes.
220fn encode_length(key: &[u8]) -> [u8; 1] {
221    [key.len().try_into().expect("key length greater than 255")]
222}
223
224#[cfg(test)]
225mod test {
226    use oasis_contract_sdk::testing::MockStore;
227
228    use super::*;
229
230    #[test]
231    fn test_map_basic() {
232        let mut store = MockStore::new();
233        let map: PublicMap<&str, u64> = PublicMap::new(b"test");
234
235        assert_eq!(map.get(&store, "foo"), None);
236        map.insert(&mut store, "foo", 42);
237        assert_eq!(map.get(&store, "foo"), Some(42));
238
239        let map: PublicMap<Int<u64>, String> = PublicMap::new(b"test2");
240
241        assert_eq!(map.get(&store, 42.into()), None);
242        map.insert(&mut store, 42.into(), "hello".to_string());
243        assert_eq!(map.get(&store, 42.into()), Some("hello".to_string()));
244
245        map.remove(&mut store, 42.into());
246        assert_eq!(map.get(&store, 42.into()), None);
247    }
248
249    #[test]
250    fn test_map_composite() {
251        let mut store = MockStore::new();
252        let map: PublicMap<(&str, &str), u64> = PublicMap::new(b"test");
253
254        assert_eq!(map.get(&store, ("foo", "bar")), None);
255        map.insert(&mut store, ("foo", "bar"), 42);
256        assert_eq!(map.get(&store, ("foo", "bar")), Some(42));
257        // Make sure we have proper key separation due to length-prefixing.
258        assert_eq!(map.get(&store, ("foob", "ar")), None);
259
260        map.remove(&mut store, ("foo", "bar"));
261        assert_eq!(map.get(&store, ("foo", "bar")), None);
262    }
263
264    #[test]
265    fn test_encode_length() {
266        assert_eq!(encode_length(b"foo"), [0x03]);
267    }
268
269    #[test]
270    #[should_panic]
271    fn test_encode_length_too_long() {
272        let v = vec![0x00; 260];
273        encode_length(&v);
274    }
275
276    #[test]
277    fn test_encode_length_prefixed_path() {
278        let four_five_six = vec![4, 5, 6];
279        let tcs = vec![
280            (vec![], vec![], vec![], vec![]),
281            (vec![1, 2, 3], vec![], vec![], vec![1, 2, 3]),
282            (vec![1, 2, 3], vec![], vec![4, 5, 6], vec![1, 2, 3, 4, 5, 6]),
283            (
284                vec![1, 2, 3],
285                vec![&four_five_six[..]],
286                vec![7, 8, 9],
287                vec![1, 2, 3, 3, 4, 5, 6, 7, 8, 9],
288            ),
289            (
290                vec![1, 2, 3],
291                vec![&four_five_six[..], &four_five_six[..]],
292                vec![7, 8, 9],
293                vec![1, 2, 3, 3, 4, 5, 6, 3, 4, 5, 6, 7, 8, 9],
294            ),
295        ];
296
297        for (front, middle, back, expected) in tcs {
298            assert_eq!(
299                encode_length_prefixed_path(&front, &middle, &back),
300                expected
301            );
302        }
303    }
304}
oasis_contract_sdk_storage/map.rs

oasis_contract_sdk_storage/
map.rs
