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