iris_ztd/
noun.rs

1use crate::belt::based_check;
2use crate::crypto::cheetah::F6lt;
3use alloc::{boxed::Box, collections::btree_map::BTreeMap, format, string::String, vec, vec::Vec};
4use bitvec::prelude::{BitSlice, BitVec, Lsb0};
5use core::fmt;
6use ibig::UBig;
7use num_traits::Zero;
8use serde::de::{Error as DeError, SeqAccess, Visitor};
9use serde::{ser::SerializeTuple, Serialize, Serializer};
10use serde::{Deserialize, Deserializer};
11
12use crate::{belt::Belt, crypto::cheetah::CheetahPoint, Digest};
13
14/// A transparent wrapper that encodes as a zero atom if the value is `None`.
15#[repr(transparent)]
16#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Default)]
17pub struct Zeroable<T>(pub Option<T>);
18
19impl<T> core::ops::Deref for Zeroable<T> {
20    type Target = Option<T>;
21
22    fn deref(&self) -> &Self::Target {
23        &self.0
24    }
25}
26
27impl<T> core::ops::DerefMut for Zeroable<T> {
28    fn deref_mut(&mut self) -> &mut Self::Target {
29        &mut self.0
30    }
31}
32
33pub fn noun_serialize<T: NounEncode, S>(v: &T, serializer: S) -> Result<S::Ok, S::Error>
34where
35    S: Serializer,
36{
37    v.to_noun().serialize(serializer)
38}
39
40pub fn noun_deserialize<'de, T: NounDecode, D>(deserializer: D) -> Result<T, D::Error>
41where
42    D: Deserializer<'de>,
43{
44    let r = Noun::deserialize(deserializer)?;
45    T::from_noun(&r).ok_or_else(|| DeError::custom("unable to parse noun"))
46}
47
48#[derive(Clone, Debug, PartialEq, Eq)]
49pub enum Noun {
50    Atom(UBig),
51    Cell(Box<Noun>, Box<Noun>),
52}
53
54impl Noun {
55    #[allow(clippy::inherent_to_string_shadow_display)]
56    pub fn to_string(&self) -> String {
57        fn autocons(cell: &Noun) -> String {
58            match cell {
59                Noun::Atom(a) => format!("{}", a),
60                Noun::Cell(head, tail) => match &**tail {
61                    Noun::Cell(_, _) => format!("{} {}", autocons(head), autocons(tail)),
62                    Noun::Atom(a) if a.is_zero() => autocons(head),
63                    Noun::Atom(_) => format!("{} . {}", autocons(head), autocons(tail)),
64                },
65            }
66        }
67
68        match self {
69            Noun::Atom(a) => format!("{}", a),
70            Noun::Cell(head, tail) => {
71                format!("[{}]", autocons(&Noun::Cell(head.clone(), tail.clone())))
72            }
73        }
74    }
75}
76
77impl Serialize for Noun {
78    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
79    where
80        S: Serializer,
81    {
82        match self {
83            Self::Atom(v) => serializer.serialize_str(&alloc::format!("{v:x}")),
84            Self::Cell(a, b) => {
85                let mut tup = serializer.serialize_tuple(2)?;
86                tup.serialize_element(a)?;
87                tup.serialize_element(b)?;
88                tup.end()
89            }
90        }
91    }
92}
93
94impl<'de> Deserialize<'de> for Noun {
95    fn deserialize<D>(de: D) -> Result<Self, D::Error>
96    where
97        D: Deserializer<'de>,
98    {
99        struct V;
100
101        impl<'de> Visitor<'de> for V {
102            type Value = Noun;
103
104            fn expecting(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
105                f.write_str("atom or cell")
106            }
107
108            fn visit_str<E>(self, s: &str) -> Result<Self::Value, E>
109            where
110                E: DeError,
111            {
112                let n = UBig::from_str_radix(s, 16).map_err(E::custom)?;
113                Ok(Noun::Atom(n))
114            }
115
116            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
117            where
118                A: SeqAccess<'de>,
119            {
120                let a = seq
121                    .next_element::<Noun>()?
122                    .ok_or_else(|| DeError::custom("cell missing car"))?;
123                let b = seq
124                    .next_element::<Noun>()?
125                    .ok_or_else(|| DeError::custom("cell missing cdr"))?;
126                Ok(Noun::Cell(Box::new(a), Box::new(b)))
127            }
128        }
129
130        de.deserialize_any(V)
131    }
132}
133
134impl fmt::Display for Noun {
135    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
136        write!(f, "{}", self.to_string())
137    }
138}
139
140pub trait NounCode: NounEncode + NounDecode {}
141impl<T: NounEncode + NounDecode> NounCode for T {}
142
143pub trait NounEncode {
144    fn to_noun(&self) -> Noun;
145}
146
147pub trait NounDecode: Sized {
148    fn from_noun(noun: &Noun) -> Option<Self>;
149}
150
151fn atom(value: u64) -> Noun {
152    Noun::Atom(UBig::from(value))
153}
154
155fn cons(left: Noun, right: Noun) -> Noun {
156    Noun::Cell(Box::new(left), Box::new(right))
157}
158
159impl<T: NounEncode + ?Sized> NounEncode for &T {
160    fn to_noun(&self) -> Noun {
161        (**self).to_noun()
162    }
163}
164
165impl<T: NounEncode + ?Sized> NounEncode for Box<T> {
166    fn to_noun(&self) -> Noun {
167        (**self).to_noun()
168    }
169}
170
171impl<T: NounDecode> NounDecode for Box<T> {
172    fn from_noun(noun: &Noun) -> Option<Self> {
173        Some(Box::new(T::from_noun(noun)?))
174    }
175}
176
177impl NounEncode for Noun {
178    fn to_noun(&self) -> Noun {
179        self.clone()
180    }
181}
182
183impl NounDecode for Noun {
184    fn from_noun(noun: &Noun) -> Option<Self> {
185        Some(noun.clone())
186    }
187}
188
189impl NounEncode for () {
190    fn to_noun(&self) -> Noun {
191        atom(0)
192    }
193}
194
195impl NounDecode for () {
196    fn from_noun(noun: &Noun) -> Option<Self> {
197        if *noun == atom(0) {
198            Some(())
199        } else {
200            None
201        }
202    }
203}
204
205impl NounEncode for Belt {
206    fn to_noun(&self) -> Noun {
207        atom(self.0)
208    }
209}
210
211impl NounDecode for Belt {
212    fn from_noun(noun: &Noun) -> Option<Self> {
213        let Noun::Atom(a) = noun else {
214            return None;
215        };
216        let v = u64::try_from(a).ok()?;
217        if based_check(v) {
218            Some(Belt(v))
219        } else {
220            None
221        }
222    }
223}
224
225impl NounEncode for Digest {
226    fn to_noun(&self) -> Noun {
227        self.0.to_noun()
228    }
229}
230
231impl NounDecode for Digest {
232    fn from_noun(noun: &Noun) -> Option<Self> {
233        Some(Digest(<[Belt; 5]>::from_noun(noun)?))
234    }
235}
236
237impl NounEncode for CheetahPoint {
238    fn to_noun(&self) -> Noun {
239        (self.x.0, self.y.0, self.inf).to_noun()
240    }
241}
242
243impl NounDecode for CheetahPoint {
244    fn from_noun(noun: &Noun) -> Option<Self> {
245        let (x, y, inf) = NounDecode::from_noun(noun)?;
246        Some(Self {
247            x: F6lt(x),
248            y: F6lt(y),
249            inf,
250        })
251    }
252}
253
254macro_rules! impl_nounable_for_int {
255    ($($ty:ty),* $(,)?) => {
256        $(
257            impl NounEncode for $ty {
258                fn to_noun(&self) -> Noun {
259                    atom(*self as u64)
260                }
261            }
262
263            impl NounDecode for $ty {
264                fn from_noun(noun: &Noun) -> Option<$ty> {
265                    let Noun::Atom(a) = noun else {
266                        return None;
267                    };
268                    <$ty>::try_from(a).ok()
269                }
270            }
271        )*
272    };
273}
274
275impl_nounable_for_int!(i32, i64, isize, u32, u64, usize);
276
277impl NounEncode for bool {
278    fn to_noun(&self) -> Noun {
279        atom(if *self { 0 } else { 1 })
280    }
281}
282
283impl NounDecode for bool {
284    fn from_noun(noun: &Noun) -> Option<Self> {
285        let Noun::Atom(a) = noun else {
286            return None;
287        };
288        if a == &UBig::from(0u64) {
289            Some(true)
290        } else if a == &UBig::from(1u64) {
291            Some(false)
292        } else {
293            None
294        }
295    }
296}
297
298impl<T: NounEncode> NounEncode for Option<T> {
299    fn to_noun(&self) -> Noun {
300        match self {
301            None => atom(0),
302            Some(value) => (0, value.to_noun()).to_noun(),
303        }
304    }
305}
306
307impl<T: NounDecode> NounDecode for Option<T> {
308    fn from_noun(noun: &Noun) -> Option<Self> {
309        match noun {
310            Noun::Cell(x, v) if **x == atom(0) => Some(Some(T::from_noun(v)?)),
311            Noun::Atom(x) if x.is_zero() => Some(None),
312            _ => None,
313        }
314    }
315}
316
317impl<T: NounEncode> NounEncode for Zeroable<T> {
318    fn to_noun(&self) -> Noun {
319        match &self.0 {
320            None => atom(0),
321            Some(value) => value.to_noun(),
322        }
323    }
324}
325
326impl<T: NounDecode> NounDecode for Zeroable<T> {
327    fn from_noun(noun: &Noun) -> Option<Self> {
328        match noun {
329            Noun::Atom(x) if x.is_zero() => Some(Zeroable(None)),
330            v => Some(Zeroable(Some(T::from_noun(v)?))),
331        }
332    }
333}
334
335macro_rules! impl_nounable_for_tuple {
336    ($T0:ident => $i0:ident) => {};
337    ($T:ident => $t:ident $( $U:ident => $u:ident )+) => {
338        impl<$T: NounEncode, $($U: NounEncode),*> NounEncode for ($T, $($U),*) {
339            fn to_noun(&self) -> Noun {
340                let ($t, $($u),*) = self;
341                cons($t.to_noun(), ($($u),*).to_noun())
342            }
343        }
344
345        impl<$T: NounDecode, $($U: NounDecode),*> NounDecode for ($T, $($U),*) {
346            fn from_noun(noun: &Noun) -> Option<($T, $($U),*)> {
347                let Noun::Cell(a, b) = noun else {
348                    return None;
349                };
350                let a = <$T>::from_noun(a)?;
351                #[allow(unused_parens)]
352                let ($($u),*) = <($($U),*)>::from_noun(b)?;
353                Some((a, $($u),*))
354            }
355        }
356
357        impl_nounable_for_tuple!($($U => $u)*);
358    };
359}
360
361impl_nounable_for_tuple!(
362    A => a
363    B => b
364    C => c
365    D => d
366    E => e
367    F => f
368    G => g
369    H => h
370    I => i
371    J => j
372    K => k
373);
374
375impl<T: NounEncode, const N: usize> NounEncode for [T; N] {
376    fn to_noun(&self) -> Noun {
377        match self.split_last() {
378            None => unreachable!(),
379            Some((last, rest)) => {
380                let mut acc = last.to_noun();
381                for item in rest.iter().rev() {
382                    acc = cons(item.to_noun(), acc);
383                }
384                acc
385            }
386        }
387    }
388}
389
390impl<T: NounDecode, const N: usize> NounDecode for [T; N] {
391    fn from_noun(mut noun: &Noun) -> Option<Self> {
392        let mut ret: [Option<T>; N] = [(); N].map(|_| None);
393        for (i, item) in ret.iter_mut().enumerate() {
394            let decode = match noun {
395                Noun::Cell(a, b) => {
396                    noun = b;
397                    a
398                }
399                v if i == N - 1 => v,
400                _ => return None,
401            };
402            *item = Some(T::from_noun(decode)?);
403        }
404
405        Some(ret.map(|v| v.unwrap()))
406    }
407}
408
409// TODO: always append ~ at the end
410impl<T: NounEncode> NounEncode for &[T] {
411    fn to_noun(&self) -> Noun {
412        match self.split_last() {
413            None => atom(0),
414            Some((last, rest)) => {
415                let mut acc = last.to_noun();
416                for item in rest.iter().rev() {
417                    acc = cons(item.to_noun(), acc);
418                }
419                acc
420            }
421        }
422    }
423}
424
425impl<T: NounEncode> NounEncode for Vec<T> {
426    fn to_noun(&self) -> Noun {
427        let mut acc = atom(0);
428        for item in self.iter().rev() {
429            acc = cons(item.to_noun(), acc);
430        }
431        acc
432    }
433}
434
435impl<T: NounDecode> NounDecode for Vec<T> {
436    fn from_noun(mut noun: &Noun) -> Option<Self> {
437        let mut ret = vec![];
438        loop {
439            match noun {
440                Noun::Cell(a, b) => {
441                    ret.push(T::from_noun(a)?);
442                    noun = b;
443                }
444                Noun::Atom(v) => {
445                    if v.is_zero() {
446                        return Some(ret);
447                    } else {
448                        return None;
449                    }
450                }
451            }
452        }
453    }
454}
455
456impl NounEncode for &str {
457    fn to_noun(&self) -> Noun {
458        Noun::Atom(UBig::from_le_bytes(self.as_bytes()))
459    }
460}
461
462impl NounEncode for String {
463    fn to_noun(&self) -> Noun {
464        self.as_str().to_noun()
465    }
466}
467
468impl NounDecode for String {
469    fn from_noun(noun: &Noun) -> Option<Self> {
470        let Noun::Atom(a) = noun else {
471            return None;
472        };
473        String::from_utf8(a.to_le_bytes()).ok()
474    }
475}
476
477pub fn jam(noun: Noun) -> Vec<u8> {
478    fn met0_u64_to_usize(value: u64) -> usize {
479        (u64::BITS - value.leading_zeros()) as usize
480    }
481
482    fn met0_atom(atom: &UBig) -> usize {
483        atom.bit_len()
484    }
485
486    fn mat_backref(buffer: &mut BitVec<u8, Lsb0>, backref: usize) {
487        if backref == 0 {
488            buffer.push(true);
489            buffer.push(true);
490            buffer.push(true);
491            return;
492        }
493        let backref_sz = met0_u64_to_usize(backref as u64);
494        let backref_sz_sz = met0_u64_to_usize(backref_sz as u64);
495        buffer.push(true);
496        buffer.push(true);
497        let buffer_len = buffer.len();
498        buffer.resize(buffer_len + backref_sz_sz, false);
499        buffer.push(true);
500        let size_bits = BitSlice::<usize, Lsb0>::from_element(&backref_sz);
501        buffer.extend_from_bitslice(&size_bits[..backref_sz_sz - 1]);
502        let backref_bits = BitSlice::<usize, Lsb0>::from_element(&backref);
503        buffer.extend_from_bitslice(&backref_bits[..backref_sz]);
504    }
505
506    fn mat_atom(buffer: &mut BitVec<u8, Lsb0>, atom: &UBig) {
507        if atom.is_zero() {
508            buffer.push(false);
509            buffer.push(true);
510            return;
511        }
512        let atom_sz = met0_atom(atom);
513        let atom_sz_sz = met0_u64_to_usize(atom_sz as u64);
514        buffer.push(false);
515        let buffer_len = buffer.len();
516        buffer.resize(buffer_len + atom_sz_sz, false);
517        buffer.push(true);
518        let size_bits = BitSlice::<usize, Lsb0>::from_element(&atom_sz);
519        buffer.extend_from_bitslice(&size_bits[..atom_sz_sz - 1]);
520        let atom_bytes = atom.to_le_bytes();
521        let atom_bits = BitSlice::<u8, Lsb0>::from_slice(&atom_bytes);
522        buffer.extend_from_bitslice(&atom_bits[..atom_sz]);
523    }
524
525    fn find_backref(backrefs: &[(Noun, usize)], target: &Noun) -> Option<usize> {
526        backrefs
527            .iter()
528            .find(|(noun, _)| noun == target)
529            .map(|(_, offset)| *offset)
530    }
531
532    let mut backrefs: Vec<(Noun, usize)> = Vec::new();
533    let mut stack = Vec::new();
534    stack.push(noun);
535    let mut buffer = BitVec::<u8, Lsb0>::new();
536
537    while let Some(current) = stack.pop() {
538        if let Some(backref) = find_backref(&backrefs, &current) {
539            match &current {
540                Noun::Atom(atom) => {
541                    if met0_u64_to_usize(backref as u64) < met0_atom(atom) {
542                        mat_backref(&mut buffer, backref);
543                    } else {
544                        mat_atom(&mut buffer, atom);
545                    }
546                }
547                Noun::Cell(_, _) => {
548                    mat_backref(&mut buffer, backref);
549                }
550            }
551        } else {
552            let offset = buffer.len();
553            backrefs.push((current.clone(), offset));
554            match current {
555                Noun::Atom(atom) => {
556                    mat_atom(&mut buffer, &atom);
557                }
558                Noun::Cell(left, right) => {
559                    buffer.push(true);
560                    buffer.push(false);
561                    stack.push(*right);
562                    stack.push(*left);
563                }
564            }
565        }
566    }
567
568    buffer.into_vec()
569}
570
571pub fn cue(bytes: &[u8]) -> Option<Noun> {
572    cue_bitslice(BitSlice::from_slice(bytes))
573}
574
575pub fn cue_bitslice(buffer: &BitSlice<u8, Lsb0>) -> Option<Noun> {
576    #[derive(Copy, Clone)]
577    enum CueStackEntry {
578        DestinationPointer(*mut Noun),
579        BackRef(u64, *mut Noun),
580    }
581
582    pub fn next_up_to_n_bits<'a>(
583        cursor: &mut usize,
584        slice: &'a BitSlice<u8, Lsb0>,
585        n: usize,
586    ) -> &'a BitSlice<u8, Lsb0> {
587        let res = if (slice).len() >= *cursor + n {
588            &slice[*cursor..*cursor + n]
589        } else if slice.len() > *cursor {
590            &slice[*cursor..]
591        } else {
592            BitSlice::<u8, Lsb0>::empty()
593        };
594        *cursor += n;
595        res
596    }
597
598    pub fn rest_bits(cursor: usize, slice: &BitSlice<u8, Lsb0>) -> &BitSlice<u8, Lsb0> {
599        if slice.len() > cursor {
600            &slice[cursor..]
601        } else {
602            BitSlice::<u8, Lsb0>::empty()
603        }
604    }
605
606    fn get_size(cursor: &mut usize, buffer: &BitSlice<u8, Lsb0>) -> Option<usize> {
607        let buff_at_cursor = rest_bits(*cursor, buffer);
608        let bitsize = buff_at_cursor.first_one()?;
609        if bitsize == 0 {
610            *cursor += 1;
611            Some(0)
612        } else {
613            let mut size = [0u8; 8];
614            *cursor += bitsize + 1;
615            let size_bits = next_up_to_n_bits(cursor, buffer, bitsize - 1);
616            BitSlice::from_slice_mut(&mut size)[0..bitsize - 1].copy_from_bitslice(size_bits);
617            Some((u64::from_le_bytes(size) as usize) + (1 << (bitsize - 1)))
618        }
619    }
620
621    fn rub_backref(cursor: &mut usize, buffer: &BitSlice<u8, Lsb0>) -> Option<u64> {
622        // TODO: What's size here usually?
623        let size = get_size(cursor, buffer)?;
624        if size == 0 {
625            Some(0)
626        } else if size <= 64 {
627            // TODO: Size <= 64, so we can fit the backref in a direct atom?
628            let mut backref = [0u8; 8];
629            BitSlice::from_slice_mut(&mut backref)[0..size]
630                .copy_from_bitslice(&buffer[*cursor..*cursor + size]);
631            *cursor += size;
632            Some(u64::from_le_bytes(backref))
633        } else {
634            None
635        }
636    }
637
638    fn rub_atom(cursor: &mut usize, buffer: &BitSlice<u8, Lsb0>) -> Option<UBig> {
639        let size = get_size(cursor, buffer)?;
640        let bits = next_up_to_n_bits(cursor, buffer, size);
641        if size == 0 {
642            Some(UBig::from(0u64))
643        } else if size < 64 {
644            // Fits in a direct atom
645            let mut direct_raw = [0u8; 8];
646            BitSlice::from_slice_mut(&mut direct_raw)[0..bits.len()].copy_from_bitslice(bits);
647            Some(UBig::from(u64::from_le_bytes(direct_raw)))
648        } else {
649            // Need an indirect atom
650            let wordsize = (size + 63) >> 6;
651            let mut bytes = vec![0u8; wordsize * 8];
652            BitSlice::from_slice_mut(&mut bytes)[0..bits.len()].copy_from_bitslice(bits);
653            Some(UBig::from_le_bytes(&bytes))
654        }
655    }
656
657    pub fn next_bit(cursor: &mut usize, slice: &BitSlice<u8, Lsb0>) -> bool {
658        if (*slice).len() > *cursor {
659            let res = slice[*cursor];
660            *cursor += 1;
661            res
662        } else {
663            false
664        }
665    }
666
667    let mut backref_map = BTreeMap::<u64, *mut Noun>::new();
668    let mut result = atom(0);
669    let mut cursor = 0;
670
671    let mut cue_stack = vec![];
672
673    cue_stack.push(CueStackEntry::DestinationPointer(&mut result as *mut Noun));
674
675    while let Some(stack_entry) = cue_stack.pop() {
676        unsafe {
677            // Capture the destination pointer and pop it off the stack
678            match stack_entry {
679                CueStackEntry::DestinationPointer(dest_ptr) => {
680                    // 1 bit
681                    if next_bit(&mut cursor, buffer) {
682                        // 11 tag: backref
683                        if next_bit(&mut cursor, buffer) {
684                            let backref = rub_backref(&mut cursor, buffer)?;
685                            *dest_ptr = (**backref_map.get(&backref)?).clone();
686                        } else {
687                            // 10 tag: cell
688                            let mut head = Box::new(atom(0));
689                            let head_ptr = (&mut *head) as *mut _;
690                            let mut tail = Box::new(atom(0));
691                            let tail_ptr = (&mut *tail) as *mut _;
692                            *dest_ptr = Noun::Cell(head, tail);
693                            let backref = (cursor - 2) as u64;
694                            backref_map.insert(backref, dest_ptr);
695                            cue_stack.push(CueStackEntry::BackRef(cursor as u64 - 2, dest_ptr));
696                            cue_stack.push(CueStackEntry::DestinationPointer(tail_ptr));
697                            cue_stack.push(CueStackEntry::DestinationPointer(head_ptr));
698                        }
699                    } else {
700                        // 0 tag: atom
701                        let backref: u64 = (cursor - 1) as u64;
702                        *dest_ptr = Noun::Atom(rub_atom(&mut cursor, buffer)?);
703                        backref_map.insert(backref, dest_ptr);
704                    }
705                }
706                CueStackEntry::BackRef(backref, noun_ptr) => {
707                    backref_map.insert(backref, noun_ptr);
708                }
709            }
710        }
711    }
712
713    Some(result)
714}