pint_abi/
key.rs

1//! Items related to key construction.
2
3use crate::types::essential::{Key, Word};
4
5/// Represents a dynamically provided key element (e.g. map key or array index).
6///
7/// ABI-generated `Mutations` builders maintain a stack of key `Elem`s as users
8/// enter nested map/array entries.
9#[derive(Debug)]
10pub enum Elem {
11    /// A key element is provided by a map key.
12    MapKey(Key),
13    /// A key element is provided by an array index.
14    ArrayIx(usize),
15}
16
17/// Represents statically provided information about a storage var's type nesting to assist with
18/// key construction.
19///
20/// This is similar to the `pint_abi_visit::Nesting` type, but only contains nesting information
21/// relevant to key construction.
22#[derive(Debug)]
23pub enum Nesting {
24    Var { ix: usize },
25    MapEntry,
26    TupleField { flat_ix: usize },
27    ArrayElem { elem_len: usize },
28}
29
30/// Construct a full key for a given nested type and the associated stack of
31/// dynamically provided key elements.
32pub fn construct(nesting: &[Nesting], key_elems: &[Elem]) -> Key {
33    let mut nesting_iter = nesting.iter().peekable();
34    let mut key_elems_iter = key_elems.iter();
35    let mut words = vec![];
36    while let Some(nesting) = nesting_iter.next() {
37        match nesting {
38            // Push the storage var index directly. This is always the first nesting.
39            Nesting::Var { ix } => {
40                let word = Word::try_from(*ix).expect("out of `Word` range");
41                words.push(word);
42            }
43            // If this is a map entry, take the map key from the key elements.
44            Nesting::MapEntry => {
45                let key = expect_map_key(key_elems_iter.next());
46                words.extend(key);
47            }
48            // If we've encountered a tuple field, flatten any further tuple
49            // nestings and use the deepest `flat_ix` which represents the
50            // flattened index of the tuple field.
51            Nesting::TupleField { flat_ix } => {
52                let word = flattened_word(0, *flat_ix, &mut nesting_iter, &mut key_elems_iter);
53                words.push(word);
54            }
55            // If this is an array element, determine the key from the array
56            // index. For directly nested arrays, detemine the flattened index
57            // by multiplying the index by the inner array lenghts.
58            Nesting::ArrayElem { elem_len } => {
59                let arr_ix = elem_len * expect_array_ix(key_elems_iter.next());
60                let word = flattened_word(arr_ix, 0, &mut nesting_iter, &mut key_elems_iter);
61                words.push(word);
62            }
63        }
64    }
65    words
66}
67
68/// Flattens directly nested tuple and array nestings into a single word for a key.
69///
70/// Flattening begins once an array or tuple nesting is encountered.
71///
72/// If the first nesting is an array, provide the `array_ix` and a zeroed
73/// `tuple_flat_ix`.
74///
75/// If the first nesting is a tuple, provide the `tuple_flat_ix` and a zeroed
76/// `array_ix`.
77///
78/// The `nestings` iterator's `next` method will only be called for each
79/// directly nested array and tuple.
80///
81/// The `key_elems` iterator's `next` method will only be called for each
82/// directly nested array.
83fn flattened_word<'a>(
84    mut array_ix: usize,
85    mut tuple_flat_ix: usize,
86    nestings: &mut std::iter::Peekable<impl Iterator<Item = &'a Nesting>>,
87    key_elems: &mut impl Iterator<Item = &'a Elem>,
88) -> Word {
89    let mut sum_ix = 0;
90    while let Some(nesting) = nestings.peek() {
91        match nesting {
92            Nesting::ArrayElem { elem_len } => {
93                // If the previous nesting was a tuple field, add its flattened ix to the sum.
94                sum_ix += tuple_flat_ix;
95                tuple_flat_ix = 0;
96                // Sum the new array index.
97                array_ix += elem_len * expect_array_ix(key_elems.next());
98            }
99            Nesting::TupleField { flat_ix } => {
100                // If the previous nesting was an array nesting, add it to the sum.
101                sum_ix += array_ix;
102                array_ix = 0;
103                // Sum the nested tuple's flat index onto the current.
104                tuple_flat_ix += *flat_ix;
105            }
106            // If we encounter a non-array, non-tuple nesting, flattening is complete.
107            _ => break,
108        }
109        nestings.next();
110    }
111    sum_ix += array_ix + tuple_flat_ix;
112    Word::try_from(sum_ix).expect("out of `Word` range")
113}
114
115fn expect_array_ix(elem: Option<&Elem>) -> usize {
116    match elem {
117        Some(Elem::ArrayIx(ix)) => *ix,
118        elem => unreachable!(
119            "invalid key `Elem` provided by ABI-generated code: \
120            expected array index, found {elem:?}",
121        ),
122    }
123}
124
125fn expect_map_key(elem: Option<&Elem>) -> Key {
126    match elem {
127        Some(Elem::MapKey(key)) => key.clone(),
128        elem => unreachable!(
129            "invalid key `Elem` provided by ABI-generated code: \
130            expected map key, found {elem:?}",
131        ),
132    }
133}
pint_abi/key.rs

pint_abi/
key.rs
