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

soroban_env_host/host/
metered_map.rs
