stwo_gpu/core/vcs/
poseidon252_merkle.rs

1#![allow(unused)]
2use num_traits::Zero;
3use serde::{Deserialize, Serialize};
4use starknet_crypto::{poseidon_hash, poseidon_hash_many};
5use starknet_ff::FieldElement as FieldElement252;
6use std_shims::Vec;
7
8use crate::core::channel::{MerkleChannel, Poseidon252Channel};
9use crate::core::fields::m31::{BaseField, M31};
10use crate::core::vcs::hash::Hash;
11use crate::core::vcs::utils::add_length_padding;
12use crate::core::vcs::MerkleHasher;
13
14pub const ELEMENTS_IN_BLOCK: usize = 8;
15
16#[derive(Copy, Clone, Debug, PartialEq, Eq, Default, Deserialize, Serialize)]
17pub struct Poseidon252MerkleHasher;
18impl MerkleHasher for Poseidon252MerkleHasher {
19    type Hash = FieldElement252;
20
21    fn hash_node(
22        children_hashes: Option<(Self::Hash, Self::Hash)>,
23        column_values: &[BaseField],
24    ) -> Self::Hash {
25        let n_column_blocks = column_values.len().div_ceil(ELEMENTS_IN_BLOCK);
26        let values_len = 2 + n_column_blocks;
27        let mut values = Vec::with_capacity(values_len);
28
29        if let Some((left, right)) = children_hashes {
30            if column_values.is_empty() {
31                return poseidon_hash(left, right);
32            }
33            values.push(left);
34            values.push(right);
35        }
36
37        for chunk in column_values.chunks(ELEMENTS_IN_BLOCK) {
38            values.push(construct_felt252_from_m31s(chunk));
39        }
40        poseidon_hash_many(&values)
41    }
42}
43
44// Performs felt252 = felt252 << 31 + limb.
45const fn append_m31(felt: &mut [u128; 2], limb: M31) {
46    // Felt = Felt << 31 + limb.
47    *felt = [
48        felt[0] << 31 | limb.0 as u128,
49        felt[1] << 31 | felt[0] >> (128 - 31),
50    ];
51}
52
53/// Constructs a felt252 from a slice of m31s.
54/// The maximum assumed word size is 8 limbs, which is usually the case except for the remainder.
55pub fn construct_felt252_from_m31s(word: &[M31]) -> FieldElement252 {
56    let mut felt_as_u256 = [0u128; 2];
57    for limb in word {
58        append_m31(&mut felt_as_u256, *limb);
59    }
60
61    let felt_bytes = [felt_as_u256[1].to_be_bytes(), felt_as_u256[0].to_be_bytes()];
62    let felt_bytes = unsafe { core::mem::transmute::<[[u8; 16]; 2], [u8; 32]>(felt_bytes) };
63    let mut felt = FieldElement252::from_bytes_be(&felt_bytes).unwrap();
64
65    // If this is the remainder, store its length in bits 248, 249 and 250.
66    // Note, you can also look at these 3 bits as word length modulo 8.
67    if word.len() < ELEMENTS_IN_BLOCK {
68        add_length_padding(&mut felt, word.len());
69    }
70    felt
71}
72
73impl Hash for FieldElement252 {}
74
75#[derive(Default)]
76pub struct Poseidon252MerkleChannel;
77
78#[cfg(all(test, feature = "prover"))]
79mod tests {
80    use std::time::Instant;
81
82    use itertools::Itertools;
83    use num_traits::Zero;
84    use rand::rngs::SmallRng;
85    use rand::{Rng, SeedableRng};
86    use starknet_ff::FieldElement as FieldElement252;
87
88    use crate::core::fields::m31::{BaseField, M31};
89    use crate::core::vcs::poseidon252_merkle::{
90        construct_felt252_from_m31s, Poseidon252MerkleHasher, ELEMENTS_IN_BLOCK,
91    };
92    use crate::core::vcs::test_utils::prepare_merkle;
93    use crate::core::vcs::utils::add_length_padding;
94    use crate::core::vcs::verifier::MerkleVerificationError;
95    use crate::core::vcs::MerkleHasher;
96    use crate::m31;
97
98    #[test]
99    fn test_vector() {
100        assert_eq!(
101            Poseidon252MerkleHasher::hash_node(None, &[m31!(0), m31!(1)]),
102            FieldElement252::from_dec_str(
103                "2978883932528585652864046122079599882777358126302490183268546077323303473078"
104            )
105            .unwrap()
106        );
107
108        assert_eq!(
109            Poseidon252MerkleHasher::hash_node(
110                Some((FieldElement252::from(1u32), FieldElement252::from(2u32))),
111                &[m31!(3)]
112            ),
113            FieldElement252::from_dec_str(
114                "3286095315900630438551061262740794783852190427874264245042874292062185873630"
115            )
116            .unwrap()
117        );
118    }
119
120    #[test]
121    fn test_merkle_success() {
122        let (queries, decommitment, values, verifier) = prepare_merkle::<Poseidon252MerkleHasher>();
123        verifier.verify(&queries, values, decommitment).unwrap();
124    }
125
126    #[test]
127    fn test_merkle_invalid_witness() {
128        let (queries, mut decommitment, values, verifier) =
129            prepare_merkle::<Poseidon252MerkleHasher>();
130        decommitment.hash_witness[4] = FieldElement252::default();
131
132        assert_eq!(
133            verifier.verify(&queries, values, decommitment).unwrap_err(),
134            MerkleVerificationError::RootMismatch
135        );
136    }
137
138    #[test]
139    fn test_merkle_invalid_value() {
140        let (queries, decommitment, mut values, verifier) =
141            prepare_merkle::<Poseidon252MerkleHasher>();
142        values[6] = BaseField::zero();
143
144        assert_eq!(
145            verifier.verify(&queries, values, decommitment).unwrap_err(),
146            MerkleVerificationError::RootMismatch
147        );
148    }
149
150    #[test]
151    fn test_merkle_witness_too_short() {
152        let (queries, mut decommitment, values, verifier) =
153            prepare_merkle::<Poseidon252MerkleHasher>();
154        decommitment.hash_witness.pop();
155
156        assert_eq!(
157            verifier.verify(&queries, values, decommitment).unwrap_err(),
158            MerkleVerificationError::WitnessTooShort
159        );
160    }
161
162    #[test]
163    fn test_merkle_witness_too_long() {
164        let (queries, mut decommitment, values, verifier) =
165            prepare_merkle::<Poseidon252MerkleHasher>();
166        decommitment.hash_witness.push(FieldElement252::default());
167
168        assert_eq!(
169            verifier.verify(&queries, values, decommitment).unwrap_err(),
170            MerkleVerificationError::WitnessTooLong
171        );
172    }
173
174    #[test]
175    fn test_merkle_values_too_long() {
176        let (queries, decommitment, mut values, verifier) =
177            prepare_merkle::<Poseidon252MerkleHasher>();
178        values.insert(3, BaseField::zero());
179
180        assert_eq!(
181            verifier.verify(&queries, values, decommitment).unwrap_err(),
182            MerkleVerificationError::TooManyQueriedValues
183        );
184    }
185
186    #[test]
187    fn test_merkle_values_too_short() {
188        let (queries, decommitment, mut values, verifier) =
189            prepare_merkle::<Poseidon252MerkleHasher>();
190        values.remove(3);
191
192        assert_eq!(
193            verifier.verify(&queries, values, decommitment).unwrap_err(),
194            MerkleVerificationError::TooFewQueriedValues
195        );
196    }
197
198    #[test]
199    fn test_construct_word() {
200        let mut rng = SmallRng::seed_from_u64(1638);
201        let random_values = (0..8 * 1000 + 5).map(|_| rng.gen::<M31>()).collect_vec();
202
203        let expected = random_values
204            .chunks(ELEMENTS_IN_BLOCK)
205            .map(|word| {
206                let mut felt = FieldElement252::default();
207                for x in word {
208                    felt = felt * FieldElement252::from(2u64.pow(31)) + FieldElement252::from(x.0);
209                }
210                if word.len() < ELEMENTS_IN_BLOCK {
211                    // felt = felt + word.len() << 248;
212                    add_length_padding(&mut felt, word.len());
213                }
214                felt
215            })
216            .collect_vec();
217
218        let result = random_values
219            .chunks(ELEMENTS_IN_BLOCK)
220            .map(construct_felt252_from_m31s)
221            .collect_vec();
222
223        assert_eq!(expected, result);
224    }
225}
stwo_gpu/core/vcs/poseidon252_merkle.rs

stwo_gpu/core/vcs/
poseidon252_merkle.rs
