From 616eb618f22a0f4a037dd897eab47e9a57ce44b4 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Thu, 4 Aug 2022 12:23:48 -0700 Subject: [PATCH] Support macro-local labels Again borrowing syntax from NASM. Example from the test: %macro spin %%start: PUSH %%start JUMP %endmacro One thing this lets us do is create "wrapper" macros which call a function, then return to the code immediately following the macro call, such as %macro decode_rlp_scalar %stack (pos) -> (pos, %%after) %jump(decode_rlp_scalar) %%after: %endmacro I used this to clean up `type_0.asm`. However, since such macros need to insert `%%after` beneath any arguments in the stack, using them will be suboptimal in some cases. I wouldn't worry about it generally, but we might want to avoid them in performance-critical code, or functions with many arguments like `memcpy`. --- evm/src/cpu/kernel/asm/rlp/decode.asm | 21 +++++ .../cpu/kernel/asm/transactions/type_0.asm | 81 +++++-------------- evm/src/cpu/kernel/assembler.rs | 69 +++++++++++++--- evm/src/cpu/kernel/ast.rs | 4 + evm/src/cpu/kernel/evm_asm.pest | 12 +-- evm/src/cpu/kernel/parser.rs | 13 ++- evm/src/cpu/kernel/stack_manipulation.rs | 8 +- 7 files changed, 125 insertions(+), 83 deletions(-) diff --git a/evm/src/cpu/kernel/asm/rlp/decode.asm b/evm/src/cpu/kernel/asm/rlp/decode.asm index 24d8d5a7..0388276a 100644 --- 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 diff --git a/evm/src/cpu/kernel/asm/transactions/type_0.asm b/evm/src/cpu/kernel/asm/transactions/type_0.asm index 4e39f6c3..8711790d 100644 --- 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) diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs index 636251a3..14ec9aa0 100644 --- 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, ¯os); + let expanded_file = expand_macros(file.body, ¯os, &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 { macros } -fn expand_macros(body: Vec, macros: &HashMap) -> Vec { +fn expand_macros( + body: Vec, + macros: &HashMap, + macro_counter: &mut u32, +) -> Vec { let mut expanded = vec![]; for item in body { match item { @@ -128,7 +134,7 @@ fn expand_macros(body: Vec, macros: &HashMap) -> Vec // 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, macros: &HashMap, + macro_counter: &mut u32, ) -> Vec { 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) -> Vec { @@ -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(&[ diff --git a/evm/src/cpu/kernel/ast.rs b/evm/src/cpu/kernel/ast.rs index a0de748a..24cf01e1 100644 --- 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), } diff --git a/evm/src/cpu/kernel/evm_asm.pest b/evm/src/cpu/kernel/evm_asm.pest index 0703798e..8ea7de4b 100644 --- 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 } diff --git a/evm/src/cpu/kernel/parser.rs b/evm/src/cpu/kernel/parser.rs index 66bf0757..9ed578d4 100644 --- a/evm/src/cpu/kernel/parser.rs +++ b/evm/src/cpu/kernel/parser.rs @@ -28,12 +28,15 @@ fn parse_item(item: Pair) -> 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) -> 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) -> 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()), diff --git a/evm/src/cpu/kernel/stack_manipulation.rs b/evm/src/cpu/kernel/stack_manipulation.rs index 71746f16..a1f02c7e 100644 --- a/evm/src/cpu/kernel/stack_manipulation.rs +++ b/evm/src/cpu/kernel/stack_manipulation.rs @@ -42,7 +42,9 @@ fn expand(names: Vec, replacements: Vec) -> Vec } } 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) } };