1extern crate alloc;
16
17use alloc::vec::Vec;
18use core::{borrow::Borrow, mem::size_of};
19
20use risc0_zkp::core::{
21 digest::{Digest, DIGEST_BYTES},
22 hash::sha::Sha256,
23};
24
25pub trait Digestible {
27 fn digest<S: Sha256>(&self) -> Digest;
29}
30
31impl Digestible for [u8] {
32 fn digest<S: Sha256>(&self) -> Digest {
33 *S::hash_bytes(self)
34 }
35}
36
37impl Digestible for Vec<u8> {
38 fn digest<S: Sha256>(&self) -> Digest {
39 *S::hash_bytes(self)
40 }
41}
42
43impl<D: Digestible> Digestible for [D] {
44 fn digest<S: Sha256>(&self) -> Digest {
51 self.iter().rfold(Digest::ZERO, |accum, item| {
52 *S::hash_bytes(&[accum.as_bytes(), item.digest::<S>().as_bytes()].concat())
53 })
54 }
55}
56
57impl<T: Digestible> Digestible for Option<T> {
58 fn digest<S: Sha256>(&self) -> Digest {
59 match self {
60 Some(val) => val.digest::<S>(),
61 None => Digest::ZERO,
62 }
63 }
64}
65
66pub fn tagged_struct<S: Sha256>(tag: &str, down: &[impl Borrow<Digest>], data: &[u32]) -> Digest {
70 let tag_digest: Digest = *S::hash_bytes(tag.as_bytes());
71 #[allow(clippy::manual_slice_size_calculation)]
72 let mut all = Vec::<u8>::with_capacity(
73 DIGEST_BYTES * (down.len() + 1) + size_of::<u32>() * data.len() + size_of::<u16>(),
74 );
75 all.extend_from_slice(tag_digest.as_bytes());
76 for digest in down {
77 all.extend_from_slice(digest.borrow().as_ref());
78 }
79 for word in data.iter().copied() {
80 all.extend_from_slice(&word.to_le_bytes());
81 }
82 let down_count: u16 = down
83 .len()
84 .try_into()
85 .expect("struct defined with more than 2^16 fields");
86 all.extend_from_slice(&down_count.to_le_bytes());
87 *S::hash_bytes(&all)
88}
89
90pub fn tagged_iter<S: Sha256>(
95 tag: &str,
96 iter: impl DoubleEndedIterator<Item = impl Borrow<Digest>>,
97) -> Digest {
98 iter.rfold(Digest::ZERO, |list_digest, elem| {
99 tagged_list_cons::<S>(tag, elem.borrow(), &list_digest)
100 })
101}
102
103pub fn tagged_list<S: Sha256>(tag: &str, list: &[impl Borrow<Digest>]) -> Digest {
108 tagged_iter::<S>(tag, list.iter().map(|x| x.borrow()))
109}
110
111pub fn tagged_list_cons<S: Sha256>(tag: &str, head: &Digest, tail: &Digest) -> Digest {
132 tagged_struct::<S>(tag, &[head, tail], &[])
133}
134
135#[cfg(test)]
136mod tests {
137 use risc0_zkp::core::hash::sha::cpu;
138
139 use super::{tagged_struct, Digest};
140
141 #[test]
142 fn test_tagged_struct() {
143 let digest1 = tagged_struct::<cpu::Impl>("foo", &Vec::<Digest>::new(), &[1, 2013265920, 3]);
144 let digest2 = tagged_struct::<cpu::Impl>("bar", &[digest1, digest1], &[2013265920, 5]);
145 let digest3 = tagged_struct::<cpu::Impl>(
146 "baz",
147 &[digest1, digest2, digest1],
148 &[6, 7, 2013265920, 9, 10],
149 );
150
151 println!("digest = {:?}", digest3);
152 assert_eq!(
153 digest3.to_string(),
154 "9ff20cc6d365efa2af09181772f49013d05cdee6da896851614cae23aa5dd442"
155 );
156 }
157}