Modify parser

evm/src/cpu/kernel/assembler.rs

@@ -7,6 +7,7 @@ use log::debug;
 use super::ast::PushTarget;
 use crate::cpu::kernel::ast::{Literal, StackReplacement};
 use crate::cpu::kernel::keccak_util::hash_kernel;
+use crate::cpu::kernel::prover_input::ProverInputFn;
 use crate::cpu::kernel::stack_manipulation::expand_stack_manipulation;
 use crate::cpu::kernel::{
     ast::{File, Item},
@@ -27,15 +28,22 @@ pub struct Kernel {
     pub(crate) code_hash: [u32; 8],
 
     pub(crate) global_labels: HashMap<String, usize>,
+
+    pub(crate) prover_inputs: HashMap<usize, ProverInputFn>,
 }
 
 impl Kernel {
-    fn new(code: Vec<u8>, global_labels: HashMap<String, usize>) -> Self {
+    fn new(
+        code: Vec<u8>,
+        global_labels: HashMap<String, usize>,
+        prover_inputs: HashMap<usize, ProverInputFn>,
+    ) -> Self {
         let code_hash = hash_kernel(&code);
         Self {
             code,
             code_hash,
             global_labels,
+            prover_inputs,
         }
     }
 }
@@ -57,6 +65,7 @@ impl Macro {
 pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Kernel {
     let macros = find_macros(&files);
     let mut global_labels = HashMap::new();
+    let mut prover_inputs = HashMap::new();
     let mut offset = 0;
     let mut expanded_files = Vec::with_capacity(files.len());
     let mut local_labels = Vec::with_capacity(files.len());
@@ -65,7 +74,12 @@ pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Ke
         let expanded_file = expand_repeats(expanded_file);
         let expanded_file = inline_constants(expanded_file, &constants);
         let expanded_file = expand_stack_manipulation(expanded_file);
-        local_labels.push(find_labels(&expanded_file, &mut offset, &mut global_labels));
+        local_labels.push(find_labels(
+            &expanded_file,
+            &mut offset,
+            &mut global_labels,
+            &mut prover_inputs,
+        ));
         expanded_files.push(expanded_file);
     }
     let mut code = vec![];
@@ -76,7 +90,7 @@ pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Ke
         debug!("Assembled file size: {} bytes", file_len);
     }
     assert_eq!(code.len(), offset, "Code length doesn't match offset.");
-    Kernel::new(code, global_labels)
+    Kernel::new(code, global_labels, prover_inputs)
 }
 
 fn find_macros(files: &[File]) -> HashMap<String, Macro> {
@@ -217,6 +231,7 @@ fn find_labels(
     body: &[Item],
     offset: &mut usize,
     global_labels: &mut HashMap<String, usize>,
+    prover_inputs: &mut HashMap<usize, ProverInputFn>,
 ) -> HashMap<String, usize> {
     // Discover the offset of each label in this file.
     let mut local_labels = HashMap::<String, usize>::new();
@@ -237,6 +252,10 @@ fn find_labels(
                 assert!(old.is_none(), "Duplicate local label: {}", label);
             }
             Item::Push(target) => *offset += 1 + push_target_size(target) as usize,
+            Item::ProverInput(prover_input_fn) => {
+                prover_inputs.insert(*offset, prover_input_fn.clone());
+                *offset += 1;
+            }
             Item::StandardOp(_) => *offset += 1,
             Item::Bytes(bytes) => *offset += bytes.len(),
         }
@@ -283,6 +302,9 @@ fn assemble_file(
                 code.push(get_push_opcode(target_bytes.len() as u8));
                 code.extend(target_bytes);
             }
+            Item::ProverInput(_) => {
+                code.push(get_opcode("PROVER_INPUT"));
+            }
             Item::StandardOp(opcode) => {
                 code.push(get_opcode(&opcode));
             }
@@ -357,7 +379,7 @@ mod tests {
         expected_global_labels.insert("function_1".to_string(), 0);
         expected_global_labels.insert("function_2".to_string(), 3);
 
-        let expected_kernel = Kernel::new(expected_code, expected_global_labels);
+        let expected_kernel = Kernel::new(expected_code, expected_global_labels, HashMap::new());
 
         let program = vec![file_1, file_2];
         assert_eq!(assemble(program, HashMap::new()), expected_kernel);

evm/src/cpu/kernel/ast.rs

@@ -1,6 +1,8 @@
 use ethereum_types::U256;
 use plonky2_util::ceil_div_usize;
 
+use crate::cpu::kernel::prover_input::ProverInputFn;
+
 #[derive(Debug)]
 pub(crate) struct File {
     pub(crate) body: Vec<Item>,
@@ -25,6 +27,8 @@ pub(crate) enum Item {
     LocalLabelDeclaration(String),
     /// A `PUSH` operation.
     Push(PushTarget),
+    /// A `ProverInput` operation.
+    ProverInput(ProverInputFn),
     /// Any opcode besides a PUSH opcode.
     StandardOp(String),
     /// Literal hex data; should contain an even number of hex chars.

evm/src/cpu/kernel/evm_asm.pest

@@ -15,7 +15,7 @@ literal = { literal_hex | literal_decimal }
 variable = ${ "$" ~ identifier }
 constant = ${ "@" ~ identifier }
 
-item = { macro_def | macro_call | repeat | stack | global_label | local_label | bytes_item | push_instruction | nullary_instruction }
+item = { macro_def | macro_call | repeat | stack | global_label | local_label | bytes_item | push_instruction | prover_input_instruction | nullary_instruction }
 macro_def = { ^"%macro" ~ identifier ~ paramlist? ~ item* ~ ^"%endmacro" }
 macro_call = ${ "%" ~ !(^"macro" | ^"endmacro" | ^"rep" | ^"endrep" | ^"stack") ~ identifier ~ macro_arglist? }
 repeat = { ^"%rep" ~ literal ~ item* ~ ^"%endrep" }
@@ -29,7 +29,7 @@ local_label = { identifier ~ ":" }
 bytes_item = { ^"BYTES " ~ literal ~ ("," ~ literal)* }
 push_instruction = { ^"PUSH " ~ push_target }
 push_target = { literal | identifier | variable | constant }
-prover_input_instruction = { ^"PROVER_INPUT " ~ "(" ~ prover_input_fn ~ ")" } // TODO: Can also support extra arguments.
+prover_input_instruction = { ^"PROVER_INPUT" ~ "(" ~ prover_input_fn ~ ")" } // TODO: Could also support extra arguments.
 prover_input_fn = { identifier ~ ("::" ~ identifier)*}
 nullary_instruction = { identifier }
 

evm/src/cpu/kernel/mod.rs

@@ -4,6 +4,7 @@ mod ast;
 pub(crate) mod keccak_util;
 mod opcodes;
 mod parser;
+mod prover_input;
 mod stack_manipulation;
 
 #[cfg(test)]

evm/src/cpu/kernel/parser.rs

@@ -33,6 +33,15 @@ fn parse_item(item: Pair<Rule>) -> Item {
         }
         Rule::bytes_item => Item::Bytes(item.into_inner().map(parse_literal).collect()),
         Rule::push_instruction => Item::Push(parse_push_target(item.into_inner().next().unwrap())),
+        Rule::prover_input_instruction => Item::ProverInput(
+            item.into_inner()
+                .next()
+                .unwrap()
+                .into_inner()
+                .map(|x| x.as_str().into())
+                .collect::<Vec<_>>()
+                .into(),
+        ),
         Rule::nullary_instruction => Item::StandardOp(item.as_str().into()),
         _ => panic!("Unexpected {:?}", item.as_rule()),
     }

evm/src/cpu/kernel/prover_input.rs

@@ -0,0 +1,8 @@
+#[derive(PartialEq, Eq, Debug, Clone)]
+pub(crate) struct ProverInputFn(Vec<String>);
+
+impl From<Vec<String>> for ProverInputFn {
+    fn from(v: Vec<String>) -> Self {
+        Self(v)
+    }
+}