openvm_circuit/system/memory/merkle/
public_values.rs

1use std::io::{self, Write};
2
3use itertools::Itertools;
4use openvm_stark_backend::{
5    codec::{DecodableConfig, EncodableConfig},
6    p3_util::log2_strict_usize,
7};
8use p3_field::Field;
9use serde::{Deserialize, Serialize};
10use thiserror::Error;
11use tracing::instrument;
12
13use crate::{
14    arch::{hasher::Hasher, MemoryCellType, ADDR_SPACE_OFFSET},
15    system::memory::{dimensions::MemoryDimensions, online::LinearMemory, MemoryImage},
16};
17
18pub const PUBLIC_VALUES_AS: u32 = 3;
19pub const PUBLIC_VALUES_ADDRESS_SPACE_OFFSET: u32 = PUBLIC_VALUES_AS - ADDR_SPACE_OFFSET;
20
21/// Merkle proof for user public values in the memory state.
22#[derive(Clone, Debug, Serialize, Deserialize)]
23#[serde(bound(
24    serialize = "F: Serialize, [F; CHUNK]: Serialize",
25    deserialize = "F: Deserialize<'de>, [F; CHUNK]: Deserialize<'de>"
26))]
27pub struct UserPublicValuesProof<const CHUNK: usize, F> {
28    /// Proof of the path from the root of public values to the memory root in the format of
29    /// sequence of sibling node hashes.
30    pub proof: Vec<[F; CHUNK]>,
31    /// Raw public values. Its length should be (a power of two) * CHUNK.
32    pub public_values: Vec<F>,
33    /// Merkle root of public values. The computation of this value follows the same logic of
34    /// `MemoryNode`. The merkle tree doesn't pad because the length `public_values` implies the
35    /// merkle tree is always a full binary tree.
36    pub public_values_commit: [F; CHUNK],
37}
38
39#[derive(Error, Debug)]
40pub enum UserPublicValuesProofError {
41    #[error("unexpected length: {0}")]
42    UnexpectedLength(usize),
43    #[error("incorrect proof length: {0} (expected {1})")]
44    IncorrectProofLength(usize, usize),
45    #[error("user public values do not match commitment")]
46    UserPublicValuesCommitMismatch,
47    #[error("final memory root mismatch")]
48    FinalMemoryRootMismatch,
49}
50
51impl<const CHUNK: usize, F: Field> UserPublicValuesProof<CHUNK, F> {
52    /// Computes the proof of the public values from the final memory state and the Merkle top
53    /// sub-tree of address space roots. This function will re-compute the empty merkle roots of
54    /// each height `0..=address_height` internally.
55    ///
56    /// Assumption:
57    /// - `num_public_values` is a power of two * CHUNK. It cannot be 0.
58    /// - `top_tree` is 0-indexed and a segment tree of length `2 * 2^addr_space_height - 1`.
59    #[instrument(name = "compute_user_public_values_proof", skip_all)]
60    pub fn compute(
61        memory_dimensions: MemoryDimensions,
62        num_public_values: usize,
63        hasher: &(impl Hasher<CHUNK, F> + Sync),
64        final_memory: &MemoryImage,
65        top_tree: &[[F; CHUNK]],
66    ) -> Self {
67        let public_values = extract_public_values(num_public_values, final_memory)
68            .iter()
69            .map(|&x| F::from_u8(x))
70            .collect_vec();
71        let public_values_commit = hasher.merkle_root(&public_values);
72        let proof = compute_merkle_proof_to_user_public_values_root(
73            memory_dimensions,
74            num_public_values,
75            hasher,
76            top_tree,
77        );
78        UserPublicValuesProof {
79            proof,
80            public_values,
81            public_values_commit,
82        }
83    }
84
85    pub fn verify(
86        &self,
87        hasher: &impl Hasher<CHUNK, F>,
88        memory_dimensions: MemoryDimensions,
89        final_memory_root: [F; CHUNK],
90    ) -> Result<(), UserPublicValuesProofError> {
91        // Verify user public values Merkle proof:
92        // 0. Get correct indices for Merkle proof based on memory dimensions
93        // 1. Verify user public values commitment with respect to the final memory root.
94        // 2. Compare user public values commitment with Merkle root of user public values.
95        let pv_commit = self.public_values_commit;
96        // 0.
97        let pv_as = PUBLIC_VALUES_AS;
98        let pv_start_idx = memory_dimensions.label_to_index((pv_as, 0));
99        let pvs = &self.public_values;
100        if !pvs.len().is_multiple_of(CHUNK) || !(pvs.len() / CHUNK).is_power_of_two() {
101            return Err(UserPublicValuesProofError::UnexpectedLength(pvs.len()));
102        }
103        let pv_height = log2_strict_usize(pvs.len() / CHUNK);
104        let proof_len = memory_dimensions.overall_height() - pv_height;
105        let idx_prefix = pv_start_idx >> pv_height;
106        // 1.
107        if self.proof.len() != proof_len {
108            return Err(UserPublicValuesProofError::IncorrectProofLength(
109                self.proof.len(),
110                proof_len,
111            ));
112        }
113        let mut curr_root = pv_commit;
114        for (i, sibling_hash) in self.proof.iter().enumerate() {
115            curr_root = if idx_prefix & (1 << i) != 0 {
116                hasher.compress(sibling_hash, &curr_root)
117            } else {
118                hasher.compress(&curr_root, sibling_hash)
119            }
120        }
121        if curr_root != final_memory_root {
122            return Err(UserPublicValuesProofError::FinalMemoryRootMismatch);
123        }
124        // 2. Compute merkle root of public values
125        if hasher.merkle_root(pvs) != pv_commit {
126            return Err(UserPublicValuesProofError::UserPublicValuesCommitMismatch);
127        }
128
129        Ok(())
130    }
131
132    pub fn encode<SC: EncodableConfig<F = F, Digest = [F; CHUNK]>, W: Write>(
133        &self,
134        writer: &mut W,
135    ) -> io::Result<()> {
136        SC::encode_digest_slice(&self.proof, writer)?;
137        SC::encode_base_field_slice(&self.public_values, writer)?;
138        SC::encode_digest(&self.public_values_commit, writer)?;
139        Ok(())
140    }
141
142    pub fn decode<SC: DecodableConfig<F = F, Digest = [F; CHUNK]>, R: io::Read>(
143        reader: &mut R,
144    ) -> io::Result<Self> {
145        let proof = SC::decode_digest_vec(reader)?;
146        let public_values = SC::decode_base_field_vec(reader)?;
147        let public_values_commit = SC::decode_digest(reader)?;
148        Ok(Self {
149            proof,
150            public_values,
151            public_values_commit,
152        })
153    }
154}
155
156fn compute_merkle_proof_to_user_public_values_root<const CHUNK: usize, F: Field>(
157    memory_dimensions: MemoryDimensions,
158    num_public_values: usize,
159    hasher: &(impl Hasher<CHUNK, F> + Sync),
160    top_tree: &[[F; CHUNK]],
161) -> Vec<[F; CHUNK]> {
162    assert_eq!(
163        num_public_values % CHUNK,
164        0,
165        "num_public_values must be a multiple of memory chunk {CHUNK}"
166    );
167    let address_height = memory_dimensions.address_height;
168    let addr_space_height = memory_dimensions.addr_space_height;
169    assert_eq!(top_tree.len(), (2 << addr_space_height) - 1);
170    let num_pv_chunks: usize = num_public_values / CHUNK;
171    // This enforces the number of public values cannot be 0.
172    assert!(
173        num_pv_chunks.is_power_of_two(),
174        "pv_height must be a power of two"
175    );
176    let pv_height = log2_strict_usize(num_pv_chunks);
177    let address_leading_zeros = address_height - pv_height;
178
179    let mut cur_node_idx = 1; // root
180    let mut proof = Vec::with_capacity(addr_space_height + address_leading_zeros);
181    let zero_nodes: Vec<_> = (0..address_height)
182        .scan(hasher.hash(&[F::ZERO; CHUNK]), |acc, _| {
183            let result = Some(*acc);
184            *acc = hasher.compress(acc, acc);
185            result
186        })
187        .collect();
188    for i in 0..addr_space_height {
189        let bit = 1 << (memory_dimensions.addr_space_height - i - 1);
190        // Recall: top_tree is 0-indexed, but cur_node_idx is 1-indexed
191        if (PUBLIC_VALUES_AS - ADDR_SPACE_OFFSET) & bit != 0 {
192            proof.push(top_tree[cur_node_idx * 2 - 1]);
193            cur_node_idx = cur_node_idx * 2 + 1;
194        } else {
195            proof.push(top_tree[cur_node_idx * 2]);
196            cur_node_idx *= 2;
197        }
198    }
199    for i in 0..address_leading_zeros {
200        // node is always on the left, the sibling is always zero node hash
201        proof.push(zero_nodes[address_height - 1 - i]);
202    }
203    proof.reverse();
204    proof
205}
206
207pub fn extract_public_values(num_public_values: usize, final_memory: &MemoryImage) -> Vec<u8> {
208    let mut public_values: Vec<u8> = {
209        assert_eq!(
210            final_memory.config[PUBLIC_VALUES_AS as usize].layout,
211            MemoryCellType::U8
212        );
213        final_memory.mem[PUBLIC_VALUES_AS as usize]
214            .as_slice()
215            .to_vec()
216    };
217
218    assert!(
219        public_values.len() >= num_public_values,
220        "Public values address space has {} elements, but configuration has num_public_values={}",
221        public_values.len(),
222        num_public_values
223    );
224    public_values.truncate(num_public_values);
225    public_values
226}
227
228#[cfg(test)]
229mod tests {
230    use openvm_stark_backend::p3_field::PrimeCharacteristicRing;
231    use openvm_stark_sdk::p3_baby_bear::BabyBear;
232
233    use super::UserPublicValuesProof;
234    use crate::{
235        arch::{hasher::poseidon2::vm_poseidon2_hasher, MemoryConfig, SystemConfig},
236        system::memory::{
237            merkle::{public_values::PUBLIC_VALUES_AS, tree::MerkleTree},
238            online::GuestMemory,
239            AddressMap, CHUNK,
240        },
241    };
242
243    type F = BabyBear;
244    #[test]
245    fn test_public_value_happy_path() {
246        let mut vm_config = SystemConfig::default();
247        let addr_space_height = 4;
248        vm_config.memory_config.addr_space_height = addr_space_height;
249        vm_config.memory_config.pointer_max_bits = 5;
250        let memory_dimensions = vm_config.memory_config.memory_dimensions();
251        let num_public_values = 16;
252        let mut addr_spaces_config = MemoryConfig::empty_address_space_configs(4);
253        addr_spaces_config[PUBLIC_VALUES_AS as usize].num_cells = num_public_values;
254        let mut memory = GuestMemory {
255            memory: AddressMap::new(addr_spaces_config),
256        };
257        unsafe {
258            memory.write::<u8, 4>(PUBLIC_VALUES_AS, 12, [0, 0, 0, 1]);
259        }
260        let mut expected_pvs = F::zero_vec(num_public_values);
261        expected_pvs[15] = F::ONE;
262
263        let hasher = vm_poseidon2_hasher();
264        let tree = MerkleTree::from_memory(&memory.memory, &memory_dimensions, &hasher);
265        let top_tree = tree.top_tree(addr_space_height);
266        let pv_proof = UserPublicValuesProof::<{ CHUNK }, F>::compute(
267            memory_dimensions,
268            num_public_values,
269            &hasher,
270            &memory.memory,
271            &top_tree,
272        );
273        assert_eq!(pv_proof.public_values, expected_pvs);
274        let final_memory_root =
275            MerkleTree::from_memory(&memory.memory, &memory_dimensions, &hasher).root();
276        pv_proof
277            .verify(&hasher, memory_dimensions, final_memory_root)
278            .unwrap();
279    }
280
281    #[test]
282    #[should_panic]
283    fn test_public_values_write_beyond_num_public_values_is_rejected() {
284        let num_public_values = 16;
285        let mut addr_spaces_config = MemoryConfig::empty_address_space_configs(4);
286        addr_spaces_config[PUBLIC_VALUES_AS as usize].num_cells = num_public_values;
287        let mut memory = GuestMemory {
288            memory: AddressMap::new(addr_spaces_config),
289        };
290        unsafe {
291            memory.write::<u8, 4>(PUBLIC_VALUES_AS, num_public_values as u32 + 4, [0, 0, 0, 1]);
292        }
293    }
294}
openvm_circuit/system/memory/merkle/public_values.rs

openvm_circuit/system/memory/merkle/
public_values.rs
