blocks in stack manipulation

evm/src/cpu/kernel/assembler.rs

@@ -551,6 +551,7 @@ mod tests {
         let dup1 = get_opcode("DUP1");
         let swap1 = get_opcode("SWAP1");
         let swap2 = get_opcode("SWAP2");
+        let swap3 = get_opcode("SWAP3");
         let push_label = get_push_opcode(BYTES_PER_OFFSET);
 
         let kernel = parse_and_assemble(&["%stack (a) -> (a)"]);
@@ -562,6 +563,17 @@ mod tests {
         let kernel = parse_and_assemble(&["%stack (a, b, c) -> (b)"]);
         assert_eq!(kernel.code, vec![pop, swap1, pop]);
 
+        let kernel = parse_and_assemble(&["%stack (a, (b: 3), c) -> (c)"]);
+        assert_eq!(kernel.code, vec![pop, pop, pop, pop]);
+
+        let kernel = parse_and_assemble(&["%stack ((a: 2), (b: 2)) -> (b, a)"]);
+        assert_eq!(kernel.code, vec![swap1, swap3, swap1, swap2]);
+
+        let kernel1 = parse_and_assemble(&["%stack ((a: 3), (b: 3), c) -> (c, b, a)"]);
+        let kernel2 =
+            parse_and_assemble(&["%stack (a, b, c, d, e, f, g) -> (g, d, e, f, a, b, c)"]);
+        assert_eq!(kernel1.code, kernel2.code);
+
         let mut consts = HashMap::new();
         consts.insert("LIFE".into(), 42.into());
         parse_and_assemble_ext(&["%stack (a, b) -> (b, @LIFE)"], consts, true);

evm/src/cpu/kernel/ast.rs

@@ -19,7 +19,7 @@ pub(crate) enum Item {
     /// The first list gives names to items on the top of the stack.
     /// The second list specifies replacement items.
     /// Example: `(a, b, c) -> (c, 5, 0x20, @SOME_CONST, a)`.
-    StackManipulation(Vec<String>, Vec<StackReplacement>),
+    StackManipulation(Vec<StackPlaceholder>, Vec<StackReplacement>),
     /// Declares a global label.
     GlobalLabelDeclaration(String),
     /// Declares a label that is local to the current file.
@@ -36,6 +36,14 @@ pub(crate) enum Item {
     Bytes(Vec<u8>),
 }
 
+/// The left hand side of a %stack stack-manipulation macro.
+#[derive(Eq, PartialEq, Clone, Debug)]
+pub(crate) enum StackPlaceholder {
+    Identifier(String),
+    Block(String, usize),
+}
+
+/// The right hand side of a %stack stack-manipulation macro.
 #[derive(Eq, PartialEq, Clone, Debug)]
 pub(crate) enum StackReplacement {
     /// Can be either a named item or a label.

evm/src/cpu/kernel/evm_asm.pest

@@ -21,7 +21,10 @@ macro_call = ${ "%" ~ !(^"macro" | ^"endmacro" | ^"rep" | ^"endrep" | ^"stack")
 repeat = { ^"%rep" ~ literal ~ item* ~ ^"%endrep" }
 paramlist = { "(" ~ identifier ~ ("," ~ identifier)* ~ ")" }
 macro_arglist = !{ "(" ~ push_target ~ ("," ~ push_target)* ~ ")" }
-stack = { ^"%stack" ~ paramlist ~ "->" ~ stack_replacements }
+stack = { ^"%stack" ~ stack_placeholders ~ "->" ~ stack_replacements }
+stack_placeholders = { "(" ~ stack_placeholder ~ ("," ~ stack_placeholder)* ~ ")" }
+stack_placeholder = { identifier | stack_block }
+stack_block = { "(" ~ identifier ~ ":" ~ literal_decimal ~ ")" }
 stack_replacements = { "(" ~ stack_replacement ~ ("," ~ stack_replacement)* ~ ")" }
 stack_replacement = { literal | identifier | constant | macro_label | variable }
 global_label_decl = ${ ^"GLOBAL " ~ identifier ~ ":" }

evm/src/cpu/kernel/parser.rs

@@ -4,6 +4,7 @@ use ethereum_types::U256;
 use pest::iterators::Pair;
 use pest::Parser;
 
+use super::ast::StackPlaceholder;
 use crate::cpu::kernel::ast::{File, Item, PushTarget, StackReplacement};
 
 /// Parses EVM assembly code.
@@ -99,14 +100,11 @@ fn parse_stack(item: Pair<Rule>) -> Item {
     let mut inner = item.into_inner();
 
     let params = inner.next().unwrap();
-    assert_eq!(params.as_rule(), Rule::paramlist);
+    assert_eq!(params.as_rule(), Rule::stack_placeholders);
     let replacements = inner.next().unwrap();
     assert_eq!(replacements.as_rule(), Rule::stack_replacements);
 
-    let params = params
+    let params = params.into_inner().map(parse_stack_placeholder).collect();
-        .into_inner()
-        .map(|param| param.as_str().to_string())
-        .collect();
     let replacements = replacements
         .into_inner()
         .map(parse_stack_replacement)
@@ -114,6 +112,21 @@ fn parse_stack(item: Pair<Rule>) -> Item {
     Item::StackManipulation(params, replacements)
 }
 
+fn parse_stack_placeholder(target: Pair<Rule>) -> StackPlaceholder {
+    assert_eq!(target.as_rule(), Rule::stack_placeholder);
+    let inner = target.into_inner().next().unwrap();
+    match inner.as_rule() {
+        Rule::identifier => StackPlaceholder::Identifier(inner.as_str().into()),
+        Rule::stack_block => {
+            let mut block = inner.into_inner();
+            let identifier = block.next().unwrap().as_str();
+            let length = block.next().unwrap().as_str().parse().unwrap();
+            StackPlaceholder::Block(identifier.to_string(), length)
+        }
+        _ => panic!("Unexpected {:?}", inner.as_rule()),
+    }
+}
+
 fn parse_stack_replacement(target: Pair<Rule>) -> StackReplacement {
     assert_eq!(target.as_rule(), Rule::stack_replacement);
     let inner = target.into_inner().next().unwrap();

evm/src/cpu/kernel/stack/stack_manipulation.rs

@@ -1,13 +1,13 @@
 use std::cmp::Ordering;
 use std::collections::hash_map::Entry::{Occupied, Vacant};
-use std::collections::{BinaryHeap, HashMap};
+use std::collections::{BinaryHeap, HashMap, HashSet};
 use std::hash::Hash;
 
 use itertools::Itertools;
 
 use crate::cpu::columns::NUM_CPU_COLUMNS;
 use crate::cpu::kernel::assembler::BYTES_PER_OFFSET;
-use crate::cpu::kernel::ast::{Item, PushTarget, StackReplacement};
+use crate::cpu::kernel::ast::{Item, PushTarget, StackPlaceholder, StackReplacement};
 use crate::cpu::kernel::stack::permutations::{get_stack_ops_for_perm, is_permutation};
 use crate::cpu::kernel::stack::stack_manipulation::StackOp::Pop;
 use crate::cpu::kernel::utils::u256_to_trimmed_be_bytes;
@@ -25,25 +25,50 @@ pub(crate) fn expand_stack_manipulation(body: Vec<Item>) -> Vec<Item> {
     expanded
 }
 
-fn expand(names: Vec<String>, replacements: Vec<StackReplacement>) -> Vec<Item> {
+fn expand(names: Vec<StackPlaceholder>, replacements: Vec<StackReplacement>) -> Vec<Item> {
+    let mut stack_blocks = HashMap::new();
+    let mut stack_names = HashSet::new();
+
     let mut src = names
         .iter()
         .cloned()
-        .map(StackItem::NamedItem)
+        .flat_map(|item| match item {
+            StackPlaceholder::Identifier(name) => {
+                stack_names.insert(name.clone());
+                vec![StackItem::NamedItem(name)]
+            }
+            StackPlaceholder::Block(name, n) => {
+                stack_blocks.insert(name.clone(), n);
+                (0..n)
+                    .map(|i| {
+                        let literal_name = format!("block_{}_{}", name, i);
+                        StackItem::NamedItem(literal_name)
+                    })
+                    .collect_vec()
+            }
+        })
         .collect_vec();
 
     let mut dst = replacements
         .into_iter()
-        .map(|item| match item {
+        .flat_map(|item| match item {
             StackReplacement::Identifier(name) => {
                 // May be either a named item or a label. Named items have precedence.
-                if names.contains(&name) {
+                if stack_blocks.contains_key(&name) {
-                    StackItem::NamedItem(name)
+                    let n = *stack_blocks.get(&name).unwrap();
+                    (0..n)
+                        .map(|i| {
+                            let literal_name = format!("block_{}_{}", name, i);
+                            StackItem::NamedItem(literal_name)
+                        })
+                        .collect_vec()
+                } else if stack_names.contains(&name) {
+                    vec![StackItem::NamedItem(name)]
                 } else {
-                    StackItem::PushTarget(PushTarget::Label(name))
+                    vec![StackItem::PushTarget(PushTarget::Label(name))]
                 }
             }
-            StackReplacement::Literal(n) => StackItem::PushTarget(PushTarget::Literal(n)),
+            StackReplacement::Literal(n) => vec![StackItem::PushTarget(PushTarget::Literal(n))],
             StackReplacement::MacroLabel(_)
             | StackReplacement::MacroVar(_)
             | StackReplacement::Constant(_) => {