tycho_types/dict/
mod.rs

1//! Dictionary implementation.
2
3use std::ops::ControlFlow;
4
5pub use self::aug::*;
6pub use self::ops::*;
7pub use self::raw::*;
8pub use self::typed::*;
9use crate::cell::*;
10use crate::error::Error;
11
12mod aug;
13mod raw;
14mod typed;
15
16mod ops {
17    pub(crate) use self::build::{
18        AugMergeStackItemMode, MergeStackItem, MergeStackItemMode, SimpleMergeStackItemMode,
19    };
20    pub use self::build::{build_aug_dict_from_sorted_iter, build_dict_from_sorted_iter};
21    pub use self::find::{
22        aug_dict_find_by_extra, dict_find_bound, dict_find_bound_owned, dict_find_owned,
23    };
24    pub use self::get::{dict_get, dict_get_owned, dict_get_subdict};
25    pub use self::insert::{aug_dict_insert, dict_insert, dict_insert_owned};
26    pub use self::modify::{aug_dict_modify_from_sorted_iter, dict_modify_from_sorted_iter};
27    pub use self::remove::{aug_dict_remove_owned, dict_remove_bound_owned, dict_remove_owned};
28    pub use self::split_merge::{
29        PartialSplitDict, aug_dict_merge_siblings, dict_merge, dict_merge_siblings,
30        dict_split_by_prefix, dict_split_raw,
31    };
32
33    mod build;
34    mod find;
35    mod get;
36    mod insert;
37    mod modify;
38    mod remove;
39    mod split_merge;
40}
41
42/// Type which can be used as a dictionary key.
43pub trait DictKey {
44    /// Length in bits for a dictionary key.
45    const BITS: u16;
46}
47
48impl<T: DictKey> DictKey for &T {
49    const BITS: u16 = T::BITS;
50}
51
52impl<T: DictKey> DictKey for &mut T {
53    const BITS: u16 = T::BITS;
54}
55
56/// Type which can be stored as a dict key.
57pub trait StoreDictKey: DictKey {
58    /// Stores key bits into a builder data.
59    fn store_into_data(&self, data: &mut CellDataBuilder) -> Result<(), Error>;
60}
61
62impl<T: StoreDictKey> StoreDictKey for &T {
63    #[inline]
64    fn store_into_data(&self, data: &mut CellDataBuilder) -> Result<(), Error> {
65        T::store_into_data(self, data)
66    }
67}
68
69/// Type which can be loaded as a dict key.
70pub trait LoadDictKey: DictKey + Sized {
71    /// Creates a key from a raw builder data.
72    fn load_from_data(data: &CellDataBuilder) -> Option<Self>;
73}
74
75macro_rules! impl_dict_key {
76    ($($ty:ty => {
77        bits: $bits:literal,
78        store: |$this:ident, $builder:ident| $store_expr:expr,
79        load: |$raw_data:ident| $load_expr:expr$(,)?
80    }),*,) => {$(
81        impl DictKey for $ty {
82            const BITS: u16 = $bits;
83        }
84
85        impl StoreDictKey for $ty {
86            #[inline]
87            fn store_into_data(&self, $builder: &mut CellDataBuilder) -> Result<(), Error> {
88                let $this = self;
89                $store_expr
90            }
91        }
92
93        impl LoadDictKey for $ty {
94            #[inline]
95            fn load_from_data(b: &CellDataBuilder) -> Option<Self> {
96                let $raw_data = b.raw_data();
97                Some($load_expr)
98            }
99        }
100    )*};
101}
102
103impl_dict_key! {
104    bool => {
105        bits: 1,
106        store: |x, b| b.store_bit(*x),
107        load: |d| d[0] & 0x80 != 0,
108    },
109    u8 => {
110        bits: 8,
111        store: |x, b| b.store_u8(*x),
112        load: |d| d[0],
113    },
114    i8 => {
115        bits: 8,
116        store: |x, b| b.store_u8(*x as u8),
117        load: |d| d[0] as i8,
118    },
119    u16 => {
120        bits: 16,
121        store: |x, b| b.store_u16(*x),
122        load: |d| u16::from_be_bytes([d[0], d[1]]),
123    },
124    i16 => {
125        bits: 16,
126        store: |x, b| b.store_u16(*x as u16),
127        load: |d| i16::from_be_bytes([d[0], d[1]]),
128    },
129    u32 => {
130        bits: 32,
131        store: |x, b| b.store_u32(*x),
132        load: |d| u32::from_be_bytes(d[..4].try_into().unwrap()),
133    },
134    i32 => {
135        bits: 32,
136        store: |x, b| b.store_u32(*x as u32),
137        load: |d| i32::from_be_bytes(d[..4].try_into().unwrap()),
138    },
139    u64 => {
140        bits: 64,
141        store: |x, b| b.store_u64(*x),
142        load: |d| u64::from_be_bytes(d[..8].try_into().unwrap()),
143    },
144    i64 => {
145        bits: 64,
146        store: |x, b| b.store_u64(*x as u64),
147        load: |d| i64::from_be_bytes(d[..8].try_into().unwrap())
148    },
149    u128 => {
150        bits: 128,
151        store: |x, b| b.store_u128(*x),
152        load: |d| u128::from_be_bytes(d[..16].try_into().unwrap()),
153    },
154    i128 => {
155        bits: 128,
156        store: |x, b| b.store_u128(*x as u128),
157        load: |d| i128::from_be_bytes(d[..16].try_into().unwrap()),
158    },
159    [u8; 16] => {
160        bits: 128,
161        store: |x, b| b.store_raw(x, 128),
162        load: |d| d[..16].try_into().unwrap(),
163    },
164    [u8; 20] => {
165        bits: 160,
166        store: |x, b| b.store_raw(x, 160),
167        load: |d| d[..20].try_into().unwrap(),
168    },
169    [u8; 32] => {
170        bits: 256,
171        store: |x, b| b.store_raw(x, 256),
172        load: |d| d[..32].try_into().unwrap(),
173    },
174    HashBytes => {
175        bits: 256,
176        store: |x, b| b.store_u256(x),
177        load: |d| HashBytes(d[..32].try_into().unwrap()),
178    },
179     (u64, u32) => {
180        bits: 96,
181        store: |x, b| {
182            ok!(b.store_u64(x.0));
183            b.store_u32(x.1)
184        },
185        load: |d| {
186            (u64::from_be_bytes(d[..8].try_into().unwrap()), u32::from_be_bytes(d[8..12].try_into().unwrap()))
187        },
188     },
189}
190
191/// `AugDict` search control flow.
192pub trait SearchByExtra<A> {
193    /// Returns if the leaf extra satisfies the condition.
194    fn on_leaf(&mut self, leaf_extra: &A) -> bool {
195        _ = leaf_extra;
196        true
197    }
198
199    /// Returns which branch satisfies the condition.
200    fn on_edge(&mut self, left_extra: &A, right_extra: &A) -> ControlFlow<(), Branch>;
201}
202
203impl<A, T: SearchByExtra<A>> SearchByExtra<A> for &mut T {
204    #[inline]
205    fn on_leaf(&mut self, leaf_extra: &A) -> bool {
206        T::on_leaf(self, leaf_extra)
207    }
208
209    #[inline]
210    fn on_edge(&mut self, left_extra: &A, right_extra: &A) -> ControlFlow<(), Branch> {
211        T::on_edge(self, left_extra, right_extra)
212    }
213}
214
215impl<A> SearchByExtra<A> for Branch {
216    #[inline]
217    fn on_edge(&mut self, _: &A, _: &A) -> ControlFlow<(), Branch> {
218        ControlFlow::Continue(*self)
219    }
220}
221
222impl<A: Ord> SearchByExtra<A> for std::cmp::Ordering {
223    #[inline]
224    fn on_edge(&mut self, left_extra: &A, right_extra: &A) -> ControlFlow<(), Branch> {
225        ControlFlow::Continue(if *self == left_extra.cmp(right_extra) {
226            Branch::Left
227        } else {
228            Branch::Right
229        })
230    }
231}
232
233/// Dictionary insertion mode.
234#[derive(Debug, Clone, Copy, Eq, PartialEq)]
235#[repr(u8)]
236pub enum SetMode {
237    /// Sets the value associated with the key in the dictionary.
238    Set = 0b11,
239    /// Sets the value associated with the key in the dictionary
240    /// only if the key was already present in it.
241    Replace = 0b01,
242    /// Sets the value associated with key in dictionary,
243    /// but only if it is not already present.
244    Add = 0b10,
245}
246
247impl SetMode {
248    /// Returns `true` if the new value can replace the old value for the same key.
249    #[inline]
250    pub const fn can_replace(self) -> bool {
251        self as u8 & 0b01 != 0
252    }
253
254    /// Returns `true` if inserting a value can add a new key to the dictionary.
255    #[inline]
256    pub const fn can_add(self) -> bool {
257        self as u8 & 0b10 != 0
258    }
259}
260
261/// A type for a comparator function for [`AugDict`].
262///
263/// ## Args
264/// - `left` - a left branch data.
265/// - `right` - a right branch data.
266/// - `builder` - a builder to write the result.
267/// - `context` - a cell context.
268pub type AugDictFn =
269    fn(&mut CellSlice, &mut CellSlice, &mut CellBuilder, &dyn CellContext) -> Result<(), Error>;
270
271/// Creates a leaf node
272fn make_leaf(
273    key: &CellSlice<'_>,
274    key_bit_len: u16,
275    value: &dyn Store,
276    context: &dyn CellContext,
277) -> Result<Cell, Error> {
278    let mut builder = CellBuilder::new();
279    ok!(write_label(key, key_bit_len, &mut builder));
280    ok!(value.store_into(&mut builder, context));
281    builder.build_ext(context)
282}
283
284/// Creates a leaf node with extra value
285fn make_leaf_with_extra(
286    key: &CellSlice<'_>,
287    key_bit_len: u16,
288    extra: &dyn Store,
289    value: &dyn Store,
290    context: &dyn CellContext,
291) -> Result<Cell, Error> {
292    let mut builder = CellBuilder::new();
293    ok!(write_label(key, key_bit_len, &mut builder));
294    ok!(extra.store_into(&mut builder, context));
295    ok!(value.store_into(&mut builder, context));
296    builder.build_ext(context)
297}
298
299/// Splits an edge or leaf
300fn split_edge(
301    data: &CellSlice,
302    prefix: &mut CellSlice,
303    lcp: &CellSlice,
304    key: &mut CellSlice,
305    value: &dyn Store,
306    context: &dyn CellContext,
307) -> Result<Cell, Error> {
308    // Advance the key
309    let prev_key_bit_len = key.size_bits();
310    ok!(key.skip_first(lcp.size_bits() + 1, 0));
311
312    // Read the next bit from the data
313    prefix.skip_first(lcp.size_bits(), 0).ok();
314    let old_to_right = ok!(prefix.load_bit());
315
316    // Create a leaf for the old value
317    let mut left = ok!(make_leaf(prefix, key.size_bits(), data, context));
318    // Create a leaf for the right value
319    let mut right = ok!(make_leaf(key, key.size_bits(), value, context));
320
321    // The part that starts with 1 goes to the right cell
322    if old_to_right {
323        std::mem::swap(&mut left, &mut right);
324    }
325
326    // Create fork
327    let mut builder = CellBuilder::new();
328    ok!(write_label(lcp, prev_key_bit_len, &mut builder));
329    ok!(builder.store_reference(left));
330    ok!(builder.store_reference(right));
331    builder.build_ext(context)
332}
333
334#[allow(clippy::too_many_arguments)]
335fn split_aug_edge(
336    data: &mut CellSlice,
337    prefix: &mut CellSlice,
338    lcp: &CellSlice,
339    key: &mut CellSlice,
340    extra: &dyn Store,
341    value: &dyn Store,
342    comparator: AugDictFn,
343    context: &dyn CellContext,
344) -> Result<Cell, Error> {
345    // Advance the key
346    let prev_key_bit_len = key.size_bits();
347    ok!(key.skip_first(lcp.size_bits() + 1, 0));
348
349    // Read the next bit from the data
350    prefix.skip_first(lcp.size_bits(), 0).ok();
351    let old_to_right = ok!(prefix.load_bit());
352
353    // Create a leaf for the old value
354    let mut left = ok!(make_leaf(prefix, key.size_bits(), data, context));
355    // Create a leaf for the new value
356    let mut right = ok!(make_leaf_with_extra(
357        key,
358        key.size_bits(),
359        extra,
360        value,
361        context
362    ));
363    // The part that starts with 1 goes to the right cell
364    if old_to_right {
365        std::mem::swap(&mut left, &mut right);
366    }
367
368    let left_slice = &mut ok!(left.as_slice());
369    let right_slice = &mut ok!(right.as_slice());
370    ok!(skip_label_full(left_slice, key.size_bits()));
371    ok!(skip_label_full(right_slice, key.size_bits()));
372
373    // Create fork edge
374    let mut builder = CellBuilder::new();
375    ok!(write_label(lcp, prev_key_bit_len, &mut builder));
376    ok!(builder.store_reference(left.clone()));
377    ok!(builder.store_reference(right.clone()));
378    ok!(comparator(left_slice, right_slice, &mut builder, context));
379    builder.build_ext(context)
380}
381
382/// Type alias for a pair of key and value as cell slice parts.
383pub type DictOwnedEntry = (CellDataBuilder, CellSliceParts);
384
385/// Dictionary bound.
386#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
387pub enum DictBound {
388    /// The lowest dictionary key.
389    Min,
390    /// The largest dictionary key.
391    Max,
392}
393
394impl DictBound {
395    fn update_direction(
396        self,
397        prefix: &CellSlice<'_>,
398        signed: bool,
399        direction: &mut Option<Branch>,
400    ) -> Branch {
401        match direction {
402            // Compute direction by the first part
403            None => {
404                let mut branch = *direction.insert(self.into_branch());
405                // Invert first bit for signed keys if starting from the empty part
406                if signed && prefix.is_data_empty() {
407                    branch = branch.reversed();
408                }
409                branch
410            }
411            // Use the same direction for all remaining parts
412            Some(direction) => *direction,
413        }
414    }
415
416    fn into_branch(self) -> Branch {
417        match self {
418            Self::Min => Branch::Left,
419            Self::Max => Branch::Right,
420        }
421    }
422}
423
424/// Loads a non-empty dictionary from the root cell.
425pub fn dict_load_from_root(
426    slice: &mut CellSlice<'_>,
427    key_bit_len: u16,
428    context: &dyn CellContext,
429) -> Result<Cell, Error> {
430    let mut root = *slice;
431
432    let label = ok!(read_label(slice, key_bit_len));
433    if label.size_bits() != key_bit_len {
434        ok!(slice.skip_first(0, 2));
435        let root_bits = root.size_bits() - slice.size_bits();
436        let root_refs = root.size_refs() - slice.size_refs();
437        root = root.get_prefix(root_bits, root_refs)
438    } else {
439        slice.load_remaining();
440    }
441
442    let mut builder = CellBuilder::new();
443    ok!(builder.store_slice(root));
444    builder.build_ext(context)
445}
446
447fn rebuild_dict_from_stack(
448    mut segments: Vec<Segment<'_>>,
449    mut leaf: Cell,
450    context: &dyn CellContext,
451) -> Result<Cell, Error> {
452    // Rebuild the tree starting from leaves
453    while let Some(last) = segments.pop() {
454        // Load the opposite branch
455        let (left, right) = match last.next_branch {
456            Branch::Left => match last.data.reference_cloned(1) {
457                Some(cell) => (leaf, cell),
458                None => return Err(Error::CellUnderflow),
459            },
460            Branch::Right => match last.data.reference_cloned(0) {
461                Some(cell) => (cell, leaf),
462                None => return Err(Error::CellUnderflow),
463            },
464        };
465
466        let mut builder = CellBuilder::new();
467        ok!(builder.store_cell_data(last.data));
468        ok!(builder.store_reference(left));
469        ok!(builder.store_reference(right));
470        leaf = ok!(builder.build_ext(context));
471    }
472
473    Ok(leaf)
474}
475
476fn rebuild_aug_dict_from_stack(
477    mut segments: Vec<Segment<'_>>,
478    mut leaf: Cell,
479    comparator: AugDictFn,
480    context: &dyn CellContext,
481) -> Result<Cell, Error> {
482    // Rebuild the tree starting from leaves
483    while let Some(last) = segments.pop() {
484        // Load the opposite branch
485        let (left, right) = match last.next_branch {
486            Branch::Left => match last.data.reference_cloned(1) {
487                Some(cell) => (leaf, cell),
488                None => return Err(Error::CellUnderflow),
489            },
490            Branch::Right => match last.data.reference_cloned(0) {
491                Some(cell) => (cell, leaf),
492                None => return Err(Error::CellUnderflow),
493            },
494        };
495
496        let last_data_slice = ok!(last.data.as_slice());
497        let last_label = ok!(read_label(&mut last_data_slice.clone(), last.key_bit_len));
498
499        let child_key_bit_len = last.key_bit_len - last_label.size_bits() - 1;
500
501        let left_slice = &mut left.as_slice()?;
502        let right_slice = &mut right.as_slice()?;
503        ok!(skip_label_full(left_slice, child_key_bit_len));
504        ok!(skip_label_full(right_slice, child_key_bit_len));
505
506        let mut builder = CellBuilder::new();
507        ok!(write_label(&last_label, last.key_bit_len, &mut builder));
508        ok!(builder.store_reference(left.clone()));
509        ok!(builder.store_reference(right.clone()));
510        ok!(comparator(left_slice, right_slice, &mut builder, context));
511        leaf = ok!(builder.build_ext(context));
512    }
513
514    Ok(leaf)
515}
516
517fn skip_label_full(slice: &mut CellSlice<'_>, key_bit_len: u16) -> Result<(), Error> {
518    let label = ok!(read_label(slice, key_bit_len));
519    if label.size_bits() < key_bit_len {
520        // Cell is not a leaf so we need to skip 2 refs.
521        slice.skip_first(0, 2)
522    } else {
523        Ok(())
524    }
525}
526
527#[derive(Clone, Copy)]
528struct Segment<'a> {
529    data: &'a DynCell,
530    next_branch: Branch,
531    key_bit_len: u16,
532}
533
534impl Segment<'_> {
535    // Returns the new leaf and the removed leaf
536    fn rebuild_as_removed(
537        self,
538        key: &CellSlice<'_>,
539        prev_key_bit_len: u16,
540        context: &dyn CellContext,
541    ) -> Result<(Cell, Cell), Error> {
542        let index = self.next_branch as u8;
543
544        // Load value branch
545        let Some(value) = self.data.reference_cloned(index) else {
546            return Err(Error::CellUnderflow);
547        };
548        // NOTE: do not use gas here as it was accounted while loading `child` in previous block.
549        // TODO: change mode to `LoadMode::Noop` if copy-on-write for libraries is not ok.
550        let value = ok!(context.load_cell(value, LoadMode::Resolve));
551
552        // Load parent label
553        let pfx = ok!(read_label(
554            &mut self.data.as_slice_allow_exotic(),
555            prev_key_bit_len
556        ));
557
558        // Load the opposite branch
559        let mut opposite = match self.data.reference(1 - index) {
560            // TODO: change mode to `LoadMode::UseGas` if copy-on-write for libraries is not ok.
561            Some(cell) => ok!(context
562                .load_dyn_cell(cell, LoadMode::Full)
563                .and_then(CellSlice::new)),
564            None => return Err(Error::CellUnderflow),
565        };
566        let rem = ok!(read_label(&mut opposite, key.size_bits()));
567
568        // Build an edge cell
569        let mut builder = CellBuilder::new();
570        ok!(write_label_parts(
571            &pfx,
572            index == 0,
573            &rem,
574            prev_key_bit_len,
575            &mut builder
576        ));
577        ok!(builder.store_slice(opposite));
578        let leaf = ok!(builder.build_ext(context));
579
580        Ok((leaf, value))
581    }
582}
583
584fn write_label(key: &CellSlice, key_bit_len: u16, label: &mut CellBuilder) -> Result<(), Error> {
585    if key_bit_len == 0 || key.is_data_empty() {
586        return write_hml_empty(label);
587    }
588
589    let bits_for_len = (16 - key_bit_len.leading_zeros()) as u16;
590
591    let remaining_bits = key.size_bits();
592
593    let hml_short_len = 2 + 2 * remaining_bits;
594    let hml_long_len = 2 + bits_for_len + remaining_bits;
595    let hml_same_len = 3 + bits_for_len;
596
597    if hml_same_len < hml_long_len
598        && hml_same_len < hml_short_len
599        && let Some(bit) = key.test_uniform()
600    {
601        ok!(write_hml_same(bit, remaining_bits, bits_for_len, label));
602        return Ok(());
603    }
604
605    if hml_short_len <= MAX_BIT_LEN && hml_short_len <= hml_long_len {
606        ok!(write_hml_short_tag(remaining_bits, label));
607    } else if hml_long_len <= MAX_BIT_LEN {
608        ok!(write_hml_long_tag(remaining_bits, bits_for_len, label));
609    } else {
610        return Err(Error::InvalidData);
611    }
612
613    ok!(label.store_slice_data(key));
614    Ok(())
615}
616
617fn write_label_parts(
618    pfx: &CellSlice,
619    bit: bool,
620    rem: &CellSlice,
621    key_bit_len: u16,
622    label: &mut CellBuilder,
623) -> Result<(), Error> {
624    if key_bit_len == 0 {
625        return write_hml_empty(label);
626    }
627
628    let bits_for_len = (16 - key_bit_len.leading_zeros()) as u16;
629
630    let remaining_bits = pfx.size_bits() + 1 + rem.size_bits();
631
632    let hml_short_len = 2 + 2 * remaining_bits;
633    let hml_long_len = 2 + bits_for_len + remaining_bits;
634    let hml_same_len = 3 + bits_for_len;
635
636    if hml_same_len < hml_long_len
637        && hml_same_len < hml_short_len
638        && (pfx.is_data_empty() || matches!(pfx.test_uniform(), Some(p) if p == bit))
639        && (rem.is_data_empty() || matches!(rem.test_uniform(), Some(r) if r == bit))
640    {
641        return write_hml_same(bit, remaining_bits, bits_for_len, label);
642    }
643
644    if hml_short_len <= MAX_BIT_LEN && hml_short_len <= hml_long_len {
645        ok!(write_hml_short_tag(remaining_bits, label));
646    } else if hml_long_len <= MAX_BIT_LEN {
647        ok!(write_hml_long_tag(remaining_bits, bits_for_len, label));
648    } else {
649        return Err(Error::InvalidData);
650    }
651    ok!(label.store_slice_data(pfx));
652    ok!(label.store_bit(bit));
653    label.store_slice_data(rem)
654}
655
656fn read_label<'a>(label: &mut CellSlice<'a>, key_bit_len: u16) -> Result<CellSlice<'a>, Error> {
657    let bits_for_len = (16 - key_bit_len.leading_zeros()) as u16;
658
659    if bits_for_len == 0 && label.is_data_empty() {
660        Ok(label.get_prefix(0, 0))
661    } else if !ok!(label.load_bit()) {
662        read_hml_short(label)
663    } else if !ok!(label.load_bit()) {
664        read_hml_long(label, bits_for_len)
665    } else {
666        read_hml_same(label, bits_for_len)
667    }
668}
669
670fn write_hml_empty(label: &mut CellBuilder) -> Result<(), Error> {
671    label.store_zeros(2)
672}
673
674fn write_hml_short_tag(len: u16, label: &mut CellBuilder) -> Result<(), Error> {
675    ok!(label.store_bit_zero());
676
677    for _ in 0..len / 32 {
678        ok!(label.store_u32(u32::MAX));
679    }
680
681    let rem = len % 32;
682    if rem != 0 {
683        ok!(label.store_uint(u64::MAX, rem));
684    }
685    label.store_bit_zero()
686}
687
688fn read_hml_short<'a>(label: &mut CellSlice<'a>) -> Result<CellSlice<'a>, Error> {
689    let mut len = 0;
690    while ok!(label.load_bit()) {
691        len += 1;
692    }
693    let result = *label;
694    ok!(label.skip_first(len, 0));
695    Ok(result.get_prefix(len, 0))
696}
697
698fn write_hml_long_tag(len: u16, bits_for_len: u16, label: &mut CellBuilder) -> Result<(), Error> {
699    ok!(label.store_bit_one());
700    ok!(label.store_bit_zero());
701    label.store_uint(len as u64, bits_for_len)
702}
703
704fn read_hml_long<'a>(label: &mut CellSlice<'a>, bits_for_len: u16) -> Result<CellSlice<'a>, Error> {
705    let len = ok!(label.load_uint(bits_for_len)) as u16;
706    let result = *label;
707    ok!(label.skip_first(len, 0));
708    Ok(result.get_prefix(len, 0))
709}
710
711fn write_hml_same(
712    bit: bool,
713    len: u16,
714    bits_for_len: u16,
715    label: &mut CellBuilder,
716) -> Result<(), Error> {
717    ok!(label.store_small_uint(0b110 | bit as u8, 3));
718    label.store_uint(len as u64, bits_for_len)
719}
720
721fn read_hml_same<'a>(label: &mut CellSlice<'a>, bits_for_len: u16) -> Result<CellSlice<'a>, Error> {
722    let cell = match ok!(label.load_bit()) {
723        false => Cell::all_zeros_ref(),
724        true => Cell::all_ones_ref(),
725    };
726    let len = ok!(label.load_uint(bits_for_len)) as u16;
727
728    Ok(cell.as_slice_allow_exotic().get_prefix(len, 0))
729}
730
731/// Which branch to take when traversing the tree.
732#[derive(Debug, Clone, Copy, Eq, PartialEq)]
733pub enum Branch {
734    /// Branch for a key part that starts with bit 0
735    Left = 0,
736    /// Branch for a key part that starts with bit 1
737    Right = 1,
738}
739
740impl Branch {
741    /// Converts the branch to a boolean value.
742    pub fn into_bit(self) -> bool {
743        self == Self::Right
744    }
745
746    /// Returns the opposite branch.
747    pub fn reversed(self) -> Self {
748        match self {
749            Self::Left => Self::Right,
750            Self::Right => Self::Left,
751        }
752    }
753}
754
755impl From<bool> for Branch {
756    fn from(value: bool) -> Self {
757        if value { Self::Right } else { Self::Left }
758    }
759}
760
761#[cfg(test)]
762mod tests {
763    use super::*;
764
765    fn build_cell<F: FnOnce(&mut CellBuilder) -> Result<(), Error>>(f: F) -> Cell {
766        let mut builder = CellBuilder::new();
767        f(&mut builder).unwrap();
768        builder.build().unwrap()
769    }
770
771    #[test]
772    fn labels() -> anyhow::Result<()> {
773        let key_bit_len = 6;
774
775        // Build key
776        let key = {
777            let mut builder = CellBuilder::new();
778            builder.store_zeros(5)?;
779            builder.store_bit_one()?;
780            builder.build()?
781        };
782
783        // Build label
784        let label = {
785            let mut builder = CellBuilder::new();
786            write_label(&key.as_slice()?, key_bit_len, &mut builder)?;
787            builder.build().unwrap()
788        };
789
790        // Parse label
791        let parsed_key = read_label(&mut label.as_slice()?, key_bit_len)?;
792        let parsed_key = {
793            let mut builder = CellBuilder::new();
794            builder.store_slice(parsed_key)?;
795            builder.build()?
796        };
797
798        // Parsed key should be equal to the original
799        assert_eq!(key.as_ref(), parsed_key.as_ref());
800
801        let label = CellBuilder::from_raw_data(&[0xcc, 0x40], 9)
802            .unwrap()
803            .build()
804            .unwrap();
805        let prefix = read_label(&mut label.as_slice()?, 32).unwrap();
806
807        println!("{}", build_cell(|b| b.store_slice(prefix)).display_tree());
808        assert_eq!(prefix.test_uniform(), Some(false));
809
810        Ok(())
811    }
812
813    #[test]
814    fn build_from_array() {
815        for entries in [
816            &[(0u32, 1u32), (1, 2), (2, 4), (2, 3), (3, 4), (4, 5)][..],
817            &[
818                (534837844, 3117028142),
819                (1421713188, 3155891450),
820                (1526242096, 2789399854),
821                (1971086295, 1228713494),
822                (4258889371, 3256452222),
823            ],
824        ] {
825            let result =
826                build_dict_from_sorted_iter(entries.iter().copied(), Cell::empty_context())
827                    .unwrap();
828
829            let mut dict = Dict::<u32, u32>::new();
830            for (k, v) in entries {
831                dict.add(k, v).unwrap();
832            }
833
834            println!("{}", result.as_ref().unwrap().display_tree());
835            println!(
836                "BOC: {}",
837                crate::boc::BocRepr::encode_base64(&result).unwrap()
838            );
839
840            assert_eq!(result, dict.root);
841        }
842    }
843
844    #[test]
845    fn build_from_any_array() {
846        for _ in 0..100 {
847            let n = 1 + rand9::random::<u32>() % 1000;
848            let mut entries = (0..n)
849                .map(|_| (rand9::random::<u32>(), rand9::random::<u32>()))
850                .collect::<Vec<_>>();
851            entries.sort_by_key(|(k, _)| *k);
852
853            let built_from_dict =
854                build_dict_from_sorted_iter(entries.iter().copied(), Cell::empty_context())
855                    .unwrap();
856
857            let mut dict = Dict::<u32, u32>::new();
858            for (k, v) in entries {
859                dict.add(k, v).unwrap();
860            }
861
862            // println!("{}", built_from_dict.as_ref().unwrap().display_tree());
863
864            assert_eq!(built_from_dict, dict.root);
865        }
866    }
867}