Stack manipulation macro

Uses a variant of Dijkstra's, with a few pruning mechanics, to find a path of instructions between the two stack states. We don't explicitly store the graph though.

The Dijkstra implementation is somewhat inspired by the `pathfinding` crate. That crate doesn't quite fit our needs though.

If we need to make it faster later, there are a lot of allocations and clones that we could probably eliminate.

evm/src/cpu/kernel/asm/curve_add.asm

@@ -111,18 +111,7 @@ ec_add_snd_zero:
     // stack: x0, y0, x1, y1, retdest
 
     // Just return (x1,y1)
-    SWAP2
-    // stack: x1, y0, x0, y1, retdest
-    POP
-    // stack: y0, x0, y1, retdest
-    SWAP2
-    // stack: y1, x0, y0, retdest
-    POP
-    // stack: x0, y0, retdest
-    SWAP1
-    // stack: y0, x0, retdest
-    SWAP2
-    // stack: retdest, x0, y0
+    %stack (x0, y0, x1, y1, retdest) -> (retdest, x0, y0)
     JUMP
 
 // BN254 elliptic curve addition.
@@ -170,16 +159,7 @@ ec_add_valid_points_with_lambda:
     // stack: y2, x2, lambda, x0, y0, x1, y1, retdest
 
     // Return x2,y2
-    SWAP5
-    // stack: x1, x2, lambda, x0, y0, y2, y1, retdest
-    POP
-    // stack: x2, lambda, x0, y0, y2, y1, retdest
-    SWAP5
-    // stack: y1, lambda, x0, y0, y2, x2, retdest
-    %pop4
-    // stack: y2, x2, retdest
-    SWAP2
-    // stack: retdest, x2, y2
+    %stack (y2, x2, lambda, x0, y0, x1, y1, retdest) -> (retdest, x2, y2)
     JUMP
 
 // BN254 elliptic curve addition.

evm/src/cpu/kernel/assembler.rs

@@ -7,6 +7,7 @@ use log::debug;
 use super::ast::PushTarget;
 use crate::cpu::kernel::ast::Literal;
 use crate::cpu::kernel::keccak_util::hash_kernel;
+use crate::cpu::kernel::stack_manipulation::expand_stack_manipulation;
 use crate::cpu::kernel::{
     ast::{File, Item},
     opcodes::{get_opcode, get_push_opcode},
@@ -63,6 +64,7 @@ pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Ke
         let expanded_file = expand_macros(file.body, &macros);
         let expanded_file = expand_repeats(expanded_file);
         let expanded_file = inline_constants(expanded_file, &constants);
+        let expanded_file = expand_stack_manipulation(expanded_file);
         local_labels.push(find_labels(&expanded_file, &mut offset, &mut global_labels));
         expanded_files.push(expanded_file);
     }
@@ -187,8 +189,11 @@ fn find_labels(
     let mut local_labels = HashMap::<String, usize>::new();
     for item in body {
         match item {
-            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) | Item::Repeat(_, _) => {
-                panic!("Macros and repeats should have been expanded already")
+            Item::MacroDef(_, _, _)
+            | Item::MacroCall(_, _)
+            | Item::Repeat(_, _)
+            | Item::StackManipulation(_, _) => {
+                panic!("Item should have been expanded already: {:?}", item);
             }
             Item::GlobalLabelDeclaration(label) => {
                 let old = global_labels.insert(label.clone(), *offset);
@@ -215,8 +220,11 @@ fn assemble_file(
     // Assemble the file.
     for item in body {
         match item {
-            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) | Item::Repeat(_, _) => {
-                panic!("Macros and repeats should have been expanded already")
+            Item::MacroDef(_, _, _)
+            | Item::MacroCall(_, _)
+            | Item::Repeat(_, _)
+            | Item::StackManipulation(_, _) => {
+                panic!("Item should have been expanded already: {:?}", item);
             }
             Item::GlobalLabelDeclaration(_) | Item::LocalLabelDeclaration(_) => {
                 // Nothing to do; we processed labels in the prior phase.
@@ -427,6 +435,13 @@ mod tests {
         assert_eq!(kernel.code, vec![add, add, add]);
     }
 
+    #[test]
+    fn stack_manipulation() {
+        let kernel = parse_and_assemble(&["%stack (a, b, c) -> (c, b, a)"]);
+        let swap2 = get_opcode("SWAP2");
+        assert_eq!(kernel.code, vec![swap2]);
+    }
+
     fn parse_and_assemble(files: &[&str]) -> Kernel {
         parse_and_assemble_with_constants(files, HashMap::new())
     }

evm/src/cpu/kernel/ast.rs

@@ -14,6 +14,11 @@ pub(crate) enum Item {
     MacroCall(String, Vec<PushTarget>),
     /// Repetition, like `%rep` in NASM.
     Repeat(Literal, Vec<Item>),
+    /// A directive to manipulate the stack according to a specified pattern.
+    /// 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>),
     /// Declares a global label.
     GlobalLabelDeclaration(String),
     /// Declares a label that is local to the current file.
@@ -26,6 +31,14 @@ pub(crate) enum Item {
     Bytes(Vec<Literal>),
 }
 
+#[derive(Clone, Debug)]
+pub(crate) enum StackReplacement {
+    NamedItem(String),
+    Literal(Literal),
+    MacroVar(String),
+    Constant(String),
+}
+
 /// The target of a `PUSH` operation.
 #[derive(Clone, Debug)]
 pub(crate) enum PushTarget {
@@ -35,7 +48,7 @@ pub(crate) enum PushTarget {
     Constant(String),
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub(crate) enum Literal {
     Decimal(String),
     Hex(String),

evm/src/cpu/kernel/evm_asm.pest

@@ -15,12 +15,15 @@ literal = { literal_hex | literal_decimal }
 variable = ${ "$" ~ identifier }
 constant = ${ "@" ~ identifier }
 
-item = { macro_def | macro_call | repeat | global_label | local_label | bytes_item | push_instruction | nullary_instruction }
-macro_def = { ^"%macro" ~ identifier ~ macro_paramlist? ~ item* ~ ^"%endmacro" }
-macro_call = ${ "%" ~ !(^"macro" | ^"endmacro" | ^"rep" | ^"endrep") ~ identifier ~ macro_arglist? }
+item = { macro_def | macro_call | repeat | stack | global_label | local_label | bytes_item | push_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" }
-macro_paramlist = { "(" ~ identifier ~ ("," ~ identifier)* ~ ")" }
+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 ~ ":" }
 bytes_item = { ^"BYTES " ~ literal ~ ("," ~ literal)* }

evm/src/cpu/kernel/mod.rs

@@ -4,6 +4,7 @@ mod ast;
 pub(crate) mod keccak_util;
 mod opcodes;
 mod parser;
+mod stack_manipulation;
 
 #[cfg(test)]
 mod interpreter;

evm/src/cpu/kernel/parser.rs

@@ -1,7 +1,7 @@
 use pest::iterators::Pair;
 use pest::Parser;
 
-use crate::cpu::kernel::ast::{File, Item, Literal, PushTarget};
+use crate::cpu::kernel::ast::{File, Item, Literal, PushTarget, StackReplacement};
 
 /// Parses EVM assembly code.
 #[derive(pest_derive::Parser)]
@@ -24,6 +24,7 @@ fn parse_item(item: Pair<Rule>) -> Item {
         Rule::macro_def => parse_macro_def(item),
         Rule::macro_call => parse_macro_call(item),
         Rule::repeat => parse_repeat(item),
+        Rule::stack => parse_stack(item),
         Rule::global_label => {
             Item::GlobalLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
         }
@@ -44,7 +45,7 @@ fn parse_macro_def(item: Pair<Rule>) -> Item {
     let name = inner.next().unwrap().as_str().into();
 
     // The parameter list is optional.
-    let params = if let Some(Rule::macro_paramlist) = inner.peek().map(|pair| pair.as_rule()) {
+    let params = if let Some(Rule::paramlist) = inner.peek().map(|pair| pair.as_rule()) {
         let params = inner.next().unwrap().into_inner();
         params.map(|param| param.as_str().to_string()).collect()
     } else {
@@ -78,6 +79,42 @@ fn parse_repeat(item: Pair<Rule>) -> Item {
     Item::Repeat(count, inner.map(parse_item).collect())
 }
 
+fn parse_stack(item: Pair<Rule>) -> Item {
+    assert_eq!(item.as_rule(), Rule::stack);
+    let mut inner = item.into_inner().peekable();
+
+    let params = inner.next().unwrap();
+    assert_eq!(params.as_rule(), Rule::paramlist);
+    let replacements = inner.next().unwrap();
+    assert_eq!(replacements.as_rule(), Rule::stack_replacements);
+
+    let params = params
+        .into_inner()
+        .map(|param| param.as_str().to_string())
+        .collect();
+    let replacements = replacements
+        .into_inner()
+        .map(parse_stack_replacement)
+        .collect();
+    Item::StackManipulation(params, replacements)
+}
+
+fn parse_stack_replacement(target: Pair<Rule>) -> StackReplacement {
+    assert_eq!(target.as_rule(), Rule::stack_replacement);
+    let inner = target.into_inner().next().unwrap();
+    match inner.as_rule() {
+        Rule::identifier => StackReplacement::NamedItem(inner.as_str().into()),
+        Rule::literal => StackReplacement::Literal(parse_literal(inner)),
+        Rule::variable => {
+            StackReplacement::MacroVar(inner.into_inner().next().unwrap().as_str().into())
+        }
+        Rule::constant => {
+            StackReplacement::Constant(inner.into_inner().next().unwrap().as_str().into())
+        }
+        _ => panic!("Unexpected {:?}", inner.as_rule()),
+    }
+}
+
 fn parse_push_target(target: Pair<Rule>) -> PushTarget {
     assert_eq!(target.as_rule(), Rule::push_target);
     let inner = target.into_inner().next().unwrap();

evm/src/cpu/kernel/stack_manipulation.rs

@@ -0,0 +1,224 @@
+use std::cmp::Ordering;
+use std::collections::hash_map::Entry::{Occupied, Vacant};
+use std::collections::{BinaryHeap, HashMap};
+
+use itertools::Itertools;
+
+use crate::cpu::columns::NUM_CPU_COLUMNS;
+use crate::cpu::kernel::ast::{Item, Literal, PushTarget, StackReplacement};
+use crate::cpu::kernel::stack_manipulation::StackOp::Pop;
+use crate::memory;
+
+pub(crate) fn expand_stack_manipulation(body: Vec<Item>) -> Vec<Item> {
+    let mut expanded = vec![];
+    for item in body {
+        if let Item::StackManipulation(names, replacements) = item {
+            expanded.extend(expand(names, replacements));
+        } else {
+            expanded.push(item);
+        }
+    }
+    expanded
+}
+
+fn expand(names: Vec<String>, replacements: Vec<StackReplacement>) -> Vec<Item> {
+    let mut src = names.into_iter().map(StackItem::NamedItem).collect_vec();
+
+    let unique_literals = replacements
+        .iter()
+        .filter_map(|item| match item {
+            StackReplacement::Literal(n) => Some(n.clone()),
+            _ => None,
+        })
+        .unique()
+        .collect_vec();
+    let all_ops = StackOp::all(unique_literals);
+
+    let mut dst = replacements
+        .into_iter()
+        .map(|item| match item {
+            StackReplacement::NamedItem(name) => StackItem::NamedItem(name),
+            StackReplacement::Literal(n) => StackItem::Literal(n),
+            StackReplacement::MacroVar(_) | StackReplacement::Constant(_) => {
+                panic!("Should have been expanded earlier")
+            }
+        })
+        .collect_vec();
+
+    // %stack uses our convention where the top item is written on the left side.
+    // `shortest_path` expects the opposite, so we reverse src and dst.
+    src.reverse();
+    dst.reverse();
+
+    let path = shortest_path(src, dst, all_ops);
+    path.into_iter().map(StackOp::into_item).collect()
+}
+
+/// Finds the lowest-cost sequence of `StackOp`s that transforms `src` to `dst`.
+/// Uses a variant of Dijkstra's algorithm.
+fn shortest_path(src: Vec<StackItem>, dst: Vec<StackItem>, all_ops: Vec<StackOp>) -> Vec<StackOp> {
+    // Nodes to visit, starting with the lowest-cost node.
+    let mut queue = BinaryHeap::new();
+    queue.push(Node {
+        stack: src.clone(),
+        cost: 0,
+    });
+
+    // For each node, stores `(best_cost, Option<(parent, op)>)`.
+    let mut node_info = HashMap::<Vec<StackItem>, (u32, Option<(Vec<StackItem>, StackOp)>)>::new();
+    node_info.insert(src.clone(), (0, None));
+
+    while let Some(node) = queue.pop() {
+        if node.stack == dst {
+            // The destination is now the lowest-cost node, so we must have found the best path.
+            let mut path = vec![];
+            let mut stack = &node.stack;
+            // Rewind back to src, recording a list of operations which will be backwards.
+            while let Some((parent, op)) = &node_info[stack].1 {
+                stack = parent;
+                path.push(op.clone());
+            }
+            assert_eq!(stack, &src);
+            path.reverse();
+            return path;
+        }
+
+        let (best_cost, _) = node_info[&node.stack];
+        if best_cost < node.cost {
+            // Since we can't efficiently remove nodes from the heap, it can contain duplicates.
+            // In this case, we've already visited this stack state with a lower cost.
+            continue;
+        }
+
+        for op in &all_ops {
+            let neighbor = match op.apply_to(node.stack.clone()) {
+                Some(n) => n,
+                None => continue,
+            };
+
+            let cost = node.cost + op.cost();
+            let entry = node_info.entry(neighbor.clone());
+            if let Occupied(e) = &entry && e.get().0 <= cost {
+                // We already found a better or equal path.
+                continue;
+            }
+
+            let neighbor_info = (cost, Some((node.stack.clone(), op.clone())));
+            match entry {
+                Occupied(mut e) => {
+                    e.insert(neighbor_info);
+                }
+                Vacant(e) => {
+                    e.insert(neighbor_info);
+                }
+            }
+
+            queue.push(Node {
+                stack: neighbor,
+                cost,
+            });
+        }
+    }
+
+    panic!("No path found from {:?} to {:?}", src, dst)
+}
+
+/// A node in the priority queue used by Dijkstra's algorithm.
+#[derive(Eq, PartialEq)]
+struct Node {
+    stack: Vec<StackItem>,
+    cost: u32,
+}
+
+impl PartialOrd for Node {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for Node {
+    fn cmp(&self, other: &Self) -> Ordering {
+        // We want a min-heap rather than the default max-heap, so this is the opposite of the
+        // natural ordering of costs.
+        other.cost.cmp(&self.cost)
+    }
+}
+
+/// Like `StackReplacement`, but without constants or macro vars, since those were expanded already.
+#[derive(Eq, PartialEq, Hash, Clone, Debug)]
+enum StackItem {
+    NamedItem(String),
+    Literal(Literal),
+}
+
+#[derive(Clone, Debug)]
+enum StackOp {
+    Push(Literal),
+    Pop,
+    Dup(u8),
+    Swap(u8),
+}
+
+fn get_ops(src: Vec<StackItem>, dst: Vec<StackItem>) -> impl Iterator<Item = StackOp> {
+
+}
+
+impl StackOp {
+    fn all(literals: Vec<Literal>) -> Vec<Self> {
+        let mut all = literals.into_iter().map(StackOp::Push).collect_vec();
+        all.push(Pop);
+        all.extend((1..=32).map(StackOp::Dup));
+        all.extend((1..=32).map(StackOp::Swap));
+        all
+    }
+
+    fn cost(&self) -> u32 {
+        let (cpu_rows, memory_rows) = match self {
+            StackOp::Push(n) => {
+                let bytes = n.to_trimmed_be_bytes().len() as u32;
+                // This is just a rough estimate; we can update it after implementing PUSH.
+                (bytes, bytes)
+            }
+            Pop => (1, 1),
+            StackOp::Dup(_) => (1, 2),
+            StackOp::Swap(_) => (1, 4),
+        };
+
+        let cpu_cost = cpu_rows * NUM_CPU_COLUMNS as u32;
+        let memory_cost = memory_rows * memory::columns::NUM_COLUMNS as u32;
+        cpu_cost + memory_cost
+    }
+
+    /// Returns an updated stack after this operation is performed, or `None` if this operation
+    /// would not be valid on the given stack.
+    fn apply_to(&self, mut stack: Vec<StackItem>) -> Option<Vec<StackItem>> {
+        let len = stack.len();
+        match self {
+            StackOp::Push(n) => {
+                stack.push(StackItem::Literal(n.clone()));
+            }
+            Pop => {
+                stack.pop()?;
+            }
+            StackOp::Dup(n) => {
+                let idx = len.checked_sub(*n as usize)?;
+                stack.push(stack[idx].clone());
+            }
+            StackOp::Swap(n) => {
+                let from = len.checked_sub(1)?;
+                let to = len.checked_sub(*n as usize + 1)?;
+                stack.swap(from, to);
+            }
+        }
+        Some(stack)
+    }
+
+    fn into_item(self) -> Item {
+        match self {
+            StackOp::Push(n) => Item::Push(PushTarget::Literal(n)),
+            Pop => Item::StandardOp("POP".into()),
+            StackOp::Dup(n) => Item::StandardOp(format!("DUP{}", n)),
+            StackOp::Swap(n) => Item::StandardOp(format!("SWAP{}", n)),
+        }
+    }
+}