Merge pull request #659 from mir-protocol/macro_labels

Support macro-local labels

evm/src/cpu/kernel/asm/rlp/decode.asm

@@ -54,6 +54,13 @@ decode_rlp_string_len_large:
     // stack: pos', len_of_len, retdest
     %jump(decode_int_given_len)
 
+// Convenience macro to call decode_rlp_string_len and return where we left off.
+%macro decode_rlp_string_len
+    %stack (pos) -> (pos, %%after)
+    %jump(decode_rlp_string_len)
+%%after:
+%endmacro
+
 // Parse a scalar from RLP memory.
 // Pre stack: pos, retdest
 // Post stack: pos', scalar
@@ -73,6 +80,13 @@ global decode_rlp_scalar:
     // to decode_int_given_len.
     %jump(decode_rlp_string_len)
 
+// Convenience macro to call decode_rlp_scalar and return where we left off.
+%macro decode_rlp_scalar
+    %stack (pos) -> (pos, %%after)
+    %jump(decode_rlp_scalar)
+%%after:
+%endmacro
+
 // Parse the length of an RLP list from memory.
 // Pre stack: pos, retdest
 // Post stack: pos', len
@@ -111,6 +125,13 @@ decode_rlp_list_len_big:
     // stack: pos', len_of_len, retdest
     %jump(decode_int_given_len)
 
+// Convenience macro to call decode_rlp_list_len and return where we left off.
+%macro decode_rlp_list_len
+    %stack (pos) -> (pos, %%after)
+    %jump(decode_rlp_list_len)
+%%after:
+%endmacro
+
 // Parse an integer of the given length. It is assumed that the integer will
 // fit in a single (256-bit) word on the stack.
 // Pre stack: pos, len, retdest

evm/src/cpu/kernel/asm/transactions/type_0.asm

@@ -14,78 +14,50 @@
 global process_type_0_txn:
     JUMPDEST
     // stack: (empty)
-    PUSH process_txn_with_len
     PUSH 0 // initial pos
-    // stack: pos, process_txn_with_len
+    // stack: pos
-    %jump(decode_rlp_list_len)
+    %decode_rlp_list_len
-
-process_txn_with_len:
     // We don't actually need the length.
     %stack (pos, len) -> (pos)
 
-    PUSH store_nonce
+    // Decode the nonce and store it.
-    SWAP1
+    // stack: pos
-    // stack: pos, store_nonce
+    %decode_rlp_scalar
-    %jump(decode_rlp_scalar)
-
-store_nonce:
     %stack (pos, nonce) -> (@TXN_FIELD_NONCE, nonce, pos)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
 
-    // stack: pos
+    // Decode the gas price and store it.
-    PUSH store_gas_price
-    SWAP1
-    // stack: pos, store_gas_price
-    %jump(decode_rlp_scalar)
-
-store_gas_price:
     // For legacy transactions, we set both the
     // TXN_FIELD_MAX_PRIORITY_FEE_PER_GAS and TXN_FIELD_MAX_FEE_PER_GAS
     // fields to gas_price.
+    // stack: pos
+    %decode_rlp_scalar
     %stack (pos, gas_price) -> (@TXN_FIELD_MAX_PRIORITY_FEE_PER_GAS, gas_price,
                                 @TXN_FIELD_MAX_FEE_PER_GAS, gas_price, pos)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
 
+    // Decode the gas limit and store it.
     // stack: pos
-    PUSH store_gas_limit
+    %decode_rlp_scalar
-    SWAP1
-    // stack: pos, store_gas_limit
-    %jump(decode_rlp_scalar)
-
-store_gas_limit:
     %stack (pos, gas_limit) -> (@TXN_FIELD_GAS_LIMIT, gas_limit, pos)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
 
+    // Decode the "to" field and store it.
     // stack: pos
-    PUSH store_to
+    %decode_rlp_scalar
-    SWAP1
-    // stack: pos, store_to
-    %jump(decode_rlp_scalar)
-
-store_to:
     %stack (pos, to) -> (@TXN_FIELD_TO, to, pos)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
-    // stack: pos
 
-parse_value:
+    // Decode the value field and store it.
     // stack: pos
-    PUSH store_value
+    %decode_rlp_scalar
-    SWAP1
-    // stack: pos, store_value
-    %jump(decode_rlp_scalar)
-
-store_value:
     %stack (pos, value) -> (@TXN_FIELD_VALUE, value, pos)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
 
+    // Decode the data length, store it, and compute new_pos after any data.
     // stack: pos
-    PUSH store_data_len
+    %decode_rlp_string_len
-    SWAP1
-    // stack: pos, store_data_len
-    %jump(decode_rlp_string_len)
-
-store_data_len:
     %stack (pos, data_len) -> (@TXN_FIELD_DATA_LEN, data_len, pos, data_len, pos, data_len)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
     // stack: pos, data_len, pos, data_len
@@ -105,12 +77,7 @@ store_data_len:
 
 parse_v:
     // stack: pos
-    PUSH process_v
+    %decode_rlp_scalar
-    SWAP1
-    // stack: pos, process_v
-    %jump(decode_rlp_scalar)
-
-process_v:
     // stack: pos, v
     SWAP1
     // stack: v, pos
@@ -154,22 +121,12 @@ process_v_new_style:
 
 parse_r:
     // stack: pos
-    PUSH store_r
+    %decode_rlp_scalar
-    SWAP1
-    // stack: pos, store_r
-    %jump(decode_rlp_scalar)
-
-store_r:
     %stack (pos, r) -> (@TXN_FIELD_R, r, pos)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
 
     // stack: pos
-    PUSH store_s
+    %decode_rlp_scalar
-    SWAP1
-    // stack: pos, store_s
-    %jump(decode_rlp_scalar)
-
-store_s:
     %stack (pos, s) -> (@TXN_FIELD_S, s)
     %mstore_current(@SEGMENT_NORMALIZED_TXN)
     // stack: (empty)

evm/src/cpu/kernel/assembler.rs

@@ -5,6 +5,7 @@ use itertools::izip;
 use log::debug;
 
 use super::ast::PushTarget;
+use crate::cpu::kernel::ast::Item::LocalLabelDeclaration;
 use crate::cpu::kernel::ast::StackReplacement;
 use crate::cpu::kernel::keccak_util::hash_kernel;
 use crate::cpu::kernel::optimizer::optimize_asm;
@@ -76,8 +77,9 @@ pub(crate) fn assemble(
     let mut offset = 0;
     let mut expanded_files = Vec::with_capacity(files.len());
     let mut local_labels = Vec::with_capacity(files.len());
+    let mut macro_counter = 0;
     for file in files {
-        let expanded_file = expand_macros(file.body, &macros);
+        let expanded_file = expand_macros(file.body, &macros, &mut macro_counter);
         let expanded_file = expand_repeats(expanded_file);
         let expanded_file = inline_constants(expanded_file, &constants);
         let mut expanded_file = expand_stack_manipulation(expanded_file);
@@ -120,7 +122,11 @@ fn find_macros(files: &[File]) -> HashMap<String, Macro> {
     macros
 }
 
-fn expand_macros(body: Vec<Item>, macros: &HashMap<String, Macro>) -> Vec<Item> {
+fn expand_macros(
+    body: Vec<Item>,
+    macros: &HashMap<String, Macro>,
+    macro_counter: &mut u32,
+) -> Vec<Item> {
     let mut expanded = vec![];
     for item in body {
         match item {
@@ -128,7 +134,7 @@ fn expand_macros(body: Vec<Item>, macros: &HashMap<String, Macro>) -> Vec<Item>
                 // At this phase, we no longer need macro definitions.
             }
             Item::MacroCall(m, args) => {
-                expanded.extend(expand_macro_call(m, args, macros));
+                expanded.extend(expand_macro_call(m, args, macros, macro_counter));
             }
             item => {
                 expanded.push(item);
@@ -142,6 +148,7 @@ fn expand_macro_call(
     name: String,
     args: Vec<PushTarget>,
     macros: &HashMap<String, Macro>,
+    macro_counter: &mut u32,
 ) -> Vec<Item> {
     let _macro = macros
         .get(&name)
@@ -156,6 +163,8 @@ fn expand_macro_call(
         args.len()
     );
 
+    let get_actual_label = |macro_label| format!("@{}.{}", macro_counter, macro_label);
+
     let get_arg = |var| {
         let param_index = _macro.get_param_index(var);
         args[param_index].clone()
@@ -164,10 +173,13 @@ fn expand_macro_call(
     let expanded_item = _macro
         .items
         .iter()
-        .map(|item| {
+        .map(|item| match item {
-            if let Item::Push(PushTarget::MacroVar(var)) = item {
+            Item::MacroLabelDeclaration(label) => LocalLabelDeclaration(get_actual_label(label)),
-                Item::Push(get_arg(var))
+            Item::Push(PushTarget::MacroLabel(label)) => {
-            } else if let Item::MacroCall(name, args) = item {
+                Item::Push(PushTarget::Label(get_actual_label(label)))
+            }
+            Item::Push(PushTarget::MacroVar(var)) => Item::Push(get_arg(var)),
+            Item::MacroCall(name, args) => {
                 let expanded_args = args
                     .iter()
                     .map(|arg| {
@@ -179,14 +191,28 @@ fn expand_macro_call(
                     })
                     .collect();
                 Item::MacroCall(name.clone(), expanded_args)
+            }
+            Item::StackManipulation(before, after) => {
+                let after = after
+                    .iter()
+                    .map(|replacement| {
+                        if let StackReplacement::MacroLabel(label) = replacement {
+                            StackReplacement::Identifier(get_actual_label(label))
                         } else {
-                item.clone()
+                            replacement.clone()
                         }
                     })
                     .collect();
+                Item::StackManipulation(before.clone(), after)
+            }
+            _ => item.clone(),
+        })
+        .collect();
+
+    *macro_counter += 1;
 
     // Recursively expand any macros in the expanded code.
-    expand_macros(expanded_item, macros)
+    expand_macros(expanded_item, macros, macro_counter)
 }
 
 fn expand_repeats(body: Vec<Item>) -> Vec<Item> {
@@ -247,7 +273,8 @@ fn find_labels(
             Item::MacroDef(_, _, _)
             | Item::MacroCall(_, _)
             | Item::Repeat(_, _)
-            | Item::StackManipulation(_, _) => {
+            | Item::StackManipulation(_, _)
+            | Item::MacroLabelDeclaration(_) => {
                 panic!("Item should have been expanded already: {:?}", item);
             }
             Item::GlobalLabelDeclaration(label) => {
@@ -282,7 +309,8 @@ fn assemble_file(
             Item::MacroDef(_, _, _)
             | Item::MacroCall(_, _)
             | Item::Repeat(_, _)
-            | Item::StackManipulation(_, _) => {
+            | Item::StackManipulation(_, _)
+            | Item::MacroLabelDeclaration(_) => {
                 panic!("Item should have been expanded already: {:?}", item);
             }
             Item::GlobalLabelDeclaration(_) | Item::LocalLabelDeclaration(_) => {
@@ -303,6 +331,7 @@ fn assemble_file(
                             .map(|i| offset.to_le_bytes()[i as usize])
                             .collect()
                     }
+                    PushTarget::MacroLabel(v) => panic!("Macro label not in a macro: {}", v),
                     PushTarget::MacroVar(v) => panic!("Variable not in a macro: {}", v),
                     PushTarget::Constant(c) => panic!("Constant wasn't inlined: {}", c),
                 };
@@ -325,6 +354,7 @@ fn push_target_size(target: &PushTarget) -> u8 {
     match target {
         PushTarget::Literal(n) => u256_to_trimmed_be_bytes(n).len() as u8,
         PushTarget::Label(_) => BYTES_PER_OFFSET,
+        PushTarget::MacroLabel(v) => panic!("Macro label not in a macro: {}", v),
         PushTarget::MacroVar(v) => panic!("Variable not in a macro: {}", v),
         PushTarget::Constant(c) => panic!("Constant wasn't inlined: {}", c),
     }
@@ -456,6 +486,21 @@ mod tests {
         assert_eq!(kernel.code, vec![push1, 2, push1, 3, add]);
     }
 
+    #[test]
+    fn macro_with_label() {
+        let files = &[
+            "%macro spin %%start: PUSH %%start JUMP %endmacro",
+            "%spin %spin",
+        ];
+        let kernel = parse_and_assemble_ext(files, HashMap::new(), false);
+        let push3 = get_push_opcode(BYTES_PER_OFFSET);
+        let jump = get_opcode("JUMP");
+        assert_eq!(
+            kernel.code,
+            vec![push3, 0, 0, 0, jump, push3, 0, 0, 5, jump]
+        );
+    }
+
     #[test]
     fn macro_in_macro_with_vars() {
         let kernel = parse_and_assemble(&[

evm/src/cpu/kernel/ast.rs

@@ -24,6 +24,8 @@ pub(crate) enum Item {
     GlobalLabelDeclaration(String),
     /// Declares a label that is local to the current file.
     LocalLabelDeclaration(String),
+    /// Declares a label that is local to the macro it's declared in.
+    MacroLabelDeclaration(String),
     /// A `PUSH` operation.
     Push(PushTarget),
     /// A `ProverInput` operation.
@@ -39,6 +41,7 @@ pub(crate) enum StackReplacement {
     /// Can be either a named item or a label.
     Identifier(String),
     Literal(U256),
+    MacroLabel(String),
     MacroVar(String),
     Constant(String),
 }
@@ -48,6 +51,7 @@ pub(crate) enum StackReplacement {
 pub(crate) enum PushTarget {
     Literal(U256),
     Label(String),
+    MacroLabel(String),
     MacroVar(String),
     Constant(String),
 }

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 | prover_input_instruction | nullary_instruction }
+item = { macro_def | macro_call | repeat | stack | global_label_decl | local_label_decl | macro_label_decl | 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" }
@@ -23,12 +23,14 @@ paramlist = { "(" ~ identifier ~ ("," ~ identifier)* ~ ")" }
 macro_arglist = !{ "(" ~ push_target ~ ("," ~ push_target)* ~ ")" }
 stack = { ^"%stack" ~ paramlist ~ "->" ~ stack_replacements }
 stack_replacements = { "(" ~ stack_replacement ~ ("," ~ stack_replacement)* ~ ")" }
-stack_replacement = { literal | identifier | constant }
+stack_replacement = { literal | identifier | constant | macro_label | variable }
-global_label = { ^"GLOBAL " ~ identifier ~ ":" }
+global_label_decl = ${ ^"GLOBAL " ~ identifier ~ ":" }
-local_label = { identifier ~ ":" }
+local_label_decl = ${ identifier ~ ":" }
+macro_label_decl = ${ "%%" ~ identifier ~ ":" }
+macro_label = ${ "%%" ~ identifier }
 bytes_item = { ^"BYTES " ~ literal ~ ("," ~ literal)* }
 push_instruction = { ^"PUSH " ~ push_target }
-push_target = { literal | identifier | variable | constant }
+push_target = { literal | identifier | macro_label | variable | constant }
 prover_input_instruction = { ^"PROVER_INPUT" ~ "(" ~ prover_input_fn ~ ")" }
 prover_input_fn = { identifier ~ ("::" ~ identifier)*}
 nullary_instruction = { identifier }

evm/src/cpu/kernel/parser.rs

@@ -28,12 +28,15 @@ fn parse_item(item: Pair<Rule>) -> Item {
         Rule::macro_call => parse_macro_call(item),
         Rule::repeat => parse_repeat(item),
         Rule::stack => parse_stack(item),
-        Rule::global_label => {
+        Rule::global_label_decl => {
             Item::GlobalLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
         }
-        Rule::local_label => {
+        Rule::local_label_decl => {
             Item::LocalLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
         }
+        Rule::macro_label_decl => {
+            Item::MacroLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
+        }
         Rule::bytes_item => Item::Bytes(item.into_inner().map(parse_literal_u8).collect()),
         Rule::push_instruction => Item::Push(parse_push_target(item.into_inner().next().unwrap())),
         Rule::prover_input_instruction => Item::ProverInput(
@@ -117,6 +120,9 @@ fn parse_stack_replacement(target: Pair<Rule>) -> StackReplacement {
     match inner.as_rule() {
         Rule::identifier => StackReplacement::Identifier(inner.as_str().into()),
         Rule::literal => StackReplacement::Literal(parse_literal_u256(inner)),
+        Rule::macro_label => {
+            StackReplacement::MacroLabel(inner.into_inner().next().unwrap().as_str().into())
+        }
         Rule::variable => {
             StackReplacement::MacroVar(inner.into_inner().next().unwrap().as_str().into())
         }
@@ -133,6 +139,9 @@ fn parse_push_target(target: Pair<Rule>) -> PushTarget {
     match inner.as_rule() {
         Rule::literal => PushTarget::Literal(parse_literal_u256(inner)),
         Rule::identifier => PushTarget::Label(inner.as_str().into()),
+        Rule::macro_label => {
+            PushTarget::MacroLabel(inner.into_inner().next().unwrap().as_str().into())
+        }
         Rule::variable => PushTarget::MacroVar(inner.into_inner().next().unwrap().as_str().into()),
         Rule::constant => PushTarget::Constant(inner.into_inner().next().unwrap().as_str().into()),
         _ => panic!("Unexpected {:?}", inner.as_rule()),

evm/src/cpu/kernel/stack_manipulation.rs

@@ -42,7 +42,9 @@ fn expand(names: Vec<String>, replacements: Vec<StackReplacement>) -> Vec<Item>
                 }
             }
             StackReplacement::Literal(n) => StackItem::PushTarget(PushTarget::Literal(n)),
-            StackReplacement::MacroVar(_) | StackReplacement::Constant(_) => {
+            StackReplacement::MacroLabel(_)
+            | StackReplacement::MacroVar(_)
+            | StackReplacement::Constant(_) => {
                 panic!("Should have been expanded already: {:?}", item)
             }
         })
@@ -230,7 +232,9 @@ impl StackOp {
                 let bytes = match target {
                     PushTarget::Literal(n) => u256_to_trimmed_be_bytes(n).len() as u32,
                     PushTarget::Label(_) => BYTES_PER_OFFSET as u32,
-                    PushTarget::MacroVar(_) | PushTarget::Constant(_) => {
+                    PushTarget::MacroLabel(_)
+                    | PushTarget::MacroVar(_)
+                    | PushTarget::Constant(_) => {
                         panic!("Target should have been expanded already: {:?}", target)
                     }
                 };