From 2b9600e50c0ba377c05819bd776d0192af59cdbc Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sun, 31 Jul 2022 13:13:39 -0700
Subject: [PATCH] Misc

---
 evm/src/cpu/kernel/ast.rs       |   4 +-
 evm/src/cpu/kernel/optimizer.rs | 116 +++++++++++++++++++++++++-------
 2 files changed, 92 insertions(+), 28 deletions(-)

diff --git a/evm/src/cpu/kernel/ast.rs b/evm/src/cpu/kernel/ast.rs
index bc2a3ec2..a0de748a 100644
--- a/evm/src/cpu/kernel/ast.rs
+++ b/evm/src/cpu/kernel/ast.rs
@@ -7,7 +7,7 @@ pub(crate) struct File {
     pub(crate) body: Vec<Item>,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Eq, PartialEq, Clone, Debug)]
 pub(crate) enum Item {
     /// Defines a new macro: name, params, body.
     MacroDef(String, Vec<String>, Vec<Item>),
@@ -34,7 +34,7 @@ pub(crate) enum Item {
     Bytes(Vec<u8>),
 }
 
-#[derive(Clone, Debug)]
+#[derive(Eq, PartialEq, Clone, Debug)]
 pub(crate) enum StackReplacement {
     /// Can be either a named item or a label.
     Identifier(String),
diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index 2b162792..5daf541f 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -2,39 +2,32 @@ use ethereum_types::U256;
 use Item::{Push, StandardOp};
 use PushTarget::Literal;
 
-use crate::cpu::kernel::ast::Item::LocalLabelDeclaration;
+use crate::cpu::kernel::ast::Item::{GlobalLabelDeclaration, LocalLabelDeclaration};
 use crate::cpu::kernel::ast::PushTarget::Label;
 use crate::cpu::kernel::ast::{Item, PushTarget};
 use crate::cpu::kernel::utils::replace_windows;
 
 pub(crate) fn optimize_asm(code: &mut Vec<Item>) {
-    constant_propagation(code);
-
-    // Remove no-op jumps: [PUSH label, JUMP, label:] -> [label:]
-    replace_windows(code, |window| {
-        if let [Push(Label(l1)), StandardOp(jump), LocalLabelDeclaration(l2)] = window
-            && l1 == l2
-            && &jump == "JUMP"
-        {
-            Some(vec![LocalLabelDeclaration(l2)])
-        } else {
-            None
+    // Run the optimizer until nothing changes.
+    loop {
+        let old_code = code.clone();
+        optimize_asm_once(code);
+        if code == &old_code {
+            break;
         }
-    });
-
-    // Remove swaps: [PUSH x, PUSH y, SWAP1] -> [PUSH y, PUSH x]
-    replace_windows(code, |window| {
-        if let [Push(Literal(x)), Push(Literal(y)), StandardOp(swap1)] = window
-                && &swap1 == "SWAP1" {
-            Some(vec![Push(Literal(y)), Push(Literal(x))])
-        } else {
-            None
-        }
-    });
+    }
 }
 
+/// A single optimization pass.
+fn optimize_asm_once(code: &mut Vec<Item>) {
+    constant_propagation(code);
+    no_op_jumps(code);
+    remove_swaps(code);
+}
+
+/// Constant propagation.
 fn constant_propagation(code: &mut Vec<Item>) {
-    // Constant propagation for unary ops: [PUSH x, UNARYOP] -> [PUSH UNARYOP(x)]
+    // Constant propagation for unary ops: `[PUSH x, UNARYOP] -> [PUSH UNARYOP(x)]`
     replace_windows(code, |window| {
         if let [Push(Literal(x)), StandardOp(op)] = window {
             match op.as_str() {
@@ -51,7 +44,7 @@ fn constant_propagation(code: &mut Vec<Item>) {
         }
     });
 
-    // Constant propagation for binary ops: [PUSH x, PUSH y, BINOP] -> [PUSH BINOP(x, y)]
+    // Constant propagation for binary ops: `[PUSH x, PUSH y, BINOP] -> [PUSH BINOP(x, y)]`
     replace_windows(code, |window| {
         if let [Push(Literal(x)), Push(Literal(y)), StandardOp(op)] = window {
             match op.as_str() {
@@ -67,5 +60,76 @@ fn constant_propagation(code: &mut Vec<Item>) {
     });
 }
 
+/// Remove no-op jumps: `[PUSH label, JUMP, label:] -> [label:]`
+fn no_op_jumps(code: &mut Vec<Item>) {
+    replace_windows(code, |window| {
+        if let [Push(Label(l)), StandardOp(jump), decl] = window
+            && &jump == "JUMP"
+            && (decl == LocalLabelDeclaration(l.clone()) || decl == GlobalLabelDeclaration(l.clone()))
+        {
+            Some(vec![LocalLabelDeclaration(l)])
+        } else {
+            None
+        }
+    });
+}
+
+/// Remove swaps: `[PUSH x, PUSH y, SWAP1] -> [PUSH y, PUSH x]`
+fn remove_swaps(code: &mut Vec<Item>) {
+    replace_windows(code, |window| {
+        if let [Push(x), Push(y), StandardOp(swap1)] = window
+            && &swap1 == "SWAP1" {
+            Some(vec![Push(y), Push(x)])
+        } else {
+            None
+        }
+    });
+}
+
 #[cfg(test)]
-mod tests {}
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_constant_propagation_iszero() {
+        let mut code = vec![Push(Literal(3.into())), StandardOp("ISZERO".into())];
+        constant_propagation(&mut code);
+        assert_eq!(code, vec![Push(Literal(0.into()))]);
+    }
+
+    #[test]
+    fn test_constant_propagation_mul() {
+        let mut code = vec![
+            Push(Literal(3.into())),
+            Push(Literal(4.into())),
+            StandardOp("MUL".into()),
+        ];
+        constant_propagation(&mut code);
+        assert_eq!(code, vec![Push(Literal(12.into()))]);
+    }
+
+    #[test]
+    fn test_no_op_jump() {
+        let mut code = vec![
+            Push(Label("mylabel".into())),
+            StandardOp("JUMP".into()),
+            LocalLabelDeclaration("mylabel".into()),
+        ];
+        no_op_jumps(&mut code);
+        assert_eq!(code, vec![LocalLabelDeclaration("mylabel".into())]);
+    }
+
+    #[test]
+    fn test_remove_swap() {
+        let mut code = vec![
+            Push(Literal("42".into())),
+            Push(Label("mylabel".into())),
+            StandardOp("SWAP1".into()),
+        ];
+        remove_swaps(&mut code);
+        assert_eq!(
+            code,
+            vec![Push(Label("mylabel".into())), Push(Literal("42".into()))]
+        );
+    }
+}