iris_ztd/
hash.rs

1use alloc::{fmt, string::String, vec, vec::Vec};
2use ibig::ops::DivRem;
3use ibig::UBig;
4use serde::{Deserialize, Serialize};
5
6use crate::{
7    belt::{Belt, PRIME},
8    tip5::hash::{hash_fixed, hash_varlen},
9    Noun, Zeroable,
10};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
13pub struct Base58Belts<const N: usize>(pub [Belt; N]);
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
16pub struct Digest(pub [Belt; 5]);
17
18impl From<[u64; 5]> for Digest {
19    fn from(belts: [u64; 5]) -> Self {
20        Digest(belts.map(Belt))
21    }
22}
23
24impl From<Digest> for Base58Belts<5> {
25    fn from(digest: Digest) -> Self {
26        Base58Belts(digest.0)
27    }
28}
29
30impl From<Base58Belts<5>> for Digest {
31    fn from(belts: Base58Belts<5>) -> Self {
32        Digest(belts.0)
33    }
34}
35
36impl<const N: usize> From<[u64; N]> for Base58Belts<N> {
37    fn from(belts: [u64; N]) -> Self {
38        Base58Belts(belts.map(Belt))
39    }
40}
41
42impl<const N: usize> Base58Belts<N> {
43    pub fn to_atom(&self) -> UBig {
44        let p = UBig::from(PRIME);
45        let mut result = UBig::from(0u64);
46        let mut power = UBig::from(1u64);
47
48        for belt in &self.0 {
49            result += UBig::from(belt.0) * &power;
50            power *= &p;
51        }
52
53        result
54    }
55
56    pub fn to_bytes(&self) -> Vec<u8> {
57        let res = self.to_atom();
58        let res_bytes = res.to_be_bytes();
59        let expected_len = N * 8; // Each belt is 8 bytes
60
61        let mut bytes = vec![0u8; expected_len];
62        let start_offset = expected_len.saturating_sub(res_bytes.len());
63        bytes[start_offset..].copy_from_slice(&res_bytes);
64        bytes
65    }
66
67    pub fn from_bytes(bytes: &[u8]) -> Self {
68        let p = UBig::from(PRIME);
69        let num = UBig::from_be_bytes(bytes);
70
71        let mut belts = Vec::with_capacity(N);
72        let mut remainder = num;
73
74        for _ in 0..N {
75            let (quotient, rem) = remainder.div_rem(&p);
76            belts.push(Belt(rem.try_into().unwrap()));
77            remainder = quotient;
78        }
79
80        // Convert Vec to array
81        let array: [Belt; N] = belts
82            .try_into()
83            .unwrap_or_else(|_| panic!("Invalid belt count"));
84        Base58Belts(array)
85    }
86}
87
88// Digest-specific implementations that delegate to Base58Belts<5>
89impl Digest {
90    pub fn to_atom(&self) -> UBig {
91        Base58Belts::<5>::from(*self).to_atom()
92    }
93
94    pub fn to_bytes(&self) -> [u8; 40] {
95        let vec = Base58Belts::<5>::from(*self).to_bytes();
96        let mut arr = [0u8; 40];
97        arr.copy_from_slice(&vec);
98        arr
99    }
100
101    pub fn from_bytes(bytes: &[u8]) -> Self {
102        Base58Belts::<5>::from_bytes(bytes).into()
103    }
104}
105
106// Display and TryFrom implementations for Base58Belts<N>
107impl<const N: usize> fmt::Display for Base58Belts<N> {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        let bytes = self.to_bytes();
110        write!(f, "{}", bs58::encode(bytes).into_string())
111    }
112}
113
114impl<const N: usize> TryFrom<&str> for Base58Belts<N> {
115    type Error = &'static str;
116
117    fn try_from(s: &str) -> Result<Self, Self::Error> {
118        Ok(Base58Belts::from_bytes(
119            &bs58::decode(s)
120                .into_vec()
121                .map_err(|_| "unable to decode base58 belts")?,
122        ))
123    }
124}
125
126// Digest implementations delegate to Base58Belts<5>
127impl fmt::Display for Digest {
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        Base58Belts::<5>::from(*self).fmt(f)
130    }
131}
132
133impl TryFrom<&str> for Digest {
134    type Error = &'static str;
135
136    fn try_from(s: &str) -> Result<Self, Self::Error> {
137        Ok(Base58Belts::<5>::try_from(s)?.into())
138    }
139}
140
141pub fn hash_noun(leaves: &[Belt], dyck: &[Belt]) -> Digest {
142    let mut combined = Vec::with_capacity(1 + leaves.len() + dyck.len());
143    combined.push(Belt(leaves.len() as u64));
144    combined.extend_from_slice(leaves);
145    combined.extend_from_slice(dyck);
146    Digest(hash_varlen(&mut combined).map(Belt))
147}
148
149pub trait Hashable {
150    fn hash(&self) -> Digest;
151}
152
153impl Hashable for Belt {
154    fn hash(&self) -> Digest {
155        hash_noun(&[*self], &[])
156    }
157}
158
159impl Hashable for u64 {
160    fn hash(&self) -> Digest {
161        Belt(*self).hash()
162    }
163}
164
165impl Hashable for usize {
166    fn hash(&self) -> Digest {
167        (*self as u64).hash()
168    }
169}
170
171impl Hashable for i32 {
172    fn hash(&self) -> Digest {
173        (*self as u64).hash()
174    }
175}
176
177impl Hashable for bool {
178    fn hash(&self) -> Digest {
179        (if *self { 0 } else { 1 }).hash()
180    }
181}
182
183impl Hashable for Digest {
184    fn hash(&self) -> Digest {
185        *self
186    }
187}
188
189impl<T: Hashable> Hashable for &T {
190    fn hash(&self) -> Digest {
191        (**self).hash()
192    }
193}
194
195impl<T: Hashable> Hashable for Option<T> {
196    fn hash(&self) -> Digest {
197        match self {
198            None => 0.hash(),
199            Some(v) => (&0, v).hash(),
200        }
201    }
202}
203
204impl<T: Hashable> Hashable for Zeroable<T> {
205    fn hash(&self) -> Digest {
206        match &self.0 {
207            None => 0.hash(),
208            Some(v) => v.hash(),
209        }
210    }
211}
212
213impl Hashable for () {
214    fn hash(&self) -> Digest {
215        0.hash()
216    }
217}
218
219macro_rules! impl_hashable_for_tuple {
220    ($T0:ident) => {};
221    ($T0:ident, $T1:ident) => {
222        impl<$T0: Hashable, $T1: Hashable> Hashable for ($T0, $T1) {
223            fn hash(&self) -> Digest {
224                let mut belts = Vec::<Belt>::with_capacity(10);
225                belts.extend_from_slice(&self.0.hash().0);
226                belts.extend_from_slice(&self.1.hash().0);
227                Digest(hash_fixed(&mut belts).map(|u| Belt(u)))
228            }
229        }
230    };
231    ($T:ident, $($U:ident),+) => {
232        impl<$T: Hashable, $($U: Hashable),*> Hashable for ($T, $($U),*) {
233            fn hash(&self) -> Digest {
234                #[allow(non_snake_case)]
235                let ($T, $($U),*) = self;
236                ($T, ($($U,)*)).hash()
237            }
238        }
239
240        impl_hashable_for_tuple!($($U),*);
241    };
242}
243
244impl_hashable_for_tuple!(A, B, C, D, E, F, G, H, I, J, K);
245
246impl<T: Hashable> Hashable for &[T] {
247    fn hash(&self) -> Digest {
248        let (first, rest) = self.split_first().unwrap();
249        if rest.is_empty() {
250            first.hash()
251        } else {
252            (first.hash(), rest.hash()).hash()
253        }
254    }
255}
256
257impl<T: Hashable> Hashable for Vec<T> {
258    fn hash(&self) -> Digest {
259        fn hash_slice<T: Hashable>(arr: &[T]) -> Digest {
260            match arr.split_first() {
261                None => 0.hash(),
262                Some((first, rest)) => (first.hash(), hash_slice(rest)).hash(),
263            }
264        }
265        hash_slice(self.as_slice())
266    }
267}
268
269impl Hashable for &str {
270    fn hash(&self) -> Digest {
271        self.bytes()
272            .enumerate()
273            .fold(0u64, |acc, (i, byte)| acc | ((byte as u64) << (i * 8)))
274            .hash()
275    }
276}
277
278impl Hashable for String {
279    fn hash(&self) -> Digest {
280        self.bytes()
281            .enumerate()
282            .fold(0u64, |acc, (i, byte)| acc | ((byte as u64) << (i * 8)))
283            .hash()
284    }
285}
286
287impl Hashable for Noun {
288    fn hash(&self) -> Digest {
289        fn visit(noun: &Noun, leaves: &mut Vec<Belt>, dyck: &mut Vec<Belt>) {
290            match noun {
291                Noun::Atom(b) => leaves.push(Belt(b.try_into().expect("atom too large"))),
292                Noun::Cell(left, right) => {
293                    dyck.push(Belt(0));
294                    visit(left, leaves, dyck);
295                    dyck.push(Belt(1));
296                    visit(right, leaves, dyck);
297                }
298            }
299        }
300
301        let mut leaves = Vec::new();
302        let mut dyck = Vec::new();
303        visit(self, &mut leaves, &mut dyck);
304        hash_noun(&leaves, &dyck)
305    }
306}