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