everscale_types/dict/
aug.rs

1use std::borrow::Borrow;
2use std::collections::BTreeMap;
3use std::marker::PhantomData;
4
5use super::{
6    aug_dict_find_by_extra, aug_dict_insert, aug_dict_remove_owned,
7    build_aug_dict_from_sorted_iter, SearchByExtra, SetMode,
8};
9use crate::cell::*;
10use crate::error::*;
11use crate::util::*;
12
13use super::raw::*;
14use super::typed::*;
15use super::{read_label, DictKey};
16
17/// A trait for values that can be used as augmented values in an augmented dictionary.
18pub trait AugDictExtra: Default {
19    /// Merges two augmented values.
20    ///
21    /// # Parameters
22    /// - `left` - The left branch (should start with `extra`).
23    /// - `right` - The right branch (should start with `extra`).
24    /// - `b` - The builder to store the result (only `extra`).
25    /// - `cx` - The cell context.
26    fn comp_add(
27        left: &mut CellSlice,
28        right: &mut CellSlice,
29        b: &mut CellBuilder,
30        cx: &mut dyn CellContext,
31    ) -> Result<(), Error>;
32}
33
34/// Typed augmented dictionary with fixed length keys.
35///
36/// # TLB scheme
37///
38/// ```text
39/// ahm_edge#_ {n:#} {V:Type} {A:Type} {l:#} {m:#}
40///   label:(HmLabel ~l n) {n = (~m) + l}
41///   node:(HashmapAugNode m V A) = HashmapAug n V A;
42///
43/// ahmn_leaf#_ {V:Type} {A:Type} extra:A value:V = HashmapAugNode 0 V A;
44/// ahmn_fork#_ {n:#} {V:Type} {A:Type} left:^(HashmapAug n V A)
45///   right:^(HashmapAug n V A) extra:A = HashmapAugNode (n + 1) V A;
46///
47/// ahme_empty$0 {n:#} {V:Type} {A:Type} extra:A = HashmapAugE n V A;
48/// ahme_root$1 {n:#} {V:Type} {A:Type} root:^(HashmapAug n V A) extra:A = HashmapAugE n V A;
49/// ```
50pub struct AugDict<K, A, V> {
51    dict: Dict<K, (A, V)>,
52    extra: A,
53    _key: PhantomData<K>,
54    _value: PhantomData<(A, V)>,
55}
56
57impl<K, A: ExactSize, V> ExactSize for AugDict<K, A, V> {
58    #[inline]
59    fn exact_size(&self) -> Size {
60        self.dict.exact_size() + self.extra.exact_size()
61    }
62}
63
64impl<'a, K, A: Load<'a>, V> Load<'a> for AugDict<K, A, V> {
65    #[inline]
66    fn load_from(slice: &mut CellSlice<'a>) -> Result<Self, Error> {
67        Ok(Self {
68            dict: ok!(Dict::load_from(slice)),
69            extra: ok!(A::load_from(slice)),
70            _key: PhantomData,
71            _value: PhantomData,
72        })
73    }
74}
75
76impl<K, A: Store, V> Store for AugDict<K, A, V> {
77    #[inline]
78    fn store_into(
79        &self,
80        builder: &mut CellBuilder,
81        context: &mut dyn CellContext,
82    ) -> Result<(), Error> {
83        ok!(self.dict.store_into(builder, context));
84        self.extra.store_into(builder, context)
85    }
86}
87
88impl<K, A: Default, V> Default for AugDict<K, A, V> {
89    #[inline]
90    fn default() -> Self {
91        Self::new()
92    }
93}
94
95impl<K, A: Clone, V> Clone for AugDict<K, A, V> {
96    fn clone(&self) -> Self {
97        Self {
98            dict: self.dict.clone(),
99            extra: self.extra.clone(),
100            _key: PhantomData,
101            _value: PhantomData,
102        }
103    }
104}
105
106impl<K, A: Eq, V> Eq for AugDict<K, A, V> {}
107
108impl<K, A: PartialEq, V> PartialEq for AugDict<K, A, V> {
109    fn eq(&self, other: &Self) -> bool {
110        self.dict.eq(&other.dict) && self.extra.eq(&other.extra)
111    }
112}
113
114impl<K, A: std::fmt::Debug, V> std::fmt::Debug for AugDict<K, A, V> {
115    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        debug_struct_field2_finish(f, "AugDict", "dict", &self.dict, "extra", &self.extra)
117    }
118}
119
120impl<K, A: Default, V> AugDict<K, A, V> {
121    /// Creates an empty dictionary
122    pub fn new() -> Self {
123        Self {
124            dict: Dict::new(),
125            extra: A::default(),
126            _key: PhantomData,
127            _value: PhantomData,
128        }
129    }
130
131    /// Manually constructs the dictionaty from parts.
132    pub const fn from_parts(dict: Dict<K, (A, V)>, extra: A) -> Self {
133        Self {
134            dict,
135            extra,
136            _key: PhantomData,
137            _value: PhantomData,
138        }
139    }
140
141    /// Returns an underlying dictionary and the extra value.
142    pub fn into_parts(self) -> (Dict<K, (A, V)>, A) {
143        (self.dict, self.extra)
144    }
145
146    /// Converts into a dictionary with an equivalent value type.
147    #[inline]
148    pub fn cast_into<Q, T>(self) -> AugDict<Q, A, T>
149    where
150        Q: EquivalentRepr<K>,
151        (A, T): EquivalentRepr<(A, V)>,
152    {
153        AugDict {
154            dict: self.dict.cast_into(),
155            extra: self.extra,
156            _key: PhantomData,
157            _value: PhantomData,
158        }
159    }
160}
161
162impl<K: DictKey, A, V> AugDict<K, A, V> {
163    #[allow(unused)]
164    pub(crate) fn load_from_root<'a>(
165        slice: &mut CellSlice<'a>,
166        context: &mut dyn CellContext,
167    ) -> Result<Self, Error>
168    where
169        A: Load<'a>,
170        V: Load<'a>,
171    {
172        let (extra, root) = ok!(load_from_root::<A, V>(slice, K::BITS, context));
173
174        Ok(Self {
175            dict: Dict::from(Some(root)),
176            extra,
177            _key: PhantomData,
178            _value: PhantomData,
179        })
180    }
181}
182
183impl<K, A, V> AugDict<K, A, V>
184where
185    K: DictKey,
186    for<'a> A: Default + Load<'a>,
187{
188    /// Recomputes the root extra value.
189    pub fn update_root_extra(&mut self) -> Result<(), Error> {
190        self.extra = match &self.dict.root {
191            Some(root) => {
192                let slice = &mut ok!(root.as_slice());
193                let prefix = ok!(read_label(slice, K::BITS));
194                if prefix.size_bits() != K::BITS {
195                    ok!(slice.skip_first(0, 2));
196                }
197                ok!(A::load_from(slice))
198            }
199            None => A::default(),
200        };
201        Ok(())
202    }
203}
204
205fn load_from_root<'a, A, V>(
206    slice: &mut CellSlice<'a>,
207    key_bit_len: u16,
208    context: &mut dyn CellContext,
209) -> Result<(A, Cell), Error>
210where
211    A: Load<'a>,
212    V: Load<'a>,
213{
214    let root = *slice;
215
216    let label = ok!(read_label(slice, key_bit_len));
217    let extra = if label.size_bits() != key_bit_len {
218        ok!(slice.skip_first(0, 2));
219        ok!(A::load_from(slice))
220    } else {
221        let extra = ok!(A::load_from(slice));
222        ok!(V::load_from(slice));
223        extra
224    };
225
226    let root_bits = root.size_bits() - slice.size_bits();
227    let root_refs = root.size_refs() - slice.size_refs();
228
229    let mut b = CellBuilder::new();
230    ok!(b.store_slice(root.get_prefix(root_bits, root_refs)));
231    match b.build_ext(context) {
232        Ok(cell) => Ok((extra, cell)),
233        Err(e) => Err(e),
234    }
235}
236
237impl<K, A, V> AugDict<K, A, V> {
238    /// Returns `true` if the dictionary contains no elements.
239    pub const fn is_empty(&self) -> bool {
240        self.dict.is_empty()
241    }
242
243    /// Returns the underlying dictionary.
244    #[inline]
245    pub const fn dict(&self) -> &Dict<K, (A, V)> {
246        &self.dict
247    }
248
249    /// Returns the root augmented value.
250    #[inline]
251    pub const fn root_extra(&self) -> &A {
252        &self.extra
253    }
254}
255
256impl<K, A, V> AugDict<K, A, V>
257where
258    K: Store + DictKey,
259{
260    /// Returns `true` if the dictionary contains a value for the specified key.
261    pub fn contains_key<Q>(&self, key: Q) -> Result<bool, Error>
262    where
263        Q: Borrow<K>,
264    {
265        self.dict.contains_key(key)
266    }
267}
268
269impl<K, A, V> AugDict<K, A, V>
270where
271    K: Store + DictKey,
272{
273    /// Returns the value corresponding to the key.
274    pub fn get<'a: 'b, 'b, Q>(&'a self, key: Q) -> Result<Option<(A, V)>, Error>
275    where
276        Q: Borrow<K> + 'b,
277        (A, V): Load<'a>,
278    {
279        self.dict.get(key)
280    }
281}
282
283impl<K, A, V> AugDict<K, A, V>
284where
285    K: DictKey,
286{
287    /// Searches for an item using a predicate on extra values.
288    ///
289    /// Used as a secondary index.
290    pub fn find_by_extra<'a, S>(&'a self, flow: S) -> Result<Option<(K, A, V)>, Error>
291    where
292        S: SearchByExtra<A>,
293        A: Load<'a>,
294        V: Load<'a>,
295    {
296        let Some((key, extra, mut value)) = ok!(aug_dict_find_by_extra::<A, S>(
297            self.dict.root.as_ref(),
298            K::BITS,
299            flow
300        )) else {
301            return Ok(None);
302        };
303
304        let Some(key) = K::from_raw_data(key.raw_data()) else {
305            return Err(Error::CellUnderflow);
306        };
307        let value = ok!(V::load_from(&mut value));
308        Ok(Some((key, extra, value)))
309    }
310}
311
312impl<K, A, V> AugDict<K, A, V>
313where
314    K: Store + DictKey,
315    for<'a> A: AugDictExtra + Store + Load<'a>,
316    V: Store,
317{
318    /// Builds a dictionary from a sorted collection.
319    pub fn try_from_btree<Q, E, T>(sorted: &BTreeMap<Q, (E, T)>) -> Result<Self, Error>
320    where
321        Q: Borrow<K>,
322        E: Borrow<A>,
323        T: Borrow<V>,
324        K: Ord,
325    {
326        let root = ok!(build_aug_dict_from_sorted_iter(
327            sorted
328                .iter()
329                .map(|(k, (a, v))| (k.borrow(), a.borrow(), v.borrow())),
330            K::BITS,
331            A::comp_add,
332            &mut Cell::empty_context()
333        ));
334
335        let mut result = Self {
336            dict: Dict::from_raw(root),
337            extra: A::default(),
338            _key: PhantomData,
339            _value: PhantomData,
340        };
341        ok!(result.update_root_extra());
342        Ok(result)
343    }
344
345    /// Builds a dictionary from a sorted slice.
346    pub fn try_from_sorted_slice<Q, E, T>(sorted: &[(Q, E, T)]) -> Result<Self, Error>
347    where
348        Q: Borrow<K>,
349        E: Borrow<A>,
350        T: Borrow<V>,
351        K: Ord,
352    {
353        let root = ok!(build_aug_dict_from_sorted_iter(
354            sorted
355                .iter()
356                .map(|(k, a, v)| (k.borrow(), a.borrow(), v.borrow())),
357            K::BITS,
358            A::comp_add,
359            &mut Cell::empty_context()
360        ));
361
362        let mut result = Self {
363            dict: Dict::from_raw(root),
364            extra: A::default(),
365            _key: PhantomData,
366            _value: PhantomData,
367        };
368        ok!(result.update_root_extra());
369        Ok(result)
370    }
371
372    /// Sets the augmented value associated with the key in the aug dictionary.
373    ///
374    /// Use [`set_ext`] if you need to use a custom cell context.
375    ///
376    /// [`set_ext`]: AugDict::set_ext
377    pub fn set<Q, E, T>(&mut self, key: Q, aug: E, value: T) -> Result<bool, Error>
378    where
379        Q: Borrow<K>,
380        E: Borrow<A>,
381        T: Borrow<V>,
382    {
383        self.set_ext(key, aug, value, &mut Cell::empty_context())
384    }
385
386    /// Sets the value associated with the key in the dictionary.
387    pub fn set_ext<Q, E, T>(
388        &mut self,
389        key: Q,
390        aug: E,
391        value: T,
392        context: &mut dyn CellContext,
393    ) -> Result<bool, Error>
394    where
395        Q: Borrow<K>,
396        E: Borrow<A>,
397        T: Borrow<V>,
398    {
399        self.insert_impl(
400            key.borrow(),
401            aug.borrow(),
402            value.borrow(),
403            SetMode::Set,
404            context,
405        )
406    }
407
408    /// Sets the augmented value associated with the key in the aug dictionary
409    /// only if the key was already present in it.
410    ///
411    /// Use [`replace_ext`] if you need to use a custom cell context.
412    ///
413    /// [`replace_ext`]: AugDict::replace_ext
414    pub fn replace<Q, E, T>(&mut self, key: Q, aug: E, value: T) -> Result<bool, Error>
415    where
416        Q: Borrow<K>,
417        E: Borrow<A>,
418        T: Borrow<V>,
419    {
420        self.replace_ext(key, aug, value, &mut Cell::empty_context())
421    }
422
423    /// Sets the value associated with the key in the dictionary
424    /// only if the key was already present in it.
425    pub fn replace_ext<Q, E, T>(
426        &mut self,
427        key: Q,
428        aug: E,
429        value: T,
430        context: &mut dyn CellContext,
431    ) -> Result<bool, Error>
432    where
433        Q: Borrow<K>,
434        E: Borrow<A>,
435        T: Borrow<V>,
436    {
437        self.insert_impl(
438            key.borrow(),
439            aug.borrow(),
440            value.borrow(),
441            SetMode::Replace,
442            context,
443        )
444    }
445
446    /// Sets the value associated with key in aug dictionary,
447    /// but only if it is not already present.
448    ///
449    /// Use [`add_ext`] if you need to use a custom cell context.
450    ///
451    /// [`add_ext`]: AugDict::add_ext
452    pub fn add<Q, E, T>(&mut self, key: Q, aug: E, value: T) -> Result<bool, Error>
453    where
454        Q: Borrow<K>,
455        E: Borrow<A>,
456        T: Borrow<V>,
457    {
458        self.add_ext(key, aug, value, &mut Cell::empty_context())
459    }
460
461    /// Sets the value associated with key in dictionary,
462    /// but only if it is not already present.
463    pub fn add_ext<Q, E, T>(
464        &mut self,
465        key: Q,
466        aug: E,
467        value: T,
468        context: &mut dyn CellContext,
469    ) -> Result<bool, Error>
470    where
471        Q: Borrow<K>,
472        E: Borrow<A>,
473        T: Borrow<V>,
474    {
475        self.insert_impl(
476            key.borrow(),
477            aug.borrow(),
478            value.borrow(),
479            SetMode::Add,
480            context,
481        )
482    }
483
484    /// Removes the value associated with key in aug dictionary.
485    /// Returns an optional removed value as cell slice parts.
486    pub fn remove<Q>(&mut self, key: Q) -> Result<Option<(A, V)>, Error>
487    where
488        Q: Borrow<K>,
489        for<'a> A: Load<'a> + 'static,
490        for<'a> V: Load<'a> + 'static,
491    {
492        match ok!(self.remove_raw_ext(key, &mut Cell::empty_context())) {
493            Some((cell, range)) => {
494                let mut slice = ok!(range.apply(&cell));
495                let extra = ok!(A::load_from(&mut slice));
496                let value = ok!(V::load_from(&mut slice));
497                Ok(Some((extra, value)))
498            }
499            None => Ok(None),
500        }
501    }
502
503    /// Removes the value associated with key in dictionary.
504    /// Returns an optional removed value as cell slice parts.
505    pub fn remove_raw_ext<Q>(
506        &mut self,
507        key: Q,
508        context: &mut dyn CellContext,
509    ) -> Result<Option<CellSliceParts>, Error>
510    where
511        Q: Borrow<K>,
512    {
513        self.remove_impl(key.borrow(), context)
514    }
515
516    fn insert_impl(
517        &mut self,
518        key: &K,
519        extra: &A,
520        value: &V,
521        mode: SetMode,
522        context: &mut dyn CellContext,
523    ) -> Result<bool, Error> {
524        let mut key_builder = CellBuilder::new();
525        ok!(key.store_into(&mut key_builder, &mut Cell::empty_context()));
526        let inserted = ok!(aug_dict_insert(
527            &mut self.dict.root,
528            &mut key_builder.as_data_slice(),
529            K::BITS,
530            extra,
531            value,
532            mode,
533            A::comp_add,
534            context,
535        ));
536
537        if inserted {
538            ok!(self.update_root_extra());
539        }
540
541        Ok(inserted)
542    }
543
544    fn remove_impl(
545        &mut self,
546        key: &K,
547        context: &mut dyn CellContext,
548    ) -> Result<Option<(Cell, CellSliceRange)>, Error> {
549        let mut key_builder = CellBuilder::new();
550        ok!(key.store_into(&mut key_builder, &mut Cell::empty_context()));
551        let res = ok!(aug_dict_remove_owned(
552            &mut self.dict.root,
553            &mut key_builder.as_data_slice(),
554            K::BITS,
555            false,
556            A::comp_add,
557            context,
558        ));
559
560        if res.is_some() {
561            ok!(self.update_root_extra());
562        }
563
564        Ok(res)
565    }
566
567    /// Split dictionary into 2 dictionaries by the first key bit.
568    pub fn split(&self) -> Result<(Self, Self), Error> {
569        self.split_by_prefix_ext(&Default::default(), &mut Cell::empty_context())
570    }
571
572    /// Split dictionary into 2 dictionaries by the first key bit.
573    pub fn split_ext(&self, context: &mut dyn CellContext) -> Result<(Self, Self), Error> {
574        self.split_by_prefix_ext(&Default::default(), context)
575    }
576
577    /// Split dictionary into 2 dictionaries at the prefix.
578    pub fn split_by_prefix(&self, key_prefix: &CellSlice<'_>) -> Result<(Self, Self), Error> {
579        self.split_by_prefix_ext(key_prefix, &mut Cell::empty_context())
580    }
581
582    /// Split dictionary into 2 dictionaries at the prefix.
583    pub fn split_by_prefix_ext(
584        &self,
585        key_prefix: &CellSlice<'_>,
586        context: &mut dyn CellContext,
587    ) -> Result<(Self, Self), Error> {
588        let (left, right) = ok!(self.dict.split_by_prefix_ext(key_prefix, context));
589
590        let mut left = Self {
591            dict: left,
592            extra: A::default(),
593            _key: PhantomData,
594            _value: PhantomData,
595        };
596        ok!(left.update_root_extra());
597
598        let mut right = Self {
599            dict: right,
600            extra: A::default(),
601            _key: PhantomData,
602            _value: PhantomData,
603        };
604        ok!(right.update_root_extra());
605
606        Ok((left, right))
607    }
608}
609
610impl<K, A, V> AugDict<K, A, V>
611where
612    K: DictKey,
613{
614    /// Gets an iterator over the entries of the dictionary, sorted by key.
615    /// The iterator element type is `Result<(K, A, V)>`.
616    ///
617    /// If the dictionary is invalid, finishes after the first invalid element,
618    /// returning an error.
619    ///
620    /// # Performance
621    ///
622    /// In the current implementation, iterating over dictionary builds a key
623    /// for each element. Use [`values`] or [`raw_values`] if you don't need keys from an iterator.
624    ///
625    /// [`values`]: Dict::values
626    /// [`raw_values`]: Dict::raw_values
627    pub fn iter<'a>(&'a self) -> AugIter<'a, K, A, V>
628    where
629        V: Load<'a>,
630    {
631        AugIter::new(self.dict.root())
632    }
633
634    /// Gets an iterator over the keys of the dictionary, in sorted order.
635    /// The iterator element type is `Result<K>`.
636    ///
637    /// If the dictionary is invalid, finishes after the first invalid element,
638    /// returning an error.
639    ///
640    /// # Performance
641    ///
642    /// In the current implementation, iterating over dictionary builds a key
643    /// for each element. Use [`values`] if you don't need keys from an iterator.
644    ///
645    /// [`values`]: Dict::values
646    pub fn keys(&'_ self) -> Keys<'_, K> {
647        Keys::new(self.dict.root())
648    }
649}
650
651impl<K, A, V> AugDict<K, A, V>
652where
653    K: DictKey,
654{
655    /// Gets an iterator over the augmented values of the dictionary, in order by key.
656    /// The iterator element type is `Result<V>`.
657    ///
658    /// If the dictionary is invalid, finishes after the first invalid element,
659    /// returning an error.
660    pub fn values<'a>(&'a self) -> Values<'a, (A, V)>
661    where
662        V: Load<'a>,
663    {
664        Values::new(self.dict.root(), K::BITS)
665    }
666}
667
668impl<K, A, V> AugDict<K, A, V>
669where
670    K: Store + DictKey,
671{
672    /// Gets an iterator over the raw entries of the dictionary, sorted by key.
673    /// The iterator element type is `Result<(CellBuilder, CellSlice)>`.
674    ///
675    /// If the dictionary is invalid, finishes after the first invalid element,
676    /// returning an error.
677    ///
678    /// # Performance
679    ///
680    /// In the current implementation, iterating over dictionary builds a key
681    /// for each element. Use [`values`] or [`raw_values`] if you don't need keys from an iterator.
682    ///
683    /// [`values`]: AugDict::values
684    /// [`raw_values`]: AugDict::raw_values
685    pub fn raw_iter(&'_ self) -> RawIter<'_> {
686        RawIter::new(self.dict.root(), K::BITS)
687    }
688
689    /// Gets an iterator over the raw keys of the dictionary, in sorted order.
690    /// The iterator element type is `Result<CellBuilder>`.
691    ///
692    /// If the dictionary is invalid, finishes after the first invalid element,
693    /// returning an error.
694    ///
695    /// # Performance
696    ///
697    /// In the current implementation, iterating over dictionary builds a key
698    /// for each element. Use [`values`] or [`raw_values`] if you don't need keys from an iterator.
699    ///
700    /// [`values`]: AugDict::values
701    /// [`raw_values`]: AugDict::raw_values
702    pub fn raw_keys(&'_ self) -> RawKeys<'_> {
703        RawKeys::new(self.dict.root(), K::BITS)
704    }
705}
706
707impl<K, A, V> AugDict<K, A, V>
708where
709    K: DictKey,
710{
711    /// Gets an iterator over the raw values of the dictionary, in order by key.
712    /// The iterator element type is `Result<CellSlice>`.
713    ///
714    /// If the dictionary is invalid, finishes after the first invalid element,
715    /// returning an error.
716    pub fn raw_values(&'_ self) -> RawValues<'_> {
717        RawValues::new(self.dict.root(), K::BITS)
718    }
719}
720
721#[cfg(feature = "serde")]
722impl<K, A, V> serde::Serialize for AugDict<K, A, V>
723where
724    K: serde::Serialize + Store + DictKey,
725    for<'a> A: serde::Serialize + Store + Load<'a>,
726    for<'a> V: serde::Serialize + Load<'a>,
727{
728    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
729    where
730        S: serde::Serializer,
731    {
732        use serde::ser::{Error, SerializeMap};
733
734        #[derive(serde::Serialize)]
735        struct AugDictHelper<'a, K, A, V>
736        where
737            K: serde::Serialize + Store + DictKey,
738            A: serde::Serialize + Store + Load<'a>,
739            V: serde::Serialize + Load<'a>,
740        {
741            #[serde(serialize_with = "serialize_dict_entries")]
742            entires: &'a AugDict<K, A, V>,
743            extra: &'a A,
744        }
745
746        fn serialize_dict_entries<'a, K, A, V, S>(
747            dict: &'a AugDict<K, A, V>,
748            serializer: S,
749        ) -> Result<S::Ok, S::Error>
750        where
751            S: serde::Serializer,
752            K: serde::Serialize + Store + DictKey,
753            A: serde::Serialize + Store + Load<'a>,
754            V: serde::Serialize + Load<'a>,
755        {
756            let mut ser = serializer.serialize_map(None)?;
757            for ref entry in dict.iter() {
758                let (key, extra, value) = match entry {
759                    Ok(entry) => entry,
760                    Err(e) => return Err(Error::custom(e)),
761                };
762                ok!(ser.serialize_entry(key, &(extra, value)));
763            }
764            ser.end()
765        }
766
767        if serializer.is_human_readable() {
768            AugDictHelper {
769                entires: self,
770                extra: &self.extra,
771            }
772            .serialize(serializer)
773        } else {
774            crate::boc::BocRepr::serialize(self, serializer)
775        }
776    }
777}
778
779/// An iterator over the entries of an [`AugDict`].
780///
781/// This struct is created by the [`iter`] method on [`AugDict`]. See its documentation for more.
782///
783/// [`iter`]: AugDict::iter
784pub struct AugIter<'a, K, A, V> {
785    inner: Iter<'a, K, (A, V)>,
786}
787
788impl<K, A, V> Clone for AugIter<'_, K, A, V> {
789    fn clone(&self) -> Self {
790        Self {
791            inner: self.inner.clone(),
792        }
793    }
794}
795
796impl<'a, K, A, V> AugIter<'a, K, A, V>
797where
798    K: DictKey,
799{
800    /// Creates an iterator over the entries of a dictionary.
801    pub fn new(root: &'a Option<Cell>) -> Self {
802        Self {
803            inner: Iter::new(root),
804        }
805    }
806
807    /// Changes the direction of the iterator to descending.
808    #[inline]
809    pub fn reversed(mut self) -> Self {
810        self.inner = self.inner.reversed();
811        self
812    }
813
814    /// Changes the behavior of the iterator to reverse the high bit.
815    #[inline]
816    pub fn signed(mut self) -> Self {
817        self.inner = self.inner.signed();
818        self
819    }
820}
821
822impl<'a, K, A, V> Iterator for AugIter<'a, K, A, V>
823where
824    K: DictKey,
825    (A, V): Load<'a>,
826{
827    type Item = Result<(K, A, V), Error>;
828
829    fn next(&mut self) -> Option<Self::Item> {
830        match self.inner.next()? {
831            Ok((key, (aug, value))) => Some(Ok((key, aug, value))),
832            Err(e) => Some(Err(e)),
833        }
834    }
835}
836
837#[cfg(test)]
838mod tests {
839    use anyhow::Context;
840
841    use super::*;
842
843    #[derive(Debug, Default, Load, Store, Eq, PartialEq)]
844    struct OrCmp(bool);
845
846    impl AugDictExtra for OrCmp {
847        fn comp_add(
848            left: &mut CellSlice,
849            right: &mut CellSlice,
850            b: &mut CellBuilder,
851            _: &mut dyn CellContext,
852        ) -> Result<(), Error> {
853            let left = left.load_bit()?;
854            let right = right.load_bit()?;
855            b.store_bit(left | right)
856        }
857    }
858
859    #[derive(Debug, Default, Load, Store, Eq, PartialEq)]
860    struct SomeValue(u32);
861
862    impl AugDictExtra for SomeValue {
863        fn comp_add(
864            left: &mut CellSlice,
865            right: &mut CellSlice,
866            b: &mut CellBuilder,
867            _: &mut dyn CellContext,
868        ) -> Result<(), Error> {
869            let left = left.load_u32()?;
870            let right = right.load_u32()?;
871            b.store_u32(left.saturating_add(right))
872        }
873    }
874
875    impl rand::distributions::Distribution<SomeValue> for rand::distributions::Standard {
876        #[inline]
877        fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> SomeValue {
878            SomeValue(rand::distributions::Standard.sample(rng))
879        }
880    }
881
882    #[test]
883    fn dict_set() {
884        let mut dict = AugDict::<u32, OrCmp, u16>::new();
885        assert_eq!(*dict.root_extra(), OrCmp(false));
886
887        dict.set(123, OrCmp(false), 0xffff).unwrap();
888        assert_eq!(dict.get(123).unwrap(), Some((OrCmp(false), 0xffff)));
889        assert_eq!(*dict.root_extra(), OrCmp(false));
890
891        dict.set(123, OrCmp(true), 0xcafe).unwrap();
892        assert_eq!(dict.get(123).unwrap(), Some((OrCmp(true), 0xcafe)));
893        assert_eq!(*dict.root_extra(), OrCmp(true));
894    }
895
896    #[test]
897    fn dict_set_complex() {
898        let mut dict = AugDict::<u32, OrCmp, u32>::new();
899        assert_eq!(*dict.root_extra(), OrCmp(false));
900
901        for i in 0..520 {
902            dict.set(i, OrCmp(true), 123).unwrap();
903        }
904        assert_eq!(*dict.root_extra(), OrCmp(true));
905    }
906
907    #[test]
908    fn dict_replace() {
909        let mut dict = AugDict::<u32, OrCmp, u16>::new();
910        assert_eq!(*dict.root_extra(), OrCmp(false));
911        dict.replace(123, OrCmp(false), 0xff).unwrap();
912        assert!(!dict.contains_key(123).unwrap());
913        assert_eq!(*dict.root_extra(), OrCmp(false));
914
915        dict.set(123, OrCmp(false), 0xff).unwrap();
916        assert_eq!(dict.get(123).unwrap(), Some((OrCmp(false), 0xff)));
917        assert_eq!(*dict.root_extra(), OrCmp(false));
918
919        dict.replace(123, OrCmp(true), 0xaa).unwrap();
920        assert_eq!(dict.get(123).unwrap(), Some((OrCmp(true), 0xaa)));
921        assert_eq!(*dict.root_extra(), OrCmp(true));
922    }
923
924    #[test]
925    fn dict_add() {
926        let mut dict = AugDict::<u32, OrCmp, u16>::new();
927        assert_eq!(*dict.root_extra(), OrCmp(false));
928
929        dict.add(123, OrCmp(false), 0x12).unwrap();
930        assert_eq!(dict.get(123).unwrap(), Some((OrCmp(false), 0x12)));
931        assert_eq!(*dict.root_extra(), OrCmp(false));
932
933        dict.add(123, OrCmp(true), 0x11).unwrap();
934        assert_eq!(dict.get(123).unwrap(), Some((OrCmp(false), 0x12)));
935        assert_eq!(*dict.root_extra(), OrCmp(false));
936    }
937
938    #[test]
939    fn dict_remove() {
940        let mut dict = AugDict::<u32, OrCmp, u32>::new();
941        assert_eq!(*dict.root_extra(), OrCmp(false));
942
943        for i in 0..10 {
944            assert!(dict.set(i, OrCmp(i % 2 == 0), i).unwrap());
945        }
946        assert_eq!(*dict.root_extra(), OrCmp(true));
947
948        let mut check_remove = |n: u32, expected: Option<(OrCmp, u32)>| -> anyhow::Result<()> {
949            let removed = dict.remove(n).context("Failed to remove")?;
950            anyhow::ensure!(removed == expected);
951            Ok(())
952        };
953
954        check_remove(0, Some((OrCmp(true), 0))).unwrap();
955
956        check_remove(4, Some((OrCmp(true), 4))).unwrap();
957
958        check_remove(9, Some((OrCmp(false), 9))).unwrap();
959        check_remove(9, None).unwrap();
960
961        check_remove(5, Some((OrCmp(false), 5))).unwrap();
962        check_remove(5, None).unwrap();
963
964        check_remove(100, None).unwrap();
965
966        check_remove(1, Some((OrCmp(false), 1))).unwrap();
967        check_remove(2, Some((OrCmp(true), 2))).unwrap();
968        check_remove(3, Some((OrCmp(false), 3))).unwrap();
969        check_remove(6, Some((OrCmp(true), 6))).unwrap();
970        check_remove(7, Some((OrCmp(false), 7))).unwrap();
971        check_remove(8, Some((OrCmp(true), 8))).unwrap();
972
973        assert!(dict.is_empty());
974    }
975
976    #[test]
977    fn dict_iter() {
978        let mut dict = AugDict::<u32, SomeValue, u32>::new();
979        assert_eq!(*dict.root_extra(), SomeValue(0));
980
981        let mut expected_extra = 0;
982        for i in 0..10 {
983            expected_extra += i;
984            dict.set(i, SomeValue(i), 9 - i).unwrap();
985        }
986        assert_eq!(*dict.root_extra(), SomeValue(expected_extra));
987
988        let size = dict.values().count();
989        assert_eq!(size, 10);
990
991        for (i, entry) in dict.iter().enumerate() {
992            let (key, aug, value) = entry.unwrap();
993            assert_eq!(SomeValue(key), aug);
994            assert_eq!(key, i as u32);
995            assert_eq!(value, 9 - i as u32);
996        }
997    }
998
999    #[cfg(feature = "models")]
1000    #[test]
1001    fn aug_test() {
1002        use crate::models::{AccountBlock, CurrencyCollection};
1003        use crate::prelude::Boc;
1004
1005        let boc = Boc::decode(include_bytes!("./tests/account_blocks_aug_dict.boc")).unwrap();
1006
1007        let original_dict = boc
1008            .parse::<AugDict<HashBytes, CurrencyCollection, AccountBlock>>()
1009            .unwrap();
1010
1011        let mut data = Vec::new();
1012        for entry in original_dict.iter() {
1013            data.push(entry.unwrap());
1014        }
1015        data.reverse();
1016
1017        let mut new_dict: AugDict<HashBytes, CurrencyCollection, AccountBlock> = AugDict::new();
1018        for (key, aug, value) in data.iter() {
1019            new_dict.add(key, aug, value).unwrap();
1020        }
1021        assert_eq!(new_dict.root_extra(), original_dict.root_extra());
1022
1023        let serialized = CellBuilder::build_from(&new_dict).unwrap();
1024        assert_eq!(serialized.repr_hash(), boc.repr_hash());
1025
1026        for (key, _, _) in data.iter() {
1027            new_dict.remove(key).unwrap();
1028        }
1029        assert!(new_dict.is_empty());
1030        assert_eq!(new_dict.root_extra(), &CurrencyCollection::ZERO);
1031    }
1032
1033    #[test]
1034    fn build_from_array() {
1035        for entries in [
1036            &[
1037                (0u32, SomeValue(123), 1u32),
1038                (1, SomeValue(10), 2),
1039                (2, SomeValue(20), 4),
1040                (2, SomeValue(20), 3),
1041                (3, SomeValue(40), 4),
1042                (4, SomeValue(50), 5),
1043            ][..],
1044            &[
1045                (534837844, SomeValue(331123), 3117028142),
1046                (1421713188, SomeValue(5345345), 3155891450),
1047                (1526242096, SomeValue(567567), 2789399854),
1048                (1971086295, SomeValue(5345), 1228713494),
1049                (4258889371, SomeValue(4956495), 3256452222),
1050            ],
1051        ] {
1052            let result = AugDict::<u32, SomeValue, u32>::try_from_sorted_slice(entries).unwrap();
1053
1054            let mut dict = AugDict::<u32, SomeValue, u32>::new();
1055            for (k, a, v) in entries {
1056                dict.add(k, a, v).unwrap();
1057            }
1058
1059            println!("{}", result.dict.root.as_ref().unwrap().display_tree());
1060            println!(
1061                "BOC: {}",
1062                crate::boc::BocRepr::encode_base64(&result).unwrap()
1063            );
1064
1065            assert_eq!(result, dict);
1066        }
1067    }
1068
1069    #[test]
1070    fn build_from_any_array() {
1071        for _ in 0..100 {
1072            let n = 1 + rand::random::<usize>() % 1000;
1073            let mut entries = (0..n)
1074                .map(|_| {
1075                    (
1076                        rand::random::<u32>(),
1077                        rand::random::<SomeValue>(),
1078                        rand::random::<u32>(),
1079                    )
1080                })
1081                .collect::<Vec<_>>();
1082            entries.sort_by_key(|(k, _, _)| *k);
1083
1084            let built_from_dict =
1085                AugDict::<u32, SomeValue, u32>::try_from_sorted_slice(&entries).unwrap();
1086
1087            let mut dict = AugDict::<u32, SomeValue, u32>::new();
1088            for (k, a, v) in entries {
1089                dict.add(k, a, v).unwrap();
1090            }
1091
1092            // println!("{}", built_from_dict.as_ref().unwrap().display_tree());
1093
1094            assert_eq!(built_from_dict, dict);
1095        }
1096    }
1097
1098    #[test]
1099    fn search_by_lt() {
1100        #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Store, Load)]
1101        struct MaxValue(u64);
1102
1103        impl AugDictExtra for MaxValue {
1104            fn comp_add(
1105                left: &mut CellSlice,
1106                right: &mut CellSlice,
1107                b: &mut CellBuilder,
1108                _: &mut dyn CellContext,
1109            ) -> Result<(), Error> {
1110                let left = left.load_u64()?;
1111                let right = right.load_u64()?;
1112                b.store_u64(left.max(right))
1113            }
1114        }
1115
1116        let mut items = AugDict::<u32, MaxValue, u32>::new();
1117        items.set(0, MaxValue(100), 123).unwrap();
1118        items.set(2, MaxValue(150), 234).unwrap();
1119        items.set(4, MaxValue(200), 345).unwrap();
1120        items.set(6, MaxValue(350), 456).unwrap();
1121        items.set(7, MaxValue(300), 567).unwrap();
1122        items.set(8, MaxValue(250), 678).unwrap();
1123
1124        // Search by ordering
1125        let highest = items.find_by_extra(std::cmp::Ordering::Greater).unwrap();
1126        assert_eq!(highest, Some((6, MaxValue(350), 456)));
1127    }
1128}
everscale_types/dict/aug.rs

everscale_types/dict/
aug.rs
