Skip to main content

miden_vm/
internal.rs

1use std::{
2    collections::{BTreeMap, HashMap},
3    fs,
4    path::{Path, PathBuf},
5};
6
7use miden_assembly::diagnostics::{IntoDiagnostic, Report, WrapErr};
8use miden_core::{
9    Felt, WORD_SIZE,
10    field::{PrimeField64, QuotientMap},
11};
12use serde::Deserialize;
13pub use tracing::{Level, event, instrument};
14
15use crate::{
16    StackInputs, Word, ZERO,
17    advice::AdviceInputs,
18    crypto::merkle::{MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree, SimpleSmt},
19};
20
21// CONSTANTS
22// ================================================================================================
23
24const SIMPLE_SMT_DEPTH: u8 = u64::BITS as u8;
25
26// MERKLE DATA
27// ================================================================================================
28
29/// Struct used to deserialize merkle data from input file. Merkle data can be represented as a
30/// merkle tree or a Sparse Merkle Tree.
31#[derive(Deserialize, Debug)]
32pub enum MerkleData {
33    /// String representation of a merkle tree. The merkle tree is represented as a vector of
34    /// 32 byte hex strings where each string represents a leaf in the tree.
35    #[serde(rename = "merkle_tree")]
36    MerkleTree(Vec<String>),
37    /// String representation of a Sparse Merkle Tree. The Sparse Merkle Tree is represented as a
38    /// vector of tuples where each tuple consists of a u64 node index and a 32 byte hex string
39    /// representing the value of the node.
40    #[serde(rename = "sparse_merkle_tree")]
41    SparseMerkleTree(Vec<(u64, String)>),
42    /// String representation of a Partial Merkle Tree. The Partial Merkle Tree is represented as a
43    /// vector of tuples where each tuple consists of a leaf index tuple (depth, index) and a 32
44    /// byte hex string representing the value of the leaf.
45    #[serde(rename = "partial_merkle_tree")]
46    PartialMerkleTree(Vec<((u8, u64), String)>),
47}
48
49// INPUT FILE
50// ================================================================================================
51
52// TODO consider using final types instead of string representations.
53/// Input file struct that is used to deserialize input data from file. It consists of four
54/// components:
55/// - operand_stack
56/// - advice_stack
57/// - advice_map
58/// - merkle_store
59#[derive(Deserialize, Debug)]
60pub struct InputFile {
61    /// String representation of the initial operand stack, composed of chained field elements.
62    pub operand_stack: Vec<String>,
63    /// Optional string representation of the initial advice stack, composed of chained field
64    /// elements.
65    pub advice_stack: Option<Vec<String>>,
66    /// Optional map of 32 byte hex strings to vectors of u64s representing the initial advice map.
67    pub advice_map: Option<HashMap<String, Vec<u64>>>,
68    /// Optional vector of merkle data which will be loaded into the initial merkle store. Merkle
69    /// data is represented as 32 byte hex strings and node indexes are represented as u64s.
70    pub merkle_store: Option<Vec<MerkleData>>,
71}
72
73impl Default for InputFile {
74    fn default() -> Self {
75        Self {
76            operand_stack: Vec::new(),
77            advice_stack: Some(Vec::new()),
78            advice_map: Some(HashMap::new()),
79            merkle_store: None,
80        }
81    }
82}
83
84/// Helper methods to interact with the input file
85impl InputFile {
86    #[instrument(name = "read_input_file", skip_all)]
87    pub fn read(inputs_path: &Option<PathBuf>, program_path: &Path) -> Result<Self, Report> {
88        // if file not specified explicitly and corresponding file with same name as program_path
89        // with '.inputs' extension does't exist, set operand_stack to empty vector
90        if !inputs_path.is_some() && !program_path.with_extension("inputs").exists() {
91            return Ok(Self::default());
92        }
93
94        // If inputs_path has been provided then use this as path. Alternatively we will
95        // replace the program_path extension with .inputs and use this as a default.
96        let path = match inputs_path {
97            Some(path) => path.clone(),
98            None => program_path.with_extension("inputs"),
99        };
100
101        // read input file to string
102        let inputs_file = fs::read_to_string(&path)
103            .into_diagnostic()
104            .wrap_err_with(|| format!("Failed to open input file {}", path.display()))?;
105
106        // deserialize input data
107        let inputs: InputFile = serde_json::from_str(&inputs_file)
108            .into_diagnostic()
109            .wrap_err("Failed to deserialize input data")?;
110
111        Ok(inputs)
112    }
113
114    /// Parse advice inputs from the input file.
115    pub fn parse_advice_inputs(&self) -> Result<AdviceInputs, String> {
116        let mut advice_inputs = AdviceInputs::default();
117
118        let stack = self
119            .parse_advice_stack()
120            .map_err(|e| format!("failed to parse advice provider: {e}"))?;
121        advice_inputs = advice_inputs.with_stack_values(stack).map_err(|e| e.to_string())?;
122
123        if let Some(map) = self
124            .parse_advice_map()
125            .map_err(|e| format!("failed to parse advice provider: {e}"))?
126        {
127            advice_inputs = advice_inputs.with_map(map);
128        }
129
130        if let Some(merkle_store) = self
131            .parse_merkle_store()
132            .map_err(|e| format!("failed to parse advice provider: {e}"))?
133        {
134            advice_inputs = advice_inputs.with_merkle_store(merkle_store);
135        }
136
137        Ok(advice_inputs)
138    }
139
140    /// Parse advice stack data from the input file.
141    fn parse_advice_stack(&self) -> Result<Vec<u64>, String> {
142        self.advice_stack
143            .as_deref()
144            .unwrap_or(&[])
145            .iter()
146            .map(|v| {
147                v.parse::<u64>()
148                    .map_err(|e| format!("failed to parse advice stack value '{v}': {e}"))
149            })
150            .collect::<Result<Vec<_>, _>>()
151    }
152
153    /// Parse advice map data from the input file.
154    fn parse_advice_map(&self) -> Result<Option<BTreeMap<Word, Vec<Felt>>>, String> {
155        let advice_map = match &self.advice_map {
156            Some(advice_map) => advice_map,
157            None => return Ok(None),
158        };
159
160        let map = advice_map
161            .iter()
162            .map(|(k, v)| {
163                // Convert key to Word
164                let key = Word::try_from(k)
165                    .map_err(|e| format!("failed to decode advice map key '{k}': {e}"))?;
166
167                // convert values to Felt
168                let values = v
169                    .iter()
170                    .map(|v| {
171                        Felt::from_canonical_checked(*v).ok_or_else(|| {
172                            format!("failed to convert advice map value '{v}' to Felt")
173                        })
174                    })
175                    .collect::<Result<Vec<_>, _>>()?;
176                Ok((key, values))
177            })
178            .collect::<Result<BTreeMap<Word, Vec<Felt>>, String>>()?;
179
180        Ok(Some(map))
181    }
182
183    /// Parse merkle store data from the input file.
184    fn parse_merkle_store(&self) -> Result<Option<MerkleStore>, String> {
185        let merkle_data = match &self.merkle_store {
186            Some(merkle_data) => merkle_data,
187            None => return Ok(None),
188        };
189
190        let mut merkle_store = MerkleStore::default();
191        for data in merkle_data {
192            match data {
193                MerkleData::MerkleTree(data) => {
194                    let leaves = Self::parse_merkle_tree(data)?;
195                    let tree = MerkleTree::new(leaves)
196                        .map_err(|e| format!("failed to parse a Merkle tree: {e}"))?;
197                    merkle_store.extend(tree.inner_nodes());
198                    event!(
199                        Level::TRACE,
200                        "Added Merkle tree with root {} to the Merkle store",
201                        tree.root()
202                    );
203                },
204                MerkleData::SparseMerkleTree(data) => {
205                    let entries = Self::parse_sparse_merkle_tree(data)?;
206                    let tree = SimpleSmt::<SIMPLE_SMT_DEPTH>::with_leaves(entries)
207                        .map_err(|e| format!("failed to parse a Sparse Merkle Tree: {e}"))?;
208                    merkle_store.extend(tree.inner_nodes());
209                    event!(
210                        Level::TRACE,
211                        "Added Sparse Merkle tree with root {} to the Merkle store",
212                        tree.root()
213                    );
214                },
215                MerkleData::PartialMerkleTree(data) => {
216                    let entries = Self::parse_partial_merkle_tree(data)?;
217                    let tree = PartialMerkleTree::with_leaves(entries)
218                        .map_err(|e| format!("failed to parse a Partial Merkle Tree: {e}"))?;
219                    merkle_store.extend(tree.inner_nodes());
220                    event!(
221                        Level::TRACE,
222                        "Added Partial Merkle tree with root {} to the Merkle store",
223                        tree.root()
224                    );
225                },
226            }
227        }
228
229        Ok(Some(merkle_store))
230    }
231
232    /// Parse and return merkle tree leaves.
233    fn parse_merkle_tree(tree: &[String]) -> Result<Vec<Word>, String> {
234        tree.iter()
235            .map(|v| {
236                let leaf = Self::parse_word(v)?;
237                Ok(leaf)
238            })
239            .collect()
240    }
241
242    /// Parse and return Sparse Merkle Tree entries.
243    fn parse_sparse_merkle_tree(tree: &[(u64, String)]) -> Result<Vec<(u64, Word)>, String> {
244        tree.iter()
245            .map(|(index, v)| {
246                let leaf = Self::parse_word(v)?;
247                Ok((*index, leaf))
248            })
249            .collect()
250    }
251
252    /// Parse and return Partial Merkle Tree entries.
253    fn parse_partial_merkle_tree(
254        tree: &[((u8, u64), String)],
255    ) -> Result<Vec<(NodeIndex, Word)>, String> {
256        tree.iter()
257            .map(|((depth, index), v)| {
258                let node_index = NodeIndex::new(*depth, *index).map_err(|e| {
259                    format!(
260                        "failed to create node index with depth {depth} and index {index} - {e}"
261                    )
262                })?;
263                let leaf = Self::parse_word(v)?;
264                Ok((node_index, leaf))
265            })
266            .collect()
267    }
268
269    /// Parse a `Word` from a hex string.
270    pub fn parse_word(word_hex: &str) -> Result<Word, String> {
271        let Some(word_value) = word_hex.strip_prefix("0x") else {
272            return Err(format!("failed to decode `Word` from hex {word_hex} - missing 0x prefix"));
273        };
274        let mut word_data = [0u8; 32];
275        hex::decode_to_slice(word_value, &mut word_data)
276            .map_err(|e| format!("failed to decode `Word` from hex {word_hex} - {e}"))?;
277        let mut word = [ZERO; WORD_SIZE];
278        for (i, value) in word_data.chunks(8).enumerate() {
279            // Convert 8-byte slice to [u8; 8] array, then to u64 (little-endian), then to Felt
280            let bytes: [u8; 8] = value.try_into().map_err(|_| {
281                format!("failed to convert `Word` data {word_hex} (element {i}) - expected 8 bytes")
282            })?;
283            let value_u64 = u64::from_le_bytes(bytes);
284            word[i] = Felt::from_canonical_checked(value_u64).ok_or_else(|| {
285                format!(
286                    "failed to convert `Word` data {word_hex} (element {i}) to Felt - \
287                     value {} exceeds field modulus {}",
288                    value_u64,
289                    Felt::ORDER_U64
290                )
291            })?;
292        }
293        Ok(word.into())
294    }
295
296    /// Parse and return the stack inputs for the program.
297    pub fn parse_stack_inputs(&self) -> Result<StackInputs, String> {
298        let stack_inputs: Vec<Felt> = self
299            .operand_stack
300            .iter()
301            .map(|v| {
302                let value = v.parse::<u64>().map_err(|e| e.to_string())?;
303                Felt::from_canonical_checked(value)
304                    .ok_or_else(|| format!("failed to convert stack input value '{v}' to Felt"))
305            })
306            .collect::<Result<_, _>>()?;
307
308        StackInputs::new(&stack_inputs).map_err(|e| e.to_string())
309    }
310}
311
312// TESTS
313// ================================================================================================
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318
319    #[test]
320    fn test_merkle_data_parsing() {
321        let program_with_pmt = "
322        {
323            \"operand_stack\": [\"1\"],
324            \"merkle_store\": [
325                {
326                    \"partial_merkle_tree\": [
327                        [
328                            [2, 0],
329                            \"0x1400000000000000000000000000000000000000000000000000000000000000\"
330                        ],
331                        [
332                            [2, 1],
333                            \"0x1500000000000000000000000000000000000000000000000000000000000000\"
334                        ],
335                        [
336                            [1, 1],
337                            \"0x0b00000000000000000000000000000000000000000000000000000000000000\"
338                        ]
339                    ]
340                }
341            ]
342        }";
343        let inputs: InputFile = serde_json::from_str(program_with_pmt).unwrap();
344        let merkle_store = inputs.parse_merkle_store().unwrap();
345        assert!(merkle_store.is_some());
346
347        let program_with_smt = "
348        {
349            \"operand_stack\": [\"1\"],
350            \"merkle_store\": [
351              {
352                \"sparse_merkle_tree\": [
353                  [
354                    0,
355                    \"0x1400000000000000000000000000000000000000000000000000000000000000\"
356                  ],
357                  [
358                    1,
359                    \"0x1500000000000000000000000000000000000000000000000000000000000000\"
360                  ],
361                  [
362                    3,
363                    \"0x1700000000000000000000000000000000000000000000000000000000000000\"
364                  ]
365                ]
366              }
367            ]
368          }";
369        let inputs: InputFile = serde_json::from_str(program_with_smt).unwrap();
370        let merkle_store = inputs.parse_merkle_store().unwrap();
371        assert!(merkle_store.is_some());
372
373        let program_with_merkle_tree = "
374        {
375            \"operand_stack\": [\"1\"],
376            \"merkle_store\": [
377                {
378                    \"merkle_tree\": [
379                        \"0x1400000000000000000000000000000000000000000000000000000000000000\",
380                        \"0x1500000000000000000000000000000000000000000000000000000000000000\",
381                        \"0x1600000000000000000000000000000000000000000000000000000000000000\",
382                        \"0x1700000000000000000000000000000000000000000000000000000000000000\"
383                    ]
384                }
385            ]
386        }";
387        let inputs: InputFile = serde_json::from_str(program_with_merkle_tree).unwrap();
388        let merkle_store = inputs.parse_merkle_store().unwrap();
389        assert!(merkle_store.is_some());
390    }
391
392    #[test]
393    fn test_parse_word_missing_0x_prefix() {
394        let result = InputFile::parse_word(
395            "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef",
396        );
397        assert!(result.is_err());
398        assert!(result.unwrap_err().contains("missing 0x prefix"));
399    }
400
401    #[test]
402    fn test_parse_word_edge_cases() {
403        // Empty string
404        let result = InputFile::parse_word("");
405        assert!(result.is_err());
406        assert!(result.unwrap_err().contains("missing 0x prefix"));
407
408        // Just "0x" without hex data
409        let result = InputFile::parse_word("0x");
410        assert!(result.is_err());
411
412        // Too short hex (less than 64 chars after 0x)
413        let result = InputFile::parse_word("0x123");
414        assert!(result.is_err());
415
416        // Hex string longer than 64 characters after 0x
417        let too_long =
418            "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
419        let result = InputFile::parse_word(too_long);
420        assert!(result.is_err());
421
422        // Invalid hex characters
423        let invalid_chars = "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdeg";
424        let result = InputFile::parse_word(invalid_chars);
425        assert!(result.is_err());
426    }
427
428    #[test]
429    fn test_parse_word_valid_hex() {
430        let valid_hex = "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
431        let result = InputFile::parse_word(valid_hex);
432        assert!(result.is_ok());
433
434        // Test that the parsed word is not zero word
435        let word = result.unwrap();
436        let zero_word = Word::from([ZERO; 4]);
437        assert_ne!(word, zero_word);
438    }
439
440    #[test]
441    fn test_parse_stack_inputs_large_numbers() {
442        // Test with felt modulus - 1 (should succeed)
443        let felt_modulus_minus_one = Felt::ORDER_U64 - 1;
444        let inputs2 = InputFile {
445            operand_stack: vec![felt_modulus_minus_one.to_string()],
446            advice_stack: None,
447            advice_map: None,
448            merkle_store: None,
449        };
450        let result2 = inputs2.parse_stack_inputs();
451        assert!(result2.is_ok());
452
453        // Test with felt modulus (should fail)
454        let felt_modulus = Felt::ORDER_U64;
455        let inputs3 = InputFile {
456            operand_stack: vec![felt_modulus.to_string()],
457            advice_stack: None,
458            advice_map: None,
459            merkle_store: None,
460        };
461        let result3 = inputs3.parse_stack_inputs();
462        assert!(result3.is_err());
463    }
464
465    #[test]
466    fn test_parse_advice_stack_large_numbers() {
467        // Test with felt modulus - 1 (should succeed)
468        let felt_modulus_minus_one = Felt::ORDER_U64 - 1;
469        let inputs2 = InputFile {
470            operand_stack: Vec::new(),
471            advice_stack: Some(vec![felt_modulus_minus_one.to_string()]),
472            advice_map: None,
473            merkle_store: None,
474        };
475        let result2 = inputs2.parse_advice_inputs();
476        assert!(result2.is_ok());
477
478        // Test with felt modulus (should fail)
479        let felt_modulus = Felt::ORDER_U64;
480        let inputs = InputFile {
481            operand_stack: Vec::new(),
482            advice_stack: Some(vec![felt_modulus.to_string()]),
483            advice_map: None,
484            merkle_store: None,
485        };
486        let result = inputs.parse_advice_inputs();
487        assert!(result.is_err());
488    }
489}