From f876a8ab024b2d943d4b5be45962b98f050d8664 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Wed, 21 Sep 2022 08:42:56 -0700
Subject: [PATCH] Fix macro vars in %stack directive

---
 evm/src/cpu/kernel/assembler.rs               | 40 ++++++++++++++++---
 evm/src/cpu/kernel/ast.rs                     | 15 ++++++-
 .../cpu/kernel/stack/stack_manipulation.rs    |  3 +-
 3 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs
index 0471bf99..4319612d 100644
--- a/evm/src/cpu/kernel/assembler.rs
+++ b/evm/src/cpu/kernel/assembler.rs
@@ -195,12 +195,12 @@ fn expand_macro_call(
             Item::StackManipulation(before, after) => {
                 let after = after
                     .iter()
-                    .map(|replacement| {
-                        if let StackReplacement::MacroLabel(label) = replacement {
+                    .map(|replacement| match replacement {
+                        StackReplacement::MacroLabel(label) => {
                             StackReplacement::Identifier(get_actual_label(label))
-                        } else {
-                            replacement.clone()
                         }
+                        StackReplacement::MacroVar(var) => get_arg(var).into(),
+                        _ => replacement.clone(),
                     })
                     .collect();
                 Item::StackManipulation(before.clone(), after)
@@ -508,8 +508,8 @@ mod tests {
             "%macro bar(y) PUSH $y %endmacro",
             "%foo(42)",
         ]);
-        let push = get_push_opcode(1);
-        assert_eq!(kernel.code, vec![push, 42, push, 42]);
+        let push1 = get_push_opcode(1);
+        assert_eq!(kernel.code, vec![push1, 42, push1, 42]);
     }
 
     #[test]
@@ -587,6 +587,34 @@ mod tests {
         assert_eq!(kernel.code, vec![dup1]);
     }
 
+    #[test]
+    fn stack_manipulation_in_macro() {
+        let pop = get_opcode("POP");
+        let push1 = get_push_opcode(1);
+
+        let kernel = parse_and_assemble(&[
+            "%macro set_top(x) %stack (a) -> ($x) %endmacro",
+            "%set_top(42)",
+        ]);
+        assert_eq!(kernel.code, vec![pop, push1, 42]);
+    }
+
+    #[test]
+    fn stack_manipulation_in_macro_with_name_collision() {
+        let pop = get_opcode("POP");
+        let push_label = get_push_opcode(BYTES_PER_OFFSET);
+
+        // In the stack directive, there's a named item `foo`.
+        // But when we invoke `%foo(foo)`, the argument refers to the `foo` label.
+        // Thus the expanded macro is `%stack (foo) -> (label foo)` (not real syntax).
+        let kernel = parse_and_assemble(&[
+            "global foo:",
+            "%macro foo(x) %stack (foo) -> ($x) %endmacro",
+            "%foo(foo)",
+        ]);
+        assert_eq!(kernel.code, vec![pop, push_label, 0, 0, 0]);
+    }
+
     fn parse_and_assemble(files: &[&str]) -> Kernel {
         parse_and_assemble_ext(files, HashMap::new(), true)
     }
diff --git a/evm/src/cpu/kernel/ast.rs b/evm/src/cpu/kernel/ast.rs
index bad60d03..19188ad9 100644
--- a/evm/src/cpu/kernel/ast.rs
+++ b/evm/src/cpu/kernel/ast.rs
@@ -46,14 +46,27 @@ pub(crate) enum StackPlaceholder {
 /// The right hand side of a %stack stack-manipulation macro.
 #[derive(Eq, PartialEq, Clone, Debug)]
 pub(crate) enum StackReplacement {
+    Literal(U256),
     /// Can be either a named item or a label.
     Identifier(String),
-    Literal(U256),
+    Label(String),
     MacroLabel(String),
     MacroVar(String),
     Constant(String),
 }
 
+impl From<PushTarget> for StackReplacement {
+    fn from(target: PushTarget) -> Self {
+        match target {
+            PushTarget::Literal(x) => Self::Literal(x),
+            PushTarget::Label(l) => Self::Label(l),
+            PushTarget::MacroLabel(l) => Self::MacroLabel(l),
+            PushTarget::MacroVar(v) => Self::MacroVar(v),
+            PushTarget::Constant(c) => Self::Constant(c),
+        }
+    }
+}
+
 /// The target of a `PUSH` operation.
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub(crate) enum PushTarget {
diff --git a/evm/src/cpu/kernel/stack/stack_manipulation.rs b/evm/src/cpu/kernel/stack/stack_manipulation.rs
index faec7e04..ebc54af1 100644
--- a/evm/src/cpu/kernel/stack/stack_manipulation.rs
+++ b/evm/src/cpu/kernel/stack/stack_manipulation.rs
@@ -52,6 +52,7 @@ fn expand(names: Vec<StackPlaceholder>, replacements: Vec<StackReplacement>) ->
     let mut dst = replacements
         .into_iter()
         .flat_map(|item| match item {
+            StackReplacement::Literal(n) => vec![StackItem::PushTarget(PushTarget::Literal(n))],
             StackReplacement::Identifier(name) => {
                 // May be either a named item or a label. Named items have precedence.
                 if stack_blocks.contains_key(&name) {
@@ -68,7 +69,7 @@ fn expand(names: Vec<StackPlaceholder>, replacements: Vec<StackReplacement>) ->
                     vec![StackItem::PushTarget(PushTarget::Label(name))]
                 }
             }
-            StackReplacement::Literal(n) => vec![StackItem::PushTarget(PushTarget::Literal(n))],
+            StackReplacement::Label(name) => vec![StackItem::PushTarget(PushTarget::Label(name))],
             StackReplacement::MacroLabel(_)
             | StackReplacement::MacroVar(_)
             | StackReplacement::Constant(_) => {