soroban_env_host_zephyr/host/
metered_map.rs

1use crate::{
2    budget::{AsBudget, Budget},
3    host::{declared_size::DeclaredSizeForMetering, MeteredClone},
4    xdr::{ContractCostType, ScErrorCode, ScErrorType},
5    Compare, Error, Host, HostError,
6};
7
8use std::{borrow::Borrow, cmp::Ordering, marker::PhantomData};
9
10const MAP_OOB: Error = Error::from_type_and_code(ScErrorType::Object, ScErrorCode::IndexBounds);
11
12pub struct MeteredOrdMap<K, V, Ctx> {
13    pub(crate) map: Vec<(K, V)>,
14    ctx: PhantomData<Ctx>,
15}
16
17/// `Clone` should not be used directly, used `MeteredClone` instead if
18/// possible. `Clone` is defined here to satisfy trait requirements.
19impl<K, V, Ctx> Clone for MeteredOrdMap<K, V, Ctx>
20where
21    K: MeteredClone,
22    V: MeteredClone,
23    Ctx: AsBudget,
24{
25    fn clone(&self) -> Self {
26        Self {
27            map: self.map.clone(),
28            ctx: Default::default(),
29        }
30    }
31}
32
33impl<K, V, Ctx> std::hash::Hash for MeteredOrdMap<K, V, Ctx>
34where
35    K: std::hash::Hash,
36    V: std::hash::Hash,
37{
38    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
39        self.map.hash(state);
40    }
41}
42
43impl<K, V, Ctx> MeteredOrdMap<K, V, Ctx>
44where
45    K: DeclaredSizeForMetering,
46    V: DeclaredSizeForMetering,
47    Ctx: AsBudget,
48{
49    const ENTRY_SIZE: u64 = <(K, V) as DeclaredSizeForMetering>::DECLARED_SIZE;
50
51    fn charge_access<B: AsBudget>(&self, count: usize, b: &B) -> Result<(), HostError> {
52        b.as_budget().charge(
53            ContractCostType::MemCpy,
54            Some(Self::ENTRY_SIZE.saturating_mul(count as u64)),
55        )
56    }
57
58    fn charge_scan<B: AsBudget>(&self, b: &B) -> Result<(), HostError> {
59        Self::charge_access(self, self.map.len(), b)
60    }
61
62    // Charge binary search includes accessing number of entries expected for
63    // finding an entry. Cost of comparison is charged separately and not
64    // covered here.
65    fn charge_binsearch<B: AsBudget>(&self, b: &B) -> Result<(), HostError> {
66        let mag: u32 = 64u32.saturating_sub((self.map.len() as u64).leading_zeros());
67        b.as_budget().charge(
68            ContractCostType::MemCpy,
69            Some(Self::ENTRY_SIZE.saturating_mul(mag.saturating_add(1u32) as u64)),
70        )
71    }
72}
73
74impl<K, V, Ctx> Default for MeteredOrdMap<K, V, Ctx>
75where
76    Ctx: Default,
77{
78    fn default() -> Self {
79        Self {
80            map: Default::default(),
81            ctx: Default::default(),
82        }
83    }
84}
85
86// We abstract over Ctx:AsBudget here so that you can operate on MeteredOrdMap
87// before you have a Host constructed -- a bit, though only with certain types,
88// for example you can't do any lookups on maps keyed by Objects -- so long as
89// you at _least_ have a Budget. This is done to allow clients to populate and
90// reuse Storage maps keyed by non-Objects such as LedgerKey while only keeping
91// a Budget alive, rather than a whole Host.
92impl<K, V, Ctx> MeteredOrdMap<K, V, Ctx>
93where
94    K: MeteredClone,
95    V: MeteredClone,
96    Ctx: AsBudget + Compare<K, Error = HostError>,
97{
98    pub fn new() -> Self {
99        MeteredOrdMap {
100            map: Vec::new(),
101            ctx: Default::default(),
102        }
103    }
104
105    pub fn from_map(map: Vec<(K, V)>, ctx: &Ctx) -> Result<Self, HostError> {
106        if u32::try_from(map.len()).is_err() {
107            return Err(MAP_OOB.into());
108        }
109        // Allocation cost already paid for by caller, here just checks that input
110        // has sorted and unique keys.
111        let m = MeteredOrdMap {
112            map,
113            ctx: Default::default(),
114        };
115        m.charge_scan(ctx)?;
116        for w in m.map.as_slice().windows(2) {
117            let [a, b] = w else {
118                return Err((ScErrorType::Object, ScErrorCode::InternalError).into());
119            };
120            match <Ctx as Compare<K>>::compare(ctx, &a.0, &b.0)? {
121                Ordering::Less => (),
122                _ => return Err((ScErrorType::Object, ScErrorCode::InvalidInput).into()),
123            }
124        }
125        Ok(m)
126    }
127
128    // This doesn't take ExactSizeIterator since that is not implemented for Chain
129    // (see https://github.com/rust-lang/rust/issues/34433) but it only works
130    // with iterators that report an exact size_hint, and it constructs a new
131    // Vec from that iterator with a single allocation-and-copy.
132    pub fn from_exact_iter<I: Iterator<Item = (K, V)>>(
133        iter: I,
134        ctx: &Ctx,
135    ) -> Result<Self, HostError> {
136        let _span = tracy_span!("new map");
137        if let (_, Some(sz)) = iter.size_hint() {
138            if u32::try_from(sz).is_err() {
139                Err(MAP_OOB.into())
140            } else {
141                // It's possible we temporarily go over-budget here before charging, but
142                // only by the cost of temporarily allocating twice the size of our largest
143                // possible object. In exchange we get to batch all charges associated with
144                // the clone into one (when A::IS_SHALLOW==true).
145                let map: Vec<(K, V)> = iter.collect();
146                map.charge_deep_clone(ctx.as_budget())?;
147                // Delegate to from_map here to recheck sort order.
148                Self::from_map(map, ctx)
149            }
150        } else {
151            // This is a logic error, we should never get here.
152            Err((ScErrorType::Object, ScErrorCode::InternalError).into())
153        }
154    }
155
156    fn find<Q>(&self, key: &Q, ctx: &Ctx) -> Result<Result<usize, usize>, HostError>
157    where
158        K: Borrow<Q>,
159        Ctx: Compare<Q, Error = HostError>,
160    {
161        let _span = tracy_span!("map lookup");
162        self.charge_binsearch(ctx)?;
163        let mut err: Option<HostError> = None;
164        let res = self.map.binary_search_by(|probe| {
165            // We've already hit an error, return Ordering::Equal
166            // to terminate search asap.
167            if err.is_some() {
168                return Ordering::Equal;
169            }
170            match <Ctx as Compare<Q>>::compare(ctx, probe.0.borrow(), key) {
171                Ok(ord) => ord,
172                Err(he) => {
173                    err = Some(he);
174                    Ordering::Equal
175                }
176            }
177        });
178        match err {
179            Some(he) => Err(he),
180            None => Ok(res),
181        }
182    }
183
184    pub fn insert(&self, key: K, value: V, ctx: &Ctx) -> Result<Self, HostError> {
185        self.charge_access(1, ctx)?;
186        match self.find(&key, ctx)? {
187            Ok(replace_pos) => {
188                // [0,1,2] replace_pos == 1
189                // take(1) + new + skip(2)
190                // [0] + new + [2]
191                if replace_pos == usize::MAX - 1 {
192                    return Err(MAP_OOB.into());
193                }
194                let init = self.map.iter().take(replace_pos).cloned();
195                let fini = self.map.iter().skip(replace_pos.saturating_add(1)).cloned();
196                let iter = init.chain([(key, value)]).chain(fini);
197                Self::from_exact_iter(iter, ctx)
198            }
199            Err(insert_pos) => {
200                // [0,1,2] insert_pos == 1
201                // take(1) + new + skip(1)
202                // [0] new [1, 2]
203                if self.len() == u32::MAX as usize {
204                    return Err(MAP_OOB.into());
205                } else {
206                    let init = self.map.iter().take(insert_pos).cloned();
207                    let fini = self.map.iter().skip(insert_pos).cloned();
208                    let iter = init.chain([(key, value)]).chain(fini);
209                    Self::from_exact_iter(iter, ctx)
210                }
211            }
212        }
213    }
214
215    pub fn get<Q>(&self, key: &Q, ctx: &Ctx) -> Result<Option<&V>, HostError>
216    where
217        K: Borrow<Q>,
218        Ctx: Compare<Q, Error = HostError>,
219    {
220        match self.find(key, ctx)? {
221            Ok(found) => {
222                self.charge_access(1, ctx)?;
223                let Some((_, v)) = self.map.get(found) else {
224                    return Err((ScErrorType::Object, ScErrorCode::InternalError).into());
225                };
226                Ok(Some(v))
227            }
228            _ => Ok(None),
229        }
230    }
231
232    pub fn get_at_index(&self, index: usize, ctx: &Ctx) -> Result<&(K, V), HostError> {
233        self.charge_access(1, ctx)?;
234        self.map.get(index).ok_or_else(|| {
235            Error::from_type_and_code(ScErrorType::Object, ScErrorCode::IndexBounds).into()
236        })
237    }
238
239    /// Returns a `Some((new_self, val))` pair where `new_self` no longer
240    /// contains an entry for `key`, if the key existed, otherwise `None` if
241    /// `key` didn't exist (in which case there's no need to clone).
242    pub fn remove<Q>(&self, key: &Q, ctx: &Ctx) -> Result<Option<(Self, V)>, HostError>
243    where
244        K: Borrow<Q>,
245        Ctx: Compare<Q, Error = HostError>,
246    {
247        match self.find(key, ctx)? {
248            Ok(found) if found > 0 => {
249                // There's a nonempty prefix to preserve.
250                // [0,1,2] remove_pos == 1
251                // take(1) + skip(2)
252                // [0] [2]
253                // `found` cannot be > `usize::MAX` - 1, since that means the map contains more than
254                // `usize::MAX` elements. Therefore `found + 1` is guaranteed to not overflow.
255                let init = self.map.iter().take(found).cloned();
256                let fini = self.map.iter().skip(found.saturating_add(1)).cloned();
257                let iter = init.chain(fini);
258                let new = Self::from_exact_iter(iter, ctx)?;
259                let Some((_, res)) = self.map.get(found) else {
260                    return Err((ScErrorType::Object, ScErrorCode::InternalError).into());
261                };
262                Ok(Some((new, res.metered_clone(ctx.as_budget())?)))
263            }
264            Ok(found) => {
265                // No prefix, removing at position 0.
266                // If the suffix is empty it's harmless.
267                let iter = self.map.iter().skip(1).cloned();
268                let new = Self::from_exact_iter(iter, ctx)?;
269                let Some((_, res)) = self.map.get(found) else {
270                    return Err((ScErrorType::Object, ScErrorCode::InternalError).into());
271                };
272                Ok(Some((new, res.metered_clone(ctx.as_budget())?)))
273            }
274            _ => Ok(None),
275        }
276    }
277
278    pub fn len(&self) -> usize {
279        self.map.len()
280    }
281
282    pub fn contains_key<Q>(&self, key: &Q, ctx: &Ctx) -> Result<bool, HostError>
283    where
284        K: Borrow<Q>,
285        Ctx: Compare<Q, Error = HostError>,
286    {
287        Ok(self.find(key, ctx)?.is_ok())
288    }
289
290    pub fn keys(&self, ctx: &Ctx) -> Result<impl Iterator<Item = &K>, HostError> {
291        self.charge_scan(ctx)?;
292        Ok(self.map.iter().map(|(k, _)| k))
293    }
294
295    pub fn values(&self, ctx: &Ctx) -> Result<impl Iterator<Item = &V>, HostError> {
296        self.charge_scan(ctx)?;
297        Ok(self.map.iter().map(|(_, v)| v))
298    }
299
300    pub fn iter(&self, ctx: &Ctx) -> Result<impl Iterator<Item = &(K, V)>, HostError> {
301        self.charge_scan(ctx)?;
302        Ok(self.map.iter())
303    }
304}
305
306impl<K, V, Ctx> DeclaredSizeForMetering for MeteredOrdMap<K, V, Ctx>
307where
308    K: DeclaredSizeForMetering,
309    V: DeclaredSizeForMetering,
310{
311    const DECLARED_SIZE: u64 = <Vec<(K, V)> as DeclaredSizeForMetering>::DECLARED_SIZE;
312}
313
314impl<K, V, Ctx> MeteredClone for MeteredOrdMap<K, V, Ctx>
315where
316    K: MeteredClone,
317    V: MeteredClone,
318    Ctx: AsBudget,
319{
320    fn charge_for_substructure(&self, budget: impl AsBudget) -> Result<(), HostError> {
321        self.map.charge_for_substructure(budget)
322    }
323}
324
325impl<K, V> Compare<MeteredOrdMap<K, V, Host>> for Host
326where
327    K: DeclaredSizeForMetering,
328    V: DeclaredSizeForMetering,
329    Host: Compare<K, Error = HostError> + Compare<V, Error = HostError>,
330{
331    type Error = HostError;
332
333    fn compare(
334        &self,
335        a: &MeteredOrdMap<K, V, Host>,
336        b: &MeteredOrdMap<K, V, Host>,
337    ) -> Result<Ordering, Self::Error> {
338        // Here covers the cost of accessing number of map entries. The cost of
339        // comparing entries is covered by the `compare` call below.
340        self.as_budget().charge(
341            ContractCostType::MemCpy,
342            Some(
343                <(K, V) as DeclaredSizeForMetering>::DECLARED_SIZE
344                    .saturating_mul(a.map.len().min(b.map.len()) as u64),
345            ),
346        )?;
347        <Self as Compare<Vec<(K, V)>>>::compare(self, &a.map, &b.map)
348    }
349}
350
351impl<K, V> Compare<MeteredOrdMap<K, V, Budget>> for Budget
352where
353    K: DeclaredSizeForMetering,
354    V: DeclaredSizeForMetering,
355    Budget: Compare<K, Error = HostError> + Compare<V, Error = HostError>,
356{
357    type Error = HostError;
358
359    fn compare(
360        &self,
361        a: &MeteredOrdMap<K, V, Budget>,
362        b: &MeteredOrdMap<K, V, Budget>,
363    ) -> Result<Ordering, Self::Error> {
364        // Here covers the cost of accessing number of map entries. The cost of
365        // comparing entries is covered by the `compare` call below.
366        self.charge(
367            ContractCostType::MemCpy,
368            Some(
369                <(K, V) as DeclaredSizeForMetering>::DECLARED_SIZE
370                    .saturating_mul(a.map.len().min(b.map.len()) as u64),
371            ),
372        )?;
373        <Self as Compare<Vec<(K, V)>>>::compare(self, &a.map, &b.map)
374    }
375}
376
377impl<'a, K, V, Ctx> IntoIterator for &'a MeteredOrdMap<K, V, Ctx> {
378    type Item = &'a (K, V);
379    type IntoIter = core::slice::Iter<'a, (K, V)>;
380
381    fn into_iter(self) -> Self::IntoIter {
382        self.map.iter()
383    }
384}
385
386impl<K, V, Ctx> IntoIterator for MeteredOrdMap<K, V, Ctx> {
387    type Item = (K, V);
388    type IntoIter = std::vec::IntoIter<(K, V)>;
389
390    fn into_iter(self) -> Self::IntoIter {
391        self.map.into_iter()
392    }
393}
soroban_env_host_zephyr/host/metered_map.rs

soroban_env_host_zephyr/host/
metered_map.rs
