Skip to main content

ssz/
collections.rs

1//! SSZ blanket impls for `alloc::collections::BTreeMap`.
2//!
3//! Wire format: equivalent to `List<(K, V), MAX_BTREE_LEN>` — i.e., a flat
4//! list of sorted `(K, V)` pairs. Decode rejects out-of-order or duplicate
5//! keys.
6//!
7//! Hash form: same as `List<(K, V), MAX_BTREE_LEN>` — merkleize the pair
8//! roots, mix in length.
9
10use alloc::collections::BTreeMap;
11use alloc::vec::Vec;
12use digest::Digest;
13use digest::typenum::U32;
14
15use crate::merkle::{merkleize, mix_in_length};
16use crate::vector::decode_var_collection;
17use crate::{BYTES_PER_LENGTH_OFFSET, Decode, DecodeError, Encode, HashTreeRoot, read_offset};
18
19/// Implicit cap on `BTreeMap` length, in elements. Chosen as `1 << 32`
20/// (matches the SCALE `u32` count-prefix cap).
21pub const MAX_BTREE_LEN: u64 = 1u64 << 32;
22
23// --------------------------------------------------------------------------
24// (K, V) tuple impls — used by BTreeMap as the element type.
25// --------------------------------------------------------------------------
26
27impl<A: Encode, B: Encode> Encode for (A, B) {
28    fn is_ssz_fixed_len() -> bool {
29        A::is_ssz_fixed_len() && B::is_ssz_fixed_len()
30    }
31    fn ssz_fixed_len() -> usize {
32        if Self::is_ssz_fixed_len() {
33            A::ssz_fixed_len() + B::ssz_fixed_len()
34        } else {
35            BYTES_PER_LENGTH_OFFSET
36        }
37    }
38    fn ssz_bytes_len(&self) -> usize {
39        if Self::is_ssz_fixed_len() {
40            A::ssz_fixed_len() + B::ssz_fixed_len()
41        } else {
42            let mut total = 0usize;
43            // Fixed slot for each field (offsets for var-length).
44            total += if A::is_ssz_fixed_len() {
45                A::ssz_fixed_len()
46            } else {
47                BYTES_PER_LENGTH_OFFSET
48            };
49            total += if B::is_ssz_fixed_len() {
50                B::ssz_fixed_len()
51            } else {
52                BYTES_PER_LENGTH_OFFSET
53            };
54            // Variable payloads.
55            if !A::is_ssz_fixed_len() {
56                total += self.0.ssz_bytes_len();
57            }
58            if !B::is_ssz_fixed_len() {
59                total += self.1.ssz_bytes_len();
60            }
61            total
62        }
63    }
64    fn ssz_append(&self, buf: &mut Vec<u8>) {
65        if Self::is_ssz_fixed_len() {
66            self.0.ssz_append(buf);
67            self.1.ssz_append(buf);
68            return;
69        }
70        encode_container_pair(&self.0, &self.1, buf);
71    }
72}
73
74fn encode_container_pair<A: Encode, B: Encode>(a: &A, b: &B, buf: &mut Vec<u8>) {
75    // SSZ container layout: fixed-region first (with placeholders for
76    // variable-field offsets), then variable-region.
77    let a_fixed = if A::is_ssz_fixed_len() {
78        A::ssz_fixed_len()
79    } else {
80        BYTES_PER_LENGTH_OFFSET
81    };
82    let b_fixed = if B::is_ssz_fixed_len() {
83        B::ssz_fixed_len()
84    } else {
85        BYTES_PER_LENGTH_OFFSET
86    };
87    let fixed_region_size = a_fixed + b_fixed;
88    let start = buf.len();
89    // Reserve fixed region.
90    buf.resize(start + fixed_region_size, 0u8);
91
92    let mut running = fixed_region_size as u32;
93
94    // Field A.
95    if A::is_ssz_fixed_len() {
96        // Encode directly into the placeholder region.
97        let mut tmp: Vec<u8> = Vec::new();
98        a.ssz_append(&mut tmp);
99        debug_assert_eq!(tmp.len(), a_fixed);
100        buf[start..start + a_fixed].copy_from_slice(&tmp);
101    } else {
102        // Write offset, append payload.
103        buf[start..start + 4].copy_from_slice(&running.to_le_bytes());
104        let before = buf.len();
105        a.ssz_append(buf);
106        let after = buf.len();
107        running = running
108            .checked_add((after - before) as u32)
109            .expect("ssz offset overflow");
110    }
111
112    // Field B.
113    let b_start = start + a_fixed;
114    if B::is_ssz_fixed_len() {
115        let mut tmp: Vec<u8> = Vec::new();
116        b.ssz_append(&mut tmp);
117        debug_assert_eq!(tmp.len(), b_fixed);
118        buf[b_start..b_start + b_fixed].copy_from_slice(&tmp);
119    } else {
120        buf[b_start..b_start + 4].copy_from_slice(&running.to_le_bytes());
121        b.ssz_append(buf);
122    }
123}
124
125impl<A: Decode, B: Decode> Decode for (A, B) {
126    fn is_ssz_fixed_len() -> bool {
127        A::is_ssz_fixed_len() && B::is_ssz_fixed_len()
128    }
129    fn ssz_fixed_len() -> usize {
130        if Self::is_ssz_fixed_len() {
131            A::ssz_fixed_len() + B::ssz_fixed_len()
132        } else {
133            BYTES_PER_LENGTH_OFFSET
134        }
135    }
136    fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, DecodeError> {
137        if Self::is_ssz_fixed_len() {
138            let af = A::ssz_fixed_len();
139            let bf = B::ssz_fixed_len();
140            if bytes.len() != af + bf {
141                return Err(DecodeError::UnexpectedEof {
142                    expected: af + bf,
143                    actual: bytes.len(),
144                });
145            }
146            let a = A::from_ssz_bytes(&bytes[..af])?;
147            let b = B::from_ssz_bytes(&bytes[af..])?;
148            return Ok((a, b));
149        }
150        decode_container_pair::<A, B>(bytes)
151    }
152}
153
154fn decode_container_pair<A: Decode, B: Decode>(bytes: &[u8]) -> Result<(A, B), DecodeError> {
155    let a_fixed = if A::is_ssz_fixed_len() {
156        A::ssz_fixed_len()
157    } else {
158        BYTES_PER_LENGTH_OFFSET
159    };
160    let b_fixed = if B::is_ssz_fixed_len() {
161        B::ssz_fixed_len()
162    } else {
163        BYTES_PER_LENGTH_OFFSET
164    };
165    let fixed_region_size = a_fixed + b_fixed;
166    if bytes.len() < fixed_region_size {
167        return Err(DecodeError::UnexpectedEof {
168            expected: fixed_region_size,
169            actual: bytes.len(),
170        });
171    }
172
173    // Compute variable-region offsets in order, then slices.
174    let mut a_var_off: Option<usize> = None;
175    let mut b_var_off: Option<usize> = None;
176    if !A::is_ssz_fixed_len() {
177        a_var_off = Some(read_offset(bytes, 0)?);
178    }
179    if !B::is_ssz_fixed_len() {
180        b_var_off = Some(read_offset(bytes, a_fixed)?);
181    }
182    if let Some(off) = a_var_off
183        && off < fixed_region_size
184    {
185        return Err(DecodeError::InvalidOffset {
186            offset: off,
187            len: bytes.len(),
188            fixed: fixed_region_size,
189        });
190    }
191    if let Some(off) = b_var_off
192        && off < fixed_region_size
193    {
194        return Err(DecodeError::InvalidOffset {
195            offset: off,
196            len: bytes.len(),
197            fixed: fixed_region_size,
198        });
199    }
200    if let (Some(a), Some(b)) = (a_var_off, b_var_off)
201        && b < a
202    {
203        return Err(DecodeError::OffsetsNotMonotonic { prev: a, curr: b });
204    }
205
206    // Decode A.
207    let a_val = if A::is_ssz_fixed_len() {
208        A::from_ssz_bytes(&bytes[..a_fixed])?
209    } else {
210        let start = a_var_off.unwrap();
211        let end = b_var_off.unwrap_or(bytes.len());
212        if end > bytes.len() || start > end {
213            return Err(DecodeError::InvalidOffset {
214                offset: start,
215                len: bytes.len(),
216                fixed: fixed_region_size,
217            });
218        }
219        A::from_ssz_bytes(&bytes[start..end])?
220    };
221
222    // Decode B.
223    let b_val = if B::is_ssz_fixed_len() {
224        B::from_ssz_bytes(&bytes[a_fixed..a_fixed + b_fixed])?
225    } else {
226        let start = b_var_off.unwrap();
227        let end = bytes.len();
228        if start > end {
229            return Err(DecodeError::InvalidOffset {
230                offset: start,
231                len: bytes.len(),
232                fixed: fixed_region_size,
233            });
234        }
235        B::from_ssz_bytes(&bytes[start..end])?
236    };
237
238    Ok((a_val, b_val))
239}
240
241impl<A: HashTreeRoot, B: HashTreeRoot> HashTreeRoot for (A, B) {
242    fn hash_tree_root<D: Digest<OutputSize = U32>>(&self) -> [u8; 32] {
243        let roots = [self.0.hash_tree_root::<D>(), self.1.hash_tree_root::<D>()];
244        merkleize::<D>(&roots, 2)
245    }
246}
247
248// --------------------------------------------------------------------------
249// BTreeMap<K, V>
250// --------------------------------------------------------------------------
251
252impl<K: Encode + Ord, V: Encode> Encode for BTreeMap<K, V> {
253    fn is_ssz_fixed_len() -> bool {
254        false
255    }
256    fn ssz_fixed_len() -> usize {
257        BYTES_PER_LENGTH_OFFSET
258    }
259    fn ssz_bytes_len(&self) -> usize {
260        let elem_fixed = <(K, V) as Encode>::is_ssz_fixed_len();
261        if elem_fixed {
262            <(K, V) as Encode>::ssz_fixed_len() * self.len()
263        } else {
264            let n = self.len();
265            let mut total = n * BYTES_PER_LENGTH_OFFSET;
266            for (k, v) in self {
267                total += pair_len(k, v);
268            }
269            total
270        }
271    }
272    fn ssz_append(&self, buf: &mut Vec<u8>) {
273        let elem_fixed = <(K, V) as Encode>::is_ssz_fixed_len();
274        if elem_fixed {
275            // List<(K, V)> with fixed element → simple concatenation.
276            for (k, v) in self {
277                k.ssz_append(buf);
278                v.ssz_append(buf);
279            }
280            return;
281        }
282        // List<(K, V)> with variable element → offset table + payloads.
283        let entries: Vec<(&K, &V)> = self.iter().collect();
284        let header = entries.len() * BYTES_PER_LENGTH_OFFSET;
285        let start = buf.len();
286        buf.resize(start + header, 0u8);
287        let mut running = header as u32;
288        for (i, (k, v)) in entries.iter().enumerate() {
289            let off_pos = start + i * BYTES_PER_LENGTH_OFFSET;
290            buf[off_pos..off_pos + 4].copy_from_slice(&running.to_le_bytes());
291            let before = buf.len();
292            encode_container_pair(*k, *v, buf);
293            let after = buf.len();
294            running = running
295                .checked_add((after - before) as u32)
296                .expect("ssz offset overflow");
297        }
298    }
299}
300
301fn pair_len<K: Encode, V: Encode>(k: &K, v: &V) -> usize {
302    // Mirror the tuple `(K, V)`'s `ssz_bytes_len` accounting.
303    let a_fixed = if K::is_ssz_fixed_len() {
304        K::ssz_fixed_len()
305    } else {
306        BYTES_PER_LENGTH_OFFSET
307    };
308    let b_fixed = if V::is_ssz_fixed_len() {
309        V::ssz_fixed_len()
310    } else {
311        BYTES_PER_LENGTH_OFFSET
312    };
313    let mut total = a_fixed + b_fixed;
314    if !K::is_ssz_fixed_len() {
315        total += k.ssz_bytes_len();
316    }
317    if !V::is_ssz_fixed_len() {
318        total += v.ssz_bytes_len();
319    }
320    total
321}
322
323impl<K: Decode + Ord, V: Decode> Decode for BTreeMap<K, V> {
324    fn is_ssz_fixed_len() -> bool {
325        false
326    }
327    fn ssz_fixed_len() -> usize {
328        BYTES_PER_LENGTH_OFFSET
329    }
330    fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, DecodeError> {
331        let entries: Vec<(K, V)> = if <(K, V) as Decode>::is_ssz_fixed_len() {
332            let elem = <(K, V) as Decode>::ssz_fixed_len();
333            if elem == 0 {
334                return Err(DecodeError::Custom(
335                    "zero-sized fixed-length BTreeMap element",
336                ));
337            }
338            if !bytes.len().is_multiple_of(elem) {
339                return Err(DecodeError::LengthMismatch {
340                    expected: bytes.len().div_ceil(elem) * elem,
341                    actual: bytes.len(),
342                });
343            }
344            let n = bytes.len() / elem;
345            let mut out = Vec::with_capacity(n);
346            for i in 0..n {
347                let s = i * elem;
348                out.push(<(K, V) as Decode>::from_ssz_bytes(&bytes[s..s + elem])?);
349            }
350            out
351        } else {
352            decode_var_collection::<(K, V)>(bytes, None)?
353        };
354
355        let mut map = BTreeMap::new();
356        let mut prev: Option<&K> = None;
357        // BTreeMap doesn't let us peek prev by reference cleanly while
358        // inserting, so do the sorted-order check before insertion.
359        let mut staged: Vec<(K, V)> = Vec::with_capacity(entries.len());
360        for (k, v) in entries {
361            if let Some(p) = prev
362                && &k <= p
363            {
364                return Err(DecodeError::NotSorted);
365            }
366            staged.push((k, v));
367            prev = staged.last().map(|(k, _)| k);
368        }
369        for (k, v) in staged {
370            map.insert(k, v);
371        }
372        Ok(map)
373    }
374}
375
376impl<K: HashTreeRoot + Ord + Encode, V: HashTreeRoot + Encode> HashTreeRoot for BTreeMap<K, V> {
377    fn hash_tree_root<D: Digest<OutputSize = U32>>(&self) -> [u8; 32] {
378        let roots: Vec<[u8; 32]> = self
379            .iter()
380            .map(|(k, v)| {
381                let pair = [k.hash_tree_root::<D>(), v.hash_tree_root::<D>()];
382                merkleize::<D>(&pair, 2)
383            })
384            .collect();
385        let inner = merkleize::<D>(&roots, (MAX_BTREE_LEN as usize).max(1));
386        mix_in_length::<D>(inner, self.len() as u64)
387    }
388}
389
390// --------------------------------------------------------------------------
391// alloc::vec::Vec<T> — blanket SSZ impl
392//
393// Treated as `List<T, MAX_VEC_LEN>` semantically: variable-length list with
394// the same cap used for `BTreeMap`. Convenient for migrating struct fields
395// that hold plain `Vec<T>` without changing their type. Use `ssz::List<T,
396// N>` directly for types where a tighter type-level cap is desired.
397// --------------------------------------------------------------------------
398
399/// Implicit cap on `alloc::vec::Vec` length, in elements. Matches the legacy
400/// SCALE `u32` count-prefix cap (`1 << 32`).
401pub const MAX_VEC_LEN: u64 = 1u64 << 32;
402
403impl<T: Encode> Encode for Vec<T> {
404    fn is_ssz_fixed_len() -> bool {
405        false
406    }
407    fn ssz_fixed_len() -> usize {
408        BYTES_PER_LENGTH_OFFSET
409    }
410    fn ssz_bytes_len(&self) -> usize {
411        if T::is_ssz_fixed_len() {
412            T::ssz_fixed_len() * self.len()
413        } else {
414            let mut total = self.len() * BYTES_PER_LENGTH_OFFSET;
415            for item in self {
416                total += item.ssz_bytes_len();
417            }
418            total
419        }
420    }
421    fn ssz_append(&self, buf: &mut Vec<u8>) {
422        if T::is_ssz_fixed_len() {
423            for item in self {
424                item.ssz_append(buf);
425            }
426            return;
427        }
428        // Variable-length element: offset table + payloads.
429        let n = self.len();
430        let header = n * BYTES_PER_LENGTH_OFFSET;
431        let start = buf.len();
432        buf.resize(start + header, 0u8);
433        let mut running = header as u32;
434        for (i, item) in self.iter().enumerate() {
435            let off_pos = start + i * BYTES_PER_LENGTH_OFFSET;
436            buf[off_pos..off_pos + 4].copy_from_slice(&running.to_le_bytes());
437            let before = buf.len();
438            item.ssz_append(buf);
439            let after = buf.len();
440            running = running
441                .checked_add((after - before) as u32)
442                .expect("ssz offset overflow");
443        }
444    }
445}
446
447impl<T: Decode> Decode for Vec<T> {
448    fn is_ssz_fixed_len() -> bool {
449        false
450    }
451    fn ssz_fixed_len() -> usize {
452        BYTES_PER_LENGTH_OFFSET
453    }
454    fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, DecodeError> {
455        if T::is_ssz_fixed_len() {
456            let elem = T::ssz_fixed_len();
457            if elem == 0 {
458                return Err(DecodeError::Custom("zero-sized fixed-length Vec element"));
459            }
460            if !bytes.len().is_multiple_of(elem) {
461                return Err(DecodeError::LengthMismatch {
462                    expected: bytes.len().div_ceil(elem) * elem,
463                    actual: bytes.len(),
464                });
465            }
466            let n = bytes.len() / elem;
467            if (n as u64) > MAX_VEC_LEN {
468                return Err(DecodeError::BoundExceeded {
469                    len: n as u64,
470                    bound: MAX_VEC_LEN,
471                });
472            }
473            let mut out = Vec::with_capacity(n);
474            for i in 0..n {
475                let s = i * elem;
476                out.push(T::from_ssz_bytes(&bytes[s..s + elem])?);
477            }
478            Ok(out)
479        } else {
480            let out: Vec<T> = decode_var_collection::<T>(bytes, None)?;
481            if (out.len() as u64) > MAX_VEC_LEN {
482                return Err(DecodeError::BoundExceeded {
483                    len: out.len() as u64,
484                    bound: MAX_VEC_LEN,
485                });
486            }
487            Ok(out)
488        }
489    }
490}
491
492impl<T: HashTreeRoot + Encode> HashTreeRoot for Vec<T> {
493    fn hash_tree_root<D: Digest<OutputSize = U32>>(&self) -> [u8; 32] {
494        let len = self.len() as u64;
495        let inner_root = if T::is_basic_type() {
496            let mut buf: Vec<u8> = Vec::new();
497            for t in self {
498                t.ssz_append(&mut buf);
499            }
500            let chunks = crate::merkle::pack_bytes(&buf);
501            let cap_bytes = (MAX_VEC_LEN as usize).saturating_mul(T::ssz_fixed_len());
502            let chunk_limit = cap_bytes.div_ceil(32).max(1);
503            merkleize::<D>(&chunks, chunk_limit)
504        } else {
505            let roots: Vec<[u8; 32]> = self.iter().map(|t| t.hash_tree_root::<D>()).collect();
506            merkleize::<D>(&roots, (MAX_VEC_LEN as usize).max(1))
507        };
508        mix_in_length::<D>(inner_root, len)
509    }
510}