Merge pull request #659 from mir-protocol/macro_labels

Support macro-local labels
2026-01-04 23:03:08 +00:00 · 2022-08-05 09:13:24 -05:00 · 2022-08-05 09:13:24 -05:00 · b741458d02
commit b741458d02
parent 9d3d3afeba 616eb618f2
7 changed files with 125 additions and 83 deletions
--- a/evm/src/cpu/kernel/asm/rlp/decode.asm
+++ b/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
--- a/evm/src/cpu/kernel/asm/transactions/type_0.asm
+++ b/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
-    %jump(decode_rlp_list_len)
-
-process_txn_with_len:
+    // stack: pos
+    %decode_rlp_list_len
    // We don't actually need the length.
    %stack (pos, len) -> (pos)

-    PUSH store_nonce
-    SWAP1
-    // stack: pos, store_nonce
-    %jump(decode_rlp_scalar)
-
-store_nonce:
+    // Decode the nonce and store it.
+    // stack: pos
+    %decode_rlp_scalar
    %stack (pos, nonce) -> (@TXN_FIELD_NONCE, nonce, pos)
    %mstore_current(@SEGMENT_NORMALIZED_TXN)

-    // stack: pos
-    PUSH store_gas_price
-    SWAP1
-    // stack: pos, store_gas_price
-    %jump(decode_rlp_scalar)
-
-store_gas_price:
+    // Decode the gas price and store it.
    // 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
-    SWAP1
-    // stack: pos, store_gas_limit
-    %jump(decode_rlp_scalar)
-
-store_gas_limit:
+    %decode_rlp_scalar
    %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
-    SWAP1
-    // stack: pos, store_to
-    %jump(decode_rlp_scalar)
-
-store_to:
+    %decode_rlp_scalar
    %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
-    SWAP1
-    // stack: pos, store_value
-    %jump(decode_rlp_scalar)
-
-store_value:
+    %decode_rlp_scalar
    %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
-    SWAP1
-    // stack: pos, store_data_len
-    %jump(decode_rlp_string_len)
-
-store_data_len:
+    %decode_rlp_string_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
-    SWAP1
-    // stack: pos, process_v
-    %jump(decode_rlp_scalar)
-
-process_v:
+    %decode_rlp_scalar
    // stack: pos, v
    SWAP1
    // stack: v, pos
@ -154,22 +121,12 @@ process_v_new_style:

 parse_r:
    // stack: pos
-    PUSH store_r
-    SWAP1
-    // stack: pos, store_r
-    %jump(decode_rlp_scalar)
-
-store_r:
+    %decode_rlp_scalar
    %stack (pos, r) -> (@TXN_FIELD_R, r, pos)
    %mstore_current(@SEGMENT_NORMALIZED_TXN)

    // stack: pos
-    PUSH store_s
-    SWAP1
-    // stack: pos, store_s
-    %jump(decode_rlp_scalar)
-
-store_s:
+    %decode_rlp_scalar
    %stack (pos, s) -> (@TXN_FIELD_S, s)
    %mstore_current(@SEGMENT_NORMALIZED_TXN)
    // stack: (empty)
--- a/evm/src/cpu/kernel/assembler.rs
+++ b/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| {
-            if let Item::Push(PushTarget::MacroVar(var)) = item {
-                Item::Push(get_arg(var))
-            } else if let Item::MacroCall(name, args) = item {
+        .map(|item| match item {
+            Item::MacroLabelDeclaration(label) => LocalLabelDeclaration(get_actual_label(label)),
+            Item::Push(PushTarget::MacroLabel(label)) => {
+                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)
-            } else {
-                item.clone()
            }
+            Item::StackManipulation(before, after) => {
+                let after = after
+                    .iter()
+                    .map(|replacement| {
+                        if let StackReplacement::MacroLabel(label) = replacement {
+                            StackReplacement::Identifier(get_actual_label(label))
+                        } else {
+                            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(&[
--- a/evm/src/cpu/kernel/ast.rs
+++ b/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),
 }
--- a/evm/src/cpu/kernel/evm_asm.pest
+++ b/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 }
-global_label = { ^"GLOBAL " ~ identifier ~ ":" }
-local_label = { identifier ~ ":" }
+stack_replacement = { literal | identifier | constant | macro_label | variable }
+global_label_decl = ${ ^"GLOBAL " ~ 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 }
--- a/evm/src/cpu/kernel/parser.rs
+++ b/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()),
--- a/evm/src/cpu/kernel/stack_manipulation.rs
+++ b/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)
                    }
                };