feat: add conditional feature support in the kernel assembly

evm_arithmetization/src/cpu/kernel/aggregator.rs

@@ -1,5 +1,7 @@
 //! Loads each kernel assembly file and concatenates them.
 
+use std::collections::HashSet;
+
 use itertools::Itertools;
 use once_cell::sync::Lazy;
 
@@ -173,7 +175,7 @@ pub static KERNEL_FILES: [&str; NUMBER_KERNEL_FILES] = [
 pub static KERNEL: Lazy<Kernel> = Lazy::new(combined_kernel);
 
 pub(crate) fn combined_kernel_from_files<const N: usize>(files: [&str; N]) -> Kernel {
-    let parsed_files = files.iter().map(|f| parse(f)).collect_vec();
+    let parsed_files = files.iter().map(|f| parse(f, HashSet::new())).collect_vec();
     assemble(parsed_files, evm_constants(), true)
 }
 

evm_arithmetization/src/cpu/kernel/assembler.rs

@@ -129,6 +129,7 @@ pub(crate) fn assemble(
     for file in files {
         let start = Instant::now();
         let mut file = file.body;
+        file = expand_conditional_blocks(file);
         file = expand_macros(file, &macros, &mut macro_counter);
         file = inline_constants(file, &constants);
         file = expand_stack_manipulation(file);
@@ -159,22 +160,44 @@ pub(crate) fn assemble(
 fn find_macros(files: &[File]) -> HashMap<MacroSignature, Macro> {
     let mut macros = HashMap::new();
     for file in files {
-        for item in &file.body {
-            if let Item::MacroDef(name, params, items) = item {
-                let signature = MacroSignature {
-                    name: name.clone(),
-                    num_params: params.len(),
-                };
-                let macro_ = Macro {
-                    params: params.clone(),
-                    items: items.clone(),
-                };
-                let old = macros.insert(signature.clone(), macro_);
-                assert!(old.is_none(), "Duplicate macro signature: {signature:?}");
+        find_macros_internal(&file.body, &mut macros);
+    }
+    macros
+}
+
+fn find_macros_internal(items: &[Item], macros: &mut HashMap<MacroSignature, Macro>) {
+    for item in items {
+        if let Item::ConditionalBlock(_, local_items) = item {
+            find_macros_internal(local_items, macros);
+        }
+        if let Item::MacroDef(name, params, local_items) = item {
+            let signature = MacroSignature {
+                name: name.clone(),
+                num_params: params.len(),
+            };
+            let macro_ = Macro {
+                params: params.clone(),
+                items: local_items.clone(),
+            };
+            let old = macros.insert(signature.clone(), macro_);
+            assert!(old.is_none(), "Duplicate macro signature: {signature:?}");
+        }
+    }
+}
+
+fn expand_conditional_blocks(body: Vec<Item>) -> Vec<Item> {
+    let mut expanded = vec![];
+    for item in body {
+        match item {
+            Item::ConditionalBlock(_, items) => {
+                expanded.extend(items);
+            }
+            _ => {
+                expanded.push(item);
             }
         }
     }
-    macros
+    expanded
 }
 
 fn expand_macros(
@@ -325,7 +348,8 @@ fn find_labels(
     let mut local_labels = HashMap::<String, usize>::new();
     for item in body {
         match item {
-            Item::MacroDef(_, _, _)
+            Item::ConditionalBlock(_, _)
+            | Item::MacroDef(_, _, _)
             | Item::MacroCall(_, _)
             | Item::Repeat(_, _)
             | Item::StackManipulation(_, _)
@@ -379,7 +403,8 @@ fn assemble_file(
     // Assemble the file.
     for item in body {
         match item {
-            Item::MacroDef(_, _, _)
+            Item::ConditionalBlock(_, _)
+            | Item::MacroDef(_, _, _)
             | Item::MacroCall(_, _)
             | Item::Repeat(_, _)
             | Item::StackManipulation(_, _)
@@ -437,6 +462,8 @@ fn push_target_size(target: &PushTarget) -> u8 {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashSet;
+
     use super::*;
     use crate::cpu::kernel::parser::parse;
 
@@ -734,7 +761,7 @@ mod tests {
         constants: HashMap<String, U256>,
         optimize: bool,
     ) -> Kernel {
-        let parsed_files = files.iter().map(|f| parse(f)).collect_vec();
+        let parsed_files = files.iter().map(|f| parse(f, HashSet::new())).collect_vec();
         assemble(parsed_files, constants, optimize)
     }
 }

evm_arithmetization/src/cpu/kernel/ast.rs

@@ -9,6 +9,8 @@ pub(crate) struct File {
 
 #[derive(Eq, PartialEq, Clone, Debug)]
 pub(crate) enum Item {
+    /// Defines a new macro: name, params, body.
+    ConditionalBlock(String, Vec<Item>),
     /// Defines a new macro: name, params, body.
     MacroDef(String, Vec<String>, Vec<Item>),
     /// Calls a macro: name, args.

evm_arithmetization/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_decl | local_label_decl | macro_label_decl | bytes_item | jumptable_item | push_instruction | prover_input_instruction | nullary_instruction }
+item = { conditional_block | macro_def | macro_call | repeat | stack | global_label_decl | local_label_decl | macro_label_decl | bytes_item | jumptable_item | push_instruction | prover_input_instruction | nullary_instruction }
 macro_def = { ^"%macro" ~ identifier ~ paramlist? ~ item* ~ ^"%endmacro" }
 macro_call = ${ "%" ~ !((^"macro" | ^"endmacro" | ^"rep" | ^"endrep" | ^"stack") ~ !identifier_char) ~ identifier ~ macro_arglist? }
 repeat = { ^"%rep" ~ literal ~ item* ~ ^"%endrep" }
@@ -43,5 +43,7 @@ prover_input_instruction = { ^"PROVER_INPUT" ~ "(" ~ prover_input_fn ~ ")" }
 prover_input_fn = { identifier ~ ("::" ~ identifier)*}
 nullary_instruction = { identifier }
 
+conditional_block = { ^"#" ~ "[" ~ "cfg" ~ "(" ~ "feature" ~ "=" ~ identifier ~ ")" ~ "]" ~ "{" ~ item* ~ ^"}"}
+
 file = { SOI ~ item* ~ silent_eoi }
 silent_eoi = _{ !ANY }

evm_arithmetization/src/cpu/kernel/mod.rs

@@ -12,6 +12,8 @@ mod utils;
 
 pub(crate) mod interpreter;
 
+use std::collections::HashSet;
+
 pub use constants::cancun_constants;
 pub use constants::global_exit_root;
 
@@ -26,7 +28,7 @@ use crate::cpu::kernel::constants::evm_constants;
 /// Assemble files, outputting bytes.
 /// This is for debugging the kernel only.
 pub fn assemble_to_bytes(files: &[String]) -> Vec<u8> {
-    let parsed_files: Vec<_> = files.iter().map(|f| parse(f)).collect();
+    let parsed_files: Vec<_> = files.iter().map(|f| parse(f, HashSet::new())).collect();
     let kernel = assemble(parsed_files, evm_constants(), true);
     kernel.code
 }