Add %rep syntax for repeating a block

Same syntax as NASM.

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

@@ -9,18 +9,21 @@
 %endmacro
 
 %macro pop2
-    pop
-    pop
+    %rep 2
+        pop
+    %endrep
 %endmacro
 
 %macro pop3
-    pop
-    %pop2
+    %rep 3
+        pop
+    %endrep
 %endmacro
 
 %macro pop4
-    %pop2
-    %pop2
+    %rep 4
+        pop
+    %endrep
 %endmacro
 
 // If pred is zero, yields z; otherwise, yields nz

evm/src/cpu/kernel/assembler.rs

@@ -44,6 +44,7 @@ pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Ke
     let mut local_labels = Vec::with_capacity(files.len());
     for file in files {
         let expanded_file = expand_macros(file.body, &macros);
+        let expanded_file = expand_repeats(expanded_file);
         let expanded_file = inline_constants(expanded_file, &constants);
         local_labels.push(find_labels(&expanded_file, &mut offset, &mut global_labels));
         expanded_files.push(expanded_file);
@@ -132,6 +133,21 @@ fn expand_macro_call(
     expand_macros(expanded_item, macros)
 }
 
+fn expand_repeats(body: Vec<Item>) -> Vec<Item> {
+    let mut expanded = vec![];
+    for item in body {
+        if let Item::Repeat(count, block) = item {
+            let reps = count.to_u256().as_usize();
+            for _ in 0..reps {
+                expanded.extend(block.clone());
+            }
+        } else {
+            expanded.push(item);
+        }
+    }
+    expanded
+}
+
 fn inline_constants(body: Vec<Item>, constants: &HashMap<String, U256>) -> Vec<Item> {
     body.into_iter()
         .map(|item| {
@@ -157,8 +173,8 @@ fn find_labels(
     let mut local_labels = HashMap::<String, usize>::new();
     for item in body {
         match item {
-            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) => {
-                panic!("Macros should have been expanded already")
+            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) | Item::Repeat(_, _) => {
+                panic!("Macros and repeats should have been expanded already")
             }
             Item::GlobalLabelDeclaration(label) => {
                 let old = global_labels.insert(label.clone(), *offset);
@@ -185,8 +201,8 @@ fn assemble_file(
     // Assemble the file.
     for item in body {
         match item {
-            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) => {
-                panic!("Macros should have been expanded already")
+            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) | Item::Repeat(_, _) => {
+                panic!("Macros and repeats should have been expanded already")
             }
             Item::GlobalLabelDeclaration(_) | Item::LocalLabelDeclaration(_) => {
                 // Nothing to do; we processed labels in the prior phase.
@@ -393,6 +409,13 @@ mod tests {
         assert_eq!(kernel.code, vec![push4, 0xDE, 0xAD, 0xBE, 0xEF]);
     }
 
+    #[test]
+    fn repeat() {
+        let kernel = parse_and_assemble(&["%rep 3 ADD %endrep"]);
+        let add = get_opcode("ADD");
+        assert_eq!(kernel.code, vec![add, add, add]);
+    }
+
     fn parse_and_assemble(files: &[&str]) -> Kernel {
         parse_and_assemble_with_constants(files, HashMap::new())
     }

evm/src/cpu/kernel/ast.rs

@@ -12,6 +12,8 @@ pub(crate) enum Item {
     MacroDef(String, Vec<String>, Vec<Item>),
     /// Calls a macro: name, args.
     MacroCall(String, Vec<PushTarget>),
+    /// Repetition, like `%rep` in NASM.
+    Repeat(Literal, Vec<Item>),
     /// Declares a global label.
     GlobalLabelDeclaration(String),
     /// Declares a label that is local to the current file.

evm/src/cpu/kernel/evm_asm.pest

@@ -15,9 +15,10 @@ literal = { literal_hex | literal_decimal }
 variable = ${ "$" ~ identifier }
 constant = ${ "@" ~ identifier }
 
-item = { macro_def | macro_call | global_label | local_label | bytes_item | push_instruction | nullary_instruction }
+item = { macro_def | macro_call | repeat | global_label | local_label | bytes_item | push_instruction | nullary_instruction }
 macro_def = { ^"%macro" ~ identifier ~ macro_paramlist? ~ item* ~ ^"%endmacro" }
-macro_call = ${ "%" ~ !(^"macro" | ^"endmacro") ~ identifier ~ macro_arglist? }
+macro_call = ${ "%" ~ !(^"macro" | ^"endmacro" | ^"rep" | ^"endrep") ~ identifier ~ macro_arglist? }
+repeat = { ^"%rep" ~ literal ~ item* ~ ^"%endrep" }
 macro_paramlist = { "(" ~ identifier ~ ("," ~ identifier)* ~ ")" }
 macro_arglist = !{ "(" ~ push_target ~ ("," ~ push_target)* ~ ")" }
 global_label = { ^"GLOBAL " ~ identifier ~ ":" }

evm/src/cpu/kernel/parser.rs

@@ -23,6 +23,7 @@ fn parse_item(item: Pair<Rule>) -> Item {
     match item.as_rule() {
         Rule::macro_def => parse_macro_def(item),
         Rule::macro_call => parse_macro_call(item),
+        Rule::repeat => parse_repeat(item),
         Rule::global_label => {
             Item::GlobalLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
         }
@@ -70,6 +71,13 @@ fn parse_macro_call(item: Pair<Rule>) -> Item {
     Item::MacroCall(name, args)
 }
 
+fn parse_repeat(item: Pair<Rule>) -> Item {
+    assert_eq!(item.as_rule(), Rule::repeat);
+    let mut inner = item.into_inner().peekable();
+    let count = parse_literal(inner.next().unwrap());
+    Item::Repeat(count, inner.map(parse_item).collect())
+}
+
 fn parse_push_target(target: Pair<Rule>) -> PushTarget {
     assert_eq!(target.as_rule(), Rule::push_target);
     let inner = target.into_inner().next().unwrap();