ASM macro support (#580)

* ASM macro support

Also recognize global labels as a PUSH target; previously it only considered local labels.

* macro test

evm/src/cpu/kernel/aggregator.rs

@@ -8,6 +8,7 @@ use crate::cpu::kernel::parser::parse;
 #[allow(dead_code)] // TODO: Should be used once witness generation is done.
 pub(crate) fn combined_kernel() -> Kernel {
     let files = vec![
+        include_str!("asm/basic_macros.asm"),
         include_str!("asm/exp.asm"),
         include_str!("asm/storage_read.asm"),
         include_str!("asm/storage_write.asm"),

evm/src/cpu/kernel/asm/basic_macros.asm

@@ -0,0 +1,28 @@
+// If pred is zero, yields z; otherwise, yields nz
+%macro select
+    // stack: pred, nz, z
+    iszero
+    // stack: pred == 0, nz, z
+    dup1
+    // stack: pred == 0, pred == 0, nz, z
+    iszero
+    // stack: pred != 0, pred == 0, nz, z
+    swap3
+    // stack: z, pred == 0, nz, pred != 0
+    mul
+    // stack: (pred == 0) * z, nz, pred != 0
+    swap2
+    // stack: pred != 0, nz, (pred == 0) * z
+    mul
+    // stack: (pred != 0) * nz, (pred == 0) * z
+    add
+    // stack: (pred != 0) * nz + (pred == 0) * z
+%endmacro
+
+%macro square
+    // stack: x
+    dup1
+    // stack: x, x
+    mul
+    // stack: x^2
+%endmacro

evm/src/cpu/kernel/asm/exp.asm

@@ -10,8 +10,6 @@
 /// Note that this correctly handles exp(0, 0) == 1.
 
 global exp:
-// We don't seem to handle global labels yet, so this function has a local label too for now:
-exp:
     // stack: x, e, retdest
     dup2
     // stack: e, x, e, retdest
@@ -41,9 +39,7 @@ step_case:
     // stack: e / 2, recursion_return, x, e, retdest
     dup3
     // stack: x, e / 2, recursion_return, x, e, retdest
-    dup1
-    // stack: x, x, e / 2, recursion_return, x, e, retdest
-    mul
+    %square
     // stack: x * x, e / 2, recursion_return, x, e, retdest
     push exp
     // stack: exp, x * x, e / 2, recursion_return, x, e, retdest

evm/src/cpu/kernel/assembler.rs

@@ -18,10 +18,12 @@ pub struct Kernel {
 }
 
 pub(crate) fn assemble(files: Vec<File>) -> Kernel {
+    let macros = find_macros(&files);
     let mut code = vec![];
     let mut global_labels = HashMap::new();
     for file in files {
-        assemble_file(file.body, &mut code, &mut global_labels);
+        let expanded_file = expand_macros(file.body, &macros);
+        assemble_file(expanded_file, &mut code, &mut global_labels);
     }
     Kernel {
         code,
@@ -29,12 +31,51 @@ pub(crate) fn assemble(files: Vec<File>) -> Kernel {
     }
 }
 
+fn find_macros(files: &[File]) -> HashMap<String, Vec<Item>> {
+    let mut macros = HashMap::new();
+    for file in files {
+        for item in &file.body {
+            if let Item::MacroDef(name, items) = item {
+                macros.insert(name.clone(), items.clone());
+            }
+        }
+    }
+    macros
+}
+
+fn expand_macros(body: Vec<Item>, macros: &HashMap<String, Vec<Item>>) -> Vec<Item> {
+    let mut expanded = vec![];
+    for item in body {
+        match item {
+            Item::MacroDef(_, _) => {
+                // At this phase, we no longer need macro definitions.
+            }
+            Item::MacroCall(m) => {
+                let mut expanded_item = macros
+                    .get(&m)
+                    .cloned()
+                    .unwrap_or_else(|| panic!("No such macro: {}", m));
+                // Recursively expand any macros in the expanded code.
+                expanded_item = expand_macros(expanded_item, macros);
+                expanded.extend(expanded_item);
+            }
+            item => {
+                expanded.push(item);
+            }
+        }
+    }
+    expanded
+}
+
 fn assemble_file(body: Vec<Item>, code: &mut Vec<u8>, global_labels: &mut HashMap<String, usize>) {
-    // First discover the offset of each label  in this file.
+    // First discover the offset of each label in this file.
     let mut local_labels = HashMap::<String, usize>::new();
     let mut offset = code.len();
     for item in &body {
         match item {
+            Item::MacroDef(_, _) | Item::MacroCall(_) => {
+                panic!("Macros should have been expanded already")
+            }
             Item::GlobalLabelDeclaration(label) => {
                 let old = global_labels.insert(label.clone(), offset);
                 assert!(old.is_none(), "Duplicate global label: {}", label);
@@ -52,6 +93,9 @@ fn assemble_file(body: Vec<Item>, code: &mut Vec<u8>, global_labels: &mut HashMa
     // Now that we have label offsets, we can assemble the file.
     for item in body {
         match item {
+            Item::MacroDef(_, _) | Item::MacroCall(_) => {
+                panic!("Macros should have been expanded already")
+            }
             Item::GlobalLabelDeclaration(_) | Item::LocalLabelDeclaration(_) => {
                 // Nothing to do; we processed labels in the prior phase.
             }
@@ -59,7 +103,10 @@ fn assemble_file(body: Vec<Item>, code: &mut Vec<u8>, global_labels: &mut HashMa
                 let target_bytes: Vec<u8> = match target {
                     PushTarget::Literal(literal) => literal.to_trimmed_be_bytes(),
                     PushTarget::Label(label) => {
-                        let offset = local_labels[&label];
+                        let offset = local_labels
+                            .get(&label)
+                            .or_else(|| global_labels.get(&label))
+                            .unwrap_or_else(|| panic!("No such label: {}", label));
                         // We want the BYTES_PER_OFFSET least significant bytes in BE order.
                         // It's easiest to rev the first BYTES_PER_OFFSET bytes of the LE encoding.
                         (0..BYTES_PER_OFFSET)
@@ -97,6 +144,9 @@ fn push_target_size(target: &PushTarget) -> u8 {
 mod tests {
     use std::collections::HashMap;
 
+    use itertools::Itertools;
+
+    use crate::cpu::kernel::parser::parse;
     use crate::cpu::kernel::{assembler::*, ast::*};
 
     #[test]
@@ -202,6 +252,22 @@ mod tests {
             ],
         };
         let code = assemble(vec![file]).code;
-        assert_eq!(code, vec![0x12, 42, 0xfe, 255])
+        assert_eq!(code, vec![0x12, 42, 0xfe, 255]);
+    }
+
+    #[test]
+    fn macro_in_macro() {
+        let kernel = parse_and_assemble(&[
+            "%macro foo %bar %bar %endmacro",
+            "%macro bar ADD %endmacro",
+            "%foo",
+        ]);
+        let add = get_opcode("ADD");
+        assert_eq!(kernel.code, vec![add, add]);
+    }
+
+    fn parse_and_assemble(files: &[&str]) -> Kernel {
+        let parsed_files = files.iter().map(|f| parse(f)).collect_vec();
+        assemble(parsed_files)
     }
 }

evm/src/cpu/kernel/ast.rs

@@ -6,8 +6,12 @@ pub(crate) struct File {
     pub(crate) body: Vec<Item>,
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) enum Item {
+    /// Defines a new macro.
+    MacroDef(String, Vec<Item>),
+    /// Calls a macro.
+    MacroCall(String),
     /// Declares a global label.
     GlobalLabelDeclaration(String),
     /// Declares a label that is local to the current file.
@@ -21,13 +25,13 @@ pub(crate) enum Item {
 }
 
 /// The target of a `PUSH` operation.
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) enum PushTarget {
     Literal(Literal),
     Label(String),
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) enum Literal {
     Decimal(String),
     Hex(String),

evm/src/cpu/kernel/evm_asm.pest

@@ -12,7 +12,9 @@ literal_decimal = @{ ASCII_DIGIT+ }
 literal_hex = @{ ^"0x" ~ ASCII_HEX_DIGIT+ }
 literal = { literal_hex | literal_decimal }
 
-item = { global_label | local_label | bytes_item | push_instruction | nullary_instruction }
+item = { macro_def | macro_call | global_label | local_label | bytes_item | push_instruction | nullary_instruction }
+macro_def = { ^"%macro" ~ identifier ~ item* ~ ^"%endmacro" }
+macro_call = ${ "%" ~ !(^"macro" | ^"endmacro") ~ identifier }
 global_label = { ^"GLOBAL " ~ identifier ~ ":" }
 local_label = { identifier ~ ":" }
 bytes_item = { ^"BYTES " ~ literal ~ ("," ~ literal)* }

evm/src/cpu/kernel/parser.rs

@@ -20,6 +20,12 @@ pub(crate) fn parse(s: &str) -> File {
 fn parse_item(item: Pair<Rule>) -> Item {
     let item = item.into_inner().next().unwrap();
     match item.as_rule() {
+        Rule::macro_def => {
+            let mut inner = item.into_inner();
+            let name = inner.next().unwrap().as_str().into();
+            Item::MacroDef(name, inner.map(parse_item).collect())
+        }
+        Rule::macro_call => Item::MacroCall(item.into_inner().next().unwrap().as_str().into()),
         Rule::global_label => {
             Item::GlobalLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
         }