Skip to main content

axiom_eth/mpt/types/
mod.rs

1use crate::{
2    rlc::types::{RlcFixedTrace, RlcTrace},
3    rlp::types::{RlpArrayWitness, RlpFieldTrace},
4};
5
6use super::*;
7
8mod input;
9pub use input::*;
10
11/// Witness for the terminal node of an MPT proof.
12/// This can be either a leaf (`ext`) or extracted from a branch (`branch`).
13/// The type is determined by `node_type`.
14#[derive(Clone, Debug)]
15pub struct TerminalWitness<F: ScalarField> {
16    pub node_type: AssignedValue<F>,
17    pub ext: LeafWitness<F>,
18    pub branch: BranchWitness<F>,
19    // pub max_leaf_bytes: usize,
20}
21
22// TODO: there is no difference structurally between `TerminalTrace` and `MPTNodeTrace` right now. Should combine somehow while still keeping the distinction between the two.
23/// The RLC traces corresponding to [`TerminalWitness`]
24pub struct TerminalTrace<F: ScalarField> {
25    pub node_type: AssignedValue<F>,
26    pub ext: LeafTrace<F>,
27    pub branch: BranchTrace<F>,
28    // pub max_leaf_bytes: usize,
29}
30
31impl<F: ScalarField> TerminalTrace<F> {
32    /// Returns the RLC of the MPT hash of the node by correcting selecting based on node type
33    pub fn mpt_hash_rlc(
34        &self,
35        ctx_gate: &mut Context<F>,
36        gate: &impl GateInstructions<F>,
37    ) -> RlcVar<F> {
38        rlc_select(
39            ctx_gate,
40            gate,
41            self.ext.rlcs.mpt_hash,
42            self.branch.rlcs.mpt_hash,
43            self.node_type,
44        )
45    }
46
47    /// Returns the RLC of the keccak of the node by correcting selecting based on node type
48    pub fn keccak_rlc(
49        &self,
50        ctx_gate: &mut Context<F>,
51        gate: &impl GateInstructions<F>,
52    ) -> AssignedValue<F> {
53        assert_eq!(self.ext.rlcs.hash.len, self.branch.rlcs.hash.len);
54        gate.select(
55            ctx_gate,
56            self.ext.rlcs.hash.rlc_val,
57            self.branch.rlcs.hash.rlc_val,
58            self.node_type,
59        )
60    }
61}
62
63#[derive(Clone, Debug)]
64pub struct LeafTrace<F: ScalarField> {
65    pub key_path: RlpFieldTrace<F>,
66    pub value: RlpFieldTrace<F>,
67    pub rlcs: MPTHashTrace<F>,
68}
69
70#[derive(Clone, Debug)]
71pub struct LeafWitness<F: ScalarField> {
72    pub rlp: RlpArrayWitness<F>,
73    pub hash_query: MPTHashWitness<F>,
74}
75
76#[derive(Clone, Debug)]
77pub struct ExtensionTrace<F: ScalarField> {
78    pub key_path: RlpFieldTrace<F>,
79    pub node_ref: RlpFieldTrace<F>,
80    pub rlcs: MPTHashTrace<F>,
81}
82
83#[derive(Clone, Debug)]
84pub struct ExtensionWitness<F: ScalarField> {
85    pub rlp: RlpArrayWitness<F>,
86    pub hash_query: MPTHashWitness<F>,
87}
88
89#[derive(Clone, Debug)]
90pub struct BranchTrace<F: ScalarField> {
91    // rlc without rlp prefix
92    pub node_refs: [RlpFieldTrace<F>; BRANCH_NUM_ITEMS],
93    pub rlcs: MPTHashTrace<F>,
94}
95
96#[derive(Clone, Debug)]
97pub struct BranchWitness<F: ScalarField> {
98    pub rlp: RlpArrayWitness<F>,
99    pub hash_query: MPTHashWitness<F>,
100}
101
102// helper types for readability
103pub type AssignedBytes<F> = Vec<AssignedValue<F>>; // TODO: use SafeByte
104pub type AssignedNibbles<F> = Vec<AssignedValue<F>>;
105
106#[derive(Clone, Debug)]
107pub struct MPTNode<F: ScalarField> {
108    pub rlp_bytes: AssignedBytes<F>,
109    /// 0 = branch, 1 = extension
110    pub node_type: AssignedValue<F>,
111}
112
113#[derive(Clone, Debug)]
114/// The `node_type` flag selects whether the node is parsed as a branch or extension node.
115pub struct MPTNodeWitness<F: ScalarField> {
116    /// 0 = branch, 1 = extension
117    pub node_type: AssignedValue<F>,
118    /// The node parsed as an extension node, or dummy extension node otherwise
119    pub ext: ExtensionWitness<F>,
120    /// The node parsed as a branch node, or dummy branch node otherwise
121    pub branch: BranchWitness<F>,
122}
123
124#[derive(Clone, Debug)]
125/// The `node_type` flag selects whether the node is parsed as a branch or extension node.
126pub struct MPTNodeTrace<F: ScalarField> {
127    /// 0 = branch, 1 = extension
128    pub node_type: AssignedValue<F>,
129    /// The node parsed as an extension node, or dummy extension node otherwise
130    pub ext: ExtensionTrace<F>,
131    /// The node parsed as a branch node, or dummy branch node otherwise
132    pub branch: BranchTrace<F>,
133}
134
135impl<F: ScalarField> MPTNodeTrace<F> {
136    /// Returns the RLC of the MPT hash of the node by correcting selecting based on node type
137    pub fn mpt_hash_rlc(
138        &self,
139        ctx_gate: &mut Context<F>,
140        gate: &impl GateInstructions<F>,
141    ) -> RlcVar<F> {
142        rlc_select(
143            ctx_gate,
144            gate,
145            self.ext.rlcs.mpt_hash,
146            self.branch.rlcs.mpt_hash,
147            self.node_type,
148        )
149    }
150
151    /// Returns the RLC of the keccak of the node by correcting selecting based on node type
152    pub fn keccak_rlc(
153        &self,
154        ctx_gate: &mut Context<F>,
155        gate: &impl GateInstructions<F>,
156    ) -> AssignedValue<F> {
157        assert_eq!(self.ext.rlcs.hash.len, self.branch.rlcs.hash.len);
158        gate.select(
159            ctx_gate,
160            self.ext.rlcs.hash.rlc_val,
161            self.branch.rlcs.hash.rlc_val,
162            self.node_type,
163        )
164    }
165}
166
167#[derive(Clone, Debug)]
168/// A fragment of the key (bytes), stored as nibbles before hex-prefix encoding
169pub struct MPTFragment<F: ScalarField> {
170    /// A variable length string of hex-numbers, resized to a fixed max length with 0s
171    pub nibbles: AssignedNibbles<F>,
172    pub is_odd: AssignedValue<F>,
173    // hex_len = 2 * byte_len + is_odd - 2
174    // if nibble for branch: byte_len = is_odd = 1
175    /// The byte length of the hex-prefix encoded fragment
176    pub byte_len: AssignedValue<F>,
177}
178
179#[derive(Clone, Debug)]
180pub struct MPTProofWitness<F: ScalarField> {
181    // we keep only the parts of the proof necessary:
182    pub value_bytes: AssignedBytes<F>,
183    pub value_byte_len: AssignedValue<F>,
184    pub root_hash_bytes: AssignedBytes<F>,
185    /// The variable length of the key
186    pub key_byte_len: Option<AssignedValue<F>>,
187    /// The variable length of the proof. This includes the leaf node in the case of an inclusion proof. There is no leaf node in the case of a non-inclusion proof.
188    pub depth: AssignedValue<F>,
189    /// The non-leaf nodes of the mpt proof, resized to `max_depth - 1`. Each node has been parsed with both a hypothetical branch and extension node. The actual type is determined by the `node_type` flag.
190    ///
191    /// The actual variable length of `nodes` is `depth - 1` if `slot_is_empty == true` (excludes leaf node), otherwise `depth`.
192    pub nodes: Vec<MPTNodeWitness<F>>,
193    /// The last node parsed
194    pub terminal_node: TerminalWitness<F>,
195
196    /// Boolean indicating whether the MPT contains a value at `key`
197    pub slot_is_empty: AssignedValue<F>,
198
199    pub max_key_byte_len: usize,
200    pub max_depth: usize,
201
202    /// The key fragments (nibbles), without encoding, provided as private inputs
203    pub key_frag: Vec<MPTFragment<F>>,
204    /// The hex-prefix encoded path for (potential) extension nodes (hex-prefix encoding has leaf vs. extension distinction).
205    /// These are derived from the nodes themselves.
206    pub key_frag_ext_bytes: Vec<AssignedBytes<F>>,
207    /// The hex-prefix encoded path for (potential) leaf nodes (hex-prefix encoding has leaf vs. extension distinction).
208    /// These are derived from the nodes themselves.
209    pub key_frag_leaf_bytes: Vec<AssignedBytes<F>>,
210    pub frag_lens: Vec<AssignedValue<F>>,
211    pub key_hexs: AssignedNibbles<F>,
212}
213
214pub type MPTHashWitness<F> = KeccakVarLenQuery<F>;
215
216#[derive(Clone, Copy, Debug)]
217pub struct MPTHashTrace<F: ScalarField> {
218    /// 32-byte keccak hash RLC
219    pub hash: RlcFixedTrace<F>,
220    /// input if len < 32, hash otherwise.
221    pub mpt_hash: RlcTrace<F>,
222}
223
224#[derive(Clone, Debug, Serialize, Deserialize, Hash)]
225pub struct PathBytes(pub Vec<u8>);
226
227impl<T: AsRef<[u8]>> From<&T> for PathBytes {
228    fn from(value: &T) -> Self {
229        Self(value.as_ref().to_vec())
230    }
231}
232
233impl From<Vec<u8>> for PathBytes {
234    fn from(value: Vec<u8>) -> Self {
235        Self(value)
236    }
237}
238
239impl From<H256> for PathBytes {
240    fn from(value: H256) -> Self {
241        Self(value.0.to_vec())
242    }
243}
244
245impl AsRef<[u8]> for PathBytes {
246    fn as_ref(&self) -> &[u8] {
247        &self.0
248    }
249}