From 497b26dee6a219b7b28df037fd35c59708d07400 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Mon, 25 Jul 2022 09:36:26 -0700
Subject: [PATCH 01/14] Some simple optimization rules

Depends on #647.
---
 evm/src/cpu/kernel/aggregator.rs             |  2 +-
 evm/src/cpu/kernel/asm/util/basic_macros.asm |  8 ++-
 evm/src/cpu/kernel/assembler.rs              | 22 ++++--
 evm/src/cpu/kernel/mod.rs                    |  3 +-
 evm/src/cpu/kernel/optimizer.rs              | 71 ++++++++++++++++++++
 evm/src/cpu/kernel/parser.rs                 |  2 +-
 evm/src/cpu/kernel/utils.rs                  | 39 +++++++++++
 7 files changed, 135 insertions(+), 12 deletions(-)
 create mode 100644 evm/src/cpu/kernel/optimizer.rs

diff --git a/evm/src/cpu/kernel/aggregator.rs b/evm/src/cpu/kernel/aggregator.rs
index e5e1f29f..73d1797b 100644
--- a/evm/src/cpu/kernel/aggregator.rs
+++ b/evm/src/cpu/kernel/aggregator.rs
@@ -74,7 +74,7 @@ pub(crate) fn combined_kernel() -> Kernel {
     ];
 
     let parsed_files = files.iter().map(|f| parse(f)).collect_vec();
-    assemble(parsed_files, evm_constants())
+    assemble(parsed_files, evm_constants(), true)
 }
 
 #[cfg(test)]
diff --git a/evm/src/cpu/kernel/asm/util/basic_macros.asm b/evm/src/cpu/kernel/asm/util/basic_macros.asm
index e266b2cb..58ab5d2c 100644
--- a/evm/src/cpu/kernel/asm/util/basic_macros.asm
+++ b/evm/src/cpu/kernel/asm/util/basic_macros.asm
@@ -120,7 +120,9 @@
     // stack: input, ...
     PUSH $c
     // stack: c, input, ...
-    GE // Check it backwards: (input <= c) == (c >= input)
+    %add_const(1) // This will be optimized out.
+    // stack: c + 1, input, ...
+    GE // Check it backwards: (input <= c) == (c + 1 > input)
     // stack: input <= c, ...
 %endmacro
 
@@ -136,7 +138,9 @@
     // stack: input, ...
     PUSH $c
     // stack: c, input, ...
-    LE // Check it backwards: (input >= c) == (c <= input)
+    %sub_const(1) // This will be optimized out.
+    // stack: c - 1, input, ...
+    LT // Check it backwards: (input >= c) == (c - 1 < input)
     // stack: input >= c, ...
 %endmacro
 
diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs
index 6e98b22c..0a1232e1 100644
--- a/evm/src/cpu/kernel/assembler.rs
+++ b/evm/src/cpu/kernel/assembler.rs
@@ -7,6 +7,7 @@ use log::debug;
 use super::ast::PushTarget;
 use crate::cpu::kernel::ast::StackReplacement;
 use crate::cpu::kernel::keccak_util::hash_kernel;
+use crate::cpu::kernel::optimizer::optimize_asm;
 use crate::cpu::kernel::prover_input::ProverInputFn;
 use crate::cpu::kernel::stack_manipulation::expand_stack_manipulation;
 use crate::cpu::kernel::utils::u256_to_trimmed_be_bytes;
@@ -64,7 +65,11 @@ impl Macro {
     }
 }
 
-pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Kernel {
+pub(crate) fn assemble(
+    files: Vec<File>,
+    constants: HashMap<String, U256>,
+    optimize: bool,
+) -> Kernel {
     let macros = find_macros(&files);
     let mut global_labels = HashMap::new();
     let mut prover_inputs = HashMap::new();
@@ -75,7 +80,10 @@ 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);
+        let mut expanded_file = expand_stack_manipulation(expanded_file);
+        if optimize {
+            optimize_asm(&mut expanded_file);
+        }
         local_labels.push(find_labels(
             &expanded_file,
             &mut offset,
@@ -381,7 +389,7 @@ mod tests {
         let expected_kernel = Kernel::new(expected_code, expected_global_labels, HashMap::new());
 
         let program = vec![file_1, file_2];
-        assert_eq!(assemble(program, HashMap::new()), expected_kernel);
+        assert_eq!(assemble(program, HashMap::new(), false), expected_kernel);
     }
 
     #[test]
@@ -399,7 +407,7 @@ mod tests {
                 Item::StandardOp("JUMPDEST".to_string()),
             ],
         };
-        assemble(vec![file_1, file_2], HashMap::new());
+        assemble(vec![file_1, file_2], HashMap::new(), false);
     }
 
     #[test]
@@ -413,7 +421,7 @@ mod tests {
                 Item::StandardOp("ADD".to_string()),
             ],
         };
-        assemble(vec![file], HashMap::new());
+        assemble(vec![file], HashMap::new(), false);
     }
 
     #[test]
@@ -421,7 +429,7 @@ mod tests {
         let file = File {
             body: vec![Item::Bytes(vec![0x12, 42]), Item::Bytes(vec![0xFE, 255])],
         };
-        let code = assemble(vec![file], HashMap::new()).code;
+        let code = assemble(vec![file], HashMap::new(), false).code;
         assert_eq!(code, vec![0x12, 42, 0xfe, 255]);
     }
 
@@ -530,6 +538,6 @@ mod tests {
         constants: HashMap<String, U256>,
     ) -> Kernel {
         let parsed_files = files.iter().map(|f| parse(f)).collect_vec();
-        assemble(parsed_files, constants)
+        assemble(parsed_files, constants, false)
     }
 }
diff --git a/evm/src/cpu/kernel/mod.rs b/evm/src/cpu/kernel/mod.rs
index f0247f93..d87c1e13 100644
--- a/evm/src/cpu/kernel/mod.rs
+++ b/evm/src/cpu/kernel/mod.rs
@@ -3,6 +3,7 @@ pub mod assembler;
 mod ast;
 pub(crate) mod keccak_util;
 mod opcodes;
+mod optimizer;
 mod parser;
 pub mod prover_input;
 mod stack_manipulation;
@@ -23,6 +24,6 @@ use crate::cpu::kernel::aggregator::evm_constants;
 /// This is for debugging the kernel only.
 pub fn assemble_to_bytes(files: &[String]) -> Vec<u8> {
     let parsed_files: Vec<_> = files.iter().map(|f| parse(f)).collect();
-    let kernel = assemble(parsed_files, evm_constants());
+    let kernel = assemble(parsed_files, evm_constants(), true);
     kernel.code
 }
diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
new file mode 100644
index 00000000..2b162792
--- /dev/null
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -0,0 +1,71 @@
+use ethereum_types::U256;
+use Item::{Push, StandardOp};
+use PushTarget::Literal;
+
+use crate::cpu::kernel::ast::Item::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
+        }
+    });
+
+    // 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
+        }
+    });
+}
+
+fn constant_propagation(code: &mut Vec<Item>) {
+    // 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() {
+                "ISZERO" => Some(vec![Push(Literal(if x.is_zero() {
+                    U256::one()
+                } else {
+                    U256::zero()
+                }))]),
+                "NOT" => Some(vec![Push(Literal(!x))]),
+                _ => None,
+            }
+        } else {
+            None
+        }
+    });
+
+    // 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() {
+                "ADD" => Some(vec![Push(Literal(x + y))]),
+                "SUB" => Some(vec![Push(Literal(x - y))]),
+                "MUL" => Some(vec![Push(Literal(x * y))]),
+                "DIV" => Some(vec![Push(Literal(x / y))]),
+                _ => None,
+            }
+        } else {
+            None
+        }
+    });
+}
+
+#[cfg(test)]
+mod tests {}
diff --git a/evm/src/cpu/kernel/parser.rs b/evm/src/cpu/kernel/parser.rs
index 860dc19d..66bf0757 100644
--- a/evm/src/cpu/kernel/parser.rs
+++ b/evm/src/cpu/kernel/parser.rs
@@ -45,7 +45,7 @@ fn parse_item(item: Pair<Rule>) -> Item {
                 .collect::<Vec<_>>()
                 .into(),
         ),
-        Rule::nullary_instruction => Item::StandardOp(item.as_str().into()),
+        Rule::nullary_instruction => Item::StandardOp(item.as_str().to_uppercase()),
         _ => panic!("Unexpected {:?}", item.as_rule()),
     }
 }
diff --git a/evm/src/cpu/kernel/utils.rs b/evm/src/cpu/kernel/utils.rs
index d9682679..cff10430 100644
--- a/evm/src/cpu/kernel/utils.rs
+++ b/evm/src/cpu/kernel/utils.rs
@@ -1,6 +1,30 @@
+use std::fmt::Debug;
+
 use ethereum_types::U256;
 use plonky2_util::ceil_div_usize;
 
+/// Enumerate the length `W` windows of `vec`, and run `maybe_replace` on each one.
+///
+/// Whenever `maybe_replace` returns `Some(replacement)`, the given replacement will be applied.
+pub(crate) fn replace_windows<const W: usize, T, F>(vec: &mut Vec<T>, maybe_replace: F)
+where
+    T: Clone + Debug,
+    F: Fn([T; W]) -> Option<Vec<T>>,
+{
+    let mut start = 0;
+    while start + W <= vec.len() {
+        let range = start..start + W;
+        let window = vec[range.clone()].to_vec().try_into().unwrap();
+        if let Some(replacement) = maybe_replace(window) {
+            vec.splice(range, replacement);
+            // Go back to the earliest window that changed.
+            start = start.saturating_sub(W - 1);
+        } else {
+            start += 1;
+        }
+    }
+}
+
 pub(crate) fn u256_to_trimmed_be_bytes(u256: &U256) -> Vec<u8> {
     let num_bytes = ceil_div_usize(u256.bits(), 8).max(1);
     // `byte` is little-endian, so we manually reverse it.
@@ -11,6 +35,21 @@ pub(crate) fn u256_to_trimmed_be_bytes(u256: &U256) -> Vec<u8> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_replace_windows() {
+        // This replacement function adds pairs of integers together.
+        let mut vec = vec![1, 2, 3, 4, 5];
+        replace_windows(&mut vec, |[x, y]| Some(vec![x + y]));
+        assert_eq!(vec, vec![15u32]);
+
+        // This replacement function splits each composite integer into two factors.
+        let mut vec = vec![9, 1, 6, 8, 15, 7, 9];
+        replace_windows(&mut vec, |[n]| {
+            (2..n).filter(|d| n % d == 0).next().map(|d| vec![d, n / d])
+        });
+        assert_eq!(vec, vec![3, 3, 1, 2, 3, 2, 2, 2, 3, 5, 7, 3, 3]);
+    }
+
     #[test]
     fn literal_to_be_bytes() {
         assert_eq!(u256_to_trimmed_be_bytes(&0.into()), vec![0x00]);

From a34a4c81842ff459920fa972dae95113da7eac06 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sun, 31 Jul 2022 13:03:07 -0700
Subject: [PATCH 02/14] fix

---
 evm/src/cpu/kernel/asm/util/basic_macros.asm | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/evm/src/cpu/kernel/asm/util/basic_macros.asm b/evm/src/cpu/kernel/asm/util/basic_macros.asm
index 58ab5d2c..a91feb05 100644
--- a/evm/src/cpu/kernel/asm/util/basic_macros.asm
+++ b/evm/src/cpu/kernel/asm/util/basic_macros.asm
@@ -122,7 +122,7 @@
     // stack: c, input, ...
     %add_const(1) // This will be optimized out.
     // stack: c + 1, input, ...
-    GE // Check it backwards: (input <= c) == (c + 1 > input)
+    GT // Check it backwards: (input <= c) == (c + 1 > input)
     // stack: input <= c, ...
 %endmacro
 

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 03/14] 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()))]
+        );
+    }
+}

From d639d0b31129ac9bae9d3594f47563e1668cc3fe Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sun, 31 Jul 2022 15:55:02 -0700
Subject: [PATCH 04/14] clippy

---
 evm/src/cpu/kernel/utils.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/evm/src/cpu/kernel/utils.rs b/evm/src/cpu/kernel/utils.rs
index cff10430..e2156259 100644
--- a/evm/src/cpu/kernel/utils.rs
+++ b/evm/src/cpu/kernel/utils.rs
@@ -45,7 +45,7 @@ mod tests {
         // This replacement function splits each composite integer into two factors.
         let mut vec = vec![9, 1, 6, 8, 15, 7, 9];
         replace_windows(&mut vec, |[n]| {
-            (2..n).filter(|d| n % d == 0).next().map(|d| vec![d, n / d])
+            (2..n).find(|d| n % d == 0).map(|d| vec![d, n / d])
         });
         assert_eq!(vec, vec![3, 3, 1, 2, 3, 2, 2, 2, 3, 5, 7, 3, 3]);
     }

From 61a9839f2fe78337d2d1ead9bfb0343dcb6647a7 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sun, 31 Jul 2022 20:37:05 -0700
Subject: [PATCH 05/14] Feedback

---
 evm/src/cpu/kernel/optimizer.rs | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index 5daf541f..9683e160 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -44,14 +44,16 @@ 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 y, PUSH x, BINOP] -> [PUSH BINOP(x, y)]`
     replace_windows(code, |window| {
-        if let [Push(Literal(x)), Push(Literal(y)), StandardOp(op)] = window {
+        if let [Push(Literal(y)), Push(Literal(x)), StandardOp(op)] = window {
             match op.as_str() {
                 "ADD" => Some(vec![Push(Literal(x + y))]),
                 "SUB" => Some(vec![Push(Literal(x - y))]),
                 "MUL" => Some(vec![Push(Literal(x * y))]),
-                "DIV" => Some(vec![Push(Literal(x / y))]),
+                "DIV" => Some(vec![Push(Literal(
+                    x.checked_div(y).unwrap_or(U256::zero()),
+                ))]),
                 _ => None,
             }
         } else {
@@ -108,6 +110,28 @@ mod tests {
         assert_eq!(code, vec![Push(Literal(12.into()))]);
     }
 
+    #[test]
+    fn test_constant_propagation_div() {
+        let mut code = vec![
+            Push(Literal(3.into())),
+            Push(Literal(8.into())),
+            StandardOp("DIV".into()),
+        ];
+        constant_propagation(&mut code);
+        assert_eq!(code, vec![Push(Literal(2.into()))]);
+    }
+
+    #[test]
+    fn test_constant_propagation_div_zero() {
+        let mut code = vec![
+            Push(Literal(0.into())),
+            Push(Literal(1.into())),
+            StandardOp("DIV".into()),
+        ];
+        constant_propagation(&mut code);
+        assert_eq!(code, vec![Push(Literal(0.into()))]);
+    }
+
     #[test]
     fn test_no_op_jump() {
         let mut code = vec![

From 94183f723c57a434e7f9b5465464357d7fa682cb Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sun, 31 Jul 2022 20:43:58 -0700
Subject: [PATCH 06/14] fixes

---
 evm/src/cpu/kernel/optimizer.rs | 28 +++++++++++++++++++++++++---
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index 9683e160..66307aec 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -48,9 +48,9 @@ fn constant_propagation(code: &mut Vec<Item>) {
     replace_windows(code, |window| {
         if let [Push(Literal(y)), Push(Literal(x)), StandardOp(op)] = window {
             match op.as_str() {
-                "ADD" => Some(vec![Push(Literal(x + y))]),
-                "SUB" => Some(vec![Push(Literal(x - y))]),
-                "MUL" => Some(vec![Push(Literal(x * y))]),
+                "ADD" => Some(vec![Push(Literal(x.overflowing_add(y).0))]),
+                "SUB" => Some(vec![Push(Literal(x.overflowing_sub(y).0))]),
+                "MUL" => Some(vec![Push(Literal(x.overflowing_mul(y).0))]),
                 "DIV" => Some(vec![Push(Literal(
                     x.checked_div(y).unwrap_or(U256::zero()),
                 ))]),
@@ -99,6 +99,28 @@ mod tests {
         assert_eq!(code, vec![Push(Literal(0.into()))]);
     }
 
+    #[test]
+    fn test_constant_propagation_add_overflowing() {
+        let mut code = vec![
+            Push(Literal(U256::max_value())),
+            Push(Literal(U256::max_value())),
+            StandardOp("ADD".into()),
+        ];
+        constant_propagation(&mut code);
+        assert_eq!(code, vec![Push(Literal(U256::max_value() - 1))]);
+    }
+
+    #[test]
+    fn test_constant_propagation_sub_underflowing() {
+        let mut code = vec![
+            Push(Literal(U256::one())),
+            Push(Literal(U256::zero())),
+            StandardOp("SUB".into()),
+        ];
+        constant_propagation(&mut code);
+        assert_eq!(code, vec![Push(Literal(U256::max_value()))]);
+    }
+
     #[test]
     fn test_constant_propagation_mul() {
         let mut code = vec![

From 63c8568b17e067eae0b254cf0c7d7f8ccbaa8dd2 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Mon, 1 Aug 2022 08:47:07 -0700
Subject: [PATCH 07/14] remove_ignored_values

---
 evm/src/cpu/kernel/optimizer.rs | 34 +++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index 66307aec..cc595191 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -23,6 +23,7 @@ fn optimize_asm_once(code: &mut Vec<Item>) {
     constant_propagation(code);
     no_op_jumps(code);
     remove_swaps(code);
+    remove_ignored_values(code);
 }
 
 /// Constant propagation.
@@ -62,7 +63,7 @@ fn constant_propagation(code: &mut Vec<Item>) {
     });
 }
 
-/// Remove no-op jumps: `[PUSH label, JUMP, label:] -> [label:]`
+/// 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
@@ -76,7 +77,7 @@ fn no_op_jumps(code: &mut Vec<Item>) {
     });
 }
 
-/// Remove swaps: `[PUSH x, PUSH y, SWAP1] -> [PUSH y, PUSH x]`
+/// 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
@@ -88,6 +89,21 @@ fn remove_swaps(code: &mut Vec<Item>) {
     });
 }
 
+/// Remove push-pop type patterns, such as: `[DUP1, POP]`.
+fn remove_ignored_values(code: &mut Vec<Item>) {
+    replace_windows(code, |[a, b]| {
+        if let StandardOp(pop) = b && &pop == "POP" {
+            match a {
+                Push(_) => Some(vec![]),
+                StandardOp(dup) if dup.starts_with("DUP") => Some(vec![]),
+                _ => None,
+            }
+        } else {
+            None
+        }
+    });
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -178,4 +194,18 @@ mod tests {
             vec![Push(Label("mylabel".into())), Push(Literal("42".into()))]
         );
     }
+
+    #[test]
+    fn test_remove_push_pop() {
+        let mut code = vec![Push(Literal("42".into())), StandardOp("POP".into())];
+        remove_ignored_values(&mut code);
+        assert_eq!(code, vec![]);
+    }
+
+    #[test]
+    fn test_remove_dup_pop() {
+        let mut code = vec![StandardOp("DUP5".into()), StandardOp("POP".into())];
+        remove_ignored_values(&mut code);
+        assert_eq!(code, vec![]);
+    }
 }

From d1c9e150b37e9b8483d321631e24967823c2f3e1 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Mon, 1 Aug 2022 09:02:39 -0700
Subject: [PATCH 08/14] remove_swaps_commutative

---
 evm/src/cpu/kernel/optimizer.rs | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index cc595191..811454e1 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -22,7 +22,8 @@ pub(crate) fn optimize_asm(code: &mut Vec<Item>) {
 fn optimize_asm_once(code: &mut Vec<Item>) {
     constant_propagation(code);
     no_op_jumps(code);
-    remove_swaps(code);
+    remove_swapped_pushes(code);
+    remove_swaps_commutative(code);
     remove_ignored_values(code);
 }
 
@@ -78,7 +79,7 @@ fn no_op_jumps(code: &mut Vec<Item>) {
 }
 
 /// Remove swaps: `[PUSH x, PUSH y, SWAP1] -> [PUSH y, PUSH x]`.
-fn remove_swaps(code: &mut Vec<Item>) {
+fn remove_swapped_pushes(code: &mut Vec<Item>) {
     replace_windows(code, |window| {
         if let [Push(x), Push(y), StandardOp(swap1)] = window
             && &swap1 == "SWAP1" {
@@ -89,6 +90,22 @@ fn remove_swaps(code: &mut Vec<Item>) {
     });
 }
 
+/// Remove SWAP1 before a commutative function.
+fn remove_swaps_commutative(code: &mut Vec<Item>) {
+    replace_windows(code, |window| {
+        if let [StandardOp(swap1), StandardOp(f)] = window && &swap1 == "SWAP1" {
+            let commutative = match f.as_str() {
+                "ADD" => true,
+                "MUL" => true,
+                _ => false,
+            };
+            commutative.then_some(vec![StandardOp(f)])
+        } else {
+            None
+        }
+    });
+}
+
 /// Remove push-pop type patterns, such as: `[DUP1, POP]`.
 fn remove_ignored_values(code: &mut Vec<Item>) {
     replace_windows(code, |[a, b]| {
@@ -182,19 +199,26 @@ mod tests {
     }
 
     #[test]
-    fn test_remove_swap() {
+    fn test_remove_swapped_pushes() {
         let mut code = vec![
             Push(Literal("42".into())),
             Push(Label("mylabel".into())),
             StandardOp("SWAP1".into()),
         ];
-        remove_swaps(&mut code);
+        remove_swapped_pushes(&mut code);
         assert_eq!(
             code,
             vec![Push(Label("mylabel".into())), Push(Literal("42".into()))]
         );
     }
 
+    #[test]
+    fn test_remove_swap_mul() {
+        let mut code = vec![StandardOp("SWAP1".into()), StandardOp("MUL".into())];
+        remove_swaps_commutative(&mut code);
+        assert_eq!(code, vec![StandardOp("MUL".into())]);
+    }
+
     #[test]
     fn test_remove_push_pop() {
         let mut code = vec![Push(Literal("42".into())), StandardOp("POP".into())];

From 98c4a372fb784351b19b80ccf9c235a893b3812e Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Mon, 1 Aug 2022 09:22:01 -0700
Subject: [PATCH 09/14] More binops

---
 evm/src/cpu/kernel/optimizer.rs | 33 ++++++++++++++++-----------------
 evm/src/cpu/kernel/utils.rs     |  8 ++++++++
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index 811454e1..6d77f0fa 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -5,7 +5,7 @@ use PushTarget::Literal;
 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;
+use crate::cpu::kernel::utils::{replace_windows, u256_from_bool};
 
 pub(crate) fn optimize_asm(code: &mut Vec<Item>) {
     // Run the optimizer until nothing changes.
@@ -33,11 +33,7 @@ fn constant_propagation(code: &mut Vec<Item>) {
     replace_windows(code, |window| {
         if let [Push(Literal(x)), StandardOp(op)] = window {
             match op.as_str() {
-                "ISZERO" => Some(vec![Push(Literal(if x.is_zero() {
-                    U256::one()
-                } else {
-                    U256::zero()
-                }))]),
+                "ISZERO" => Some(vec![Push(Literal(u256_from_bool(x.is_zero())))]),
                 "NOT" => Some(vec![Push(Literal(!x))]),
                 _ => None,
             }
@@ -50,14 +46,21 @@ fn constant_propagation(code: &mut Vec<Item>) {
     replace_windows(code, |window| {
         if let [Push(Literal(y)), Push(Literal(x)), StandardOp(op)] = window {
             match op.as_str() {
-                "ADD" => Some(vec![Push(Literal(x.overflowing_add(y).0))]),
-                "SUB" => Some(vec![Push(Literal(x.overflowing_sub(y).0))]),
-                "MUL" => Some(vec![Push(Literal(x.overflowing_mul(y).0))]),
-                "DIV" => Some(vec![Push(Literal(
-                    x.checked_div(y).unwrap_or(U256::zero()),
-                ))]),
+                "ADD" => Some(x.overflowing_add(y).0),
+                "SUB" => Some(x.overflowing_sub(y).0),
+                "MUL" => Some(x.overflowing_mul(y).0),
+                "DIV" => Some(x.checked_div(y).unwrap_or(U256::zero())),
+                "SHL" => Some(x << y),
+                "SHR" => Some(x >> y),
+                "AND" => Some(x & y),
+                "OR" => Some(x | y),
+                "XOR" => Some(x ^ y),
+                "LT" => Some(u256_from_bool(x < y)),
+                "GT" => Some(u256_from_bool(x > y)),
+                "EQ" => Some(u256_from_bool(x == y)),
                 _ => None,
             }
+            .map(|res| vec![Push(Literal(res))])
         } else {
             None
         }
@@ -94,11 +97,7 @@ fn remove_swapped_pushes(code: &mut Vec<Item>) {
 fn remove_swaps_commutative(code: &mut Vec<Item>) {
     replace_windows(code, |window| {
         if let [StandardOp(swap1), StandardOp(f)] = window && &swap1 == "SWAP1" {
-            let commutative = match f.as_str() {
-                "ADD" => true,
-                "MUL" => true,
-                _ => false,
-            };
+            let commutative = matches!(f.as_str(), "ADD" | "MUL");
             commutative.then_some(vec![StandardOp(f)])
         } else {
             None
diff --git a/evm/src/cpu/kernel/utils.rs b/evm/src/cpu/kernel/utils.rs
index e2156259..8900b8e2 100644
--- a/evm/src/cpu/kernel/utils.rs
+++ b/evm/src/cpu/kernel/utils.rs
@@ -31,6 +31,14 @@ pub(crate) fn u256_to_trimmed_be_bytes(u256: &U256) -> Vec<u8> {
     (0..num_bytes).rev().map(|i| u256.byte(i)).collect()
 }
 
+pub(crate) fn u256_from_bool(b: bool) -> U256 {
+    if b {
+        U256::one()
+    } else {
+        U256::zero()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

From cb2df9fa03dd334e0e3373d2fff857893f084da5 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Tue, 2 Aug 2022 09:22:06 -0700
Subject: [PATCH 10/14] More commutative fns

---
 evm/src/cpu/kernel/optimizer.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index 6d77f0fa..ac46011a 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -97,7 +97,7 @@ fn remove_swapped_pushes(code: &mut Vec<Item>) {
 fn remove_swaps_commutative(code: &mut Vec<Item>) {
     replace_windows(code, |window| {
         if let [StandardOp(swap1), StandardOp(f)] = window && &swap1 == "SWAP1" {
-            let commutative = matches!(f.as_str(), "ADD" | "MUL");
+            let commutative = matches!(f.as_str(), "ADD" | "MUL" | "AND" | "OR" | "XOR" | "EQ");
             commutative.then_some(vec![StandardOp(f)])
         } else {
             None

From 6416826643ae906bd1ec0b59510e0922e092d6b6 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Tue, 2 Aug 2022 15:44:50 -0700
Subject: [PATCH 11/14] Feedback

---
 evm/src/cpu/kernel/asm/util/basic_macros.asm | 8 ++------
 evm/src/cpu/kernel/optimizer.rs              | 7 +++++++
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/util/basic_macros.asm b/evm/src/cpu/kernel/asm/util/basic_macros.asm
index a91feb05..4dd93d14 100644
--- a/evm/src/cpu/kernel/asm/util/basic_macros.asm
+++ b/evm/src/cpu/kernel/asm/util/basic_macros.asm
@@ -120,9 +120,7 @@
     // stack: input, ...
     PUSH $c
     // stack: c, input, ...
-    %add_const(1) // This will be optimized out.
-    // stack: c + 1, input, ...
-    GT // Check it backwards: (input <= c) == (c + 1 > input)
+    LT ISZERO // Check it backwards: (input <= c) == !(c < input)
     // stack: input <= c, ...
 %endmacro
 
@@ -138,9 +136,7 @@
     // stack: input, ...
     PUSH $c
     // stack: c, input, ...
-    %sub_const(1) // This will be optimized out.
-    // stack: c - 1, input, ...
-    LT // Check it backwards: (input >= c) == (c - 1 < input)
+    GT ISZERO // Check it backwards: (input >= c) == !(c > input)
     // stack: input >= c, ...
 %endmacro
 
diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index ac46011a..a334f2d8 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -50,6 +50,8 @@ fn constant_propagation(code: &mut Vec<Item>) {
                 "SUB" => Some(x.overflowing_sub(y).0),
                 "MUL" => Some(x.overflowing_mul(y).0),
                 "DIV" => Some(x.checked_div(y).unwrap_or(U256::zero())),
+                "MOD" => Some(x.checked_rem(y).unwrap_or(U256::zero())),
+                "EXP" => Some(x.overflowing_pow(y).0),
                 "SHL" => Some(x << y),
                 "SHR" => Some(x >> y),
                 "AND" => Some(x & y),
@@ -58,6 +60,11 @@ fn constant_propagation(code: &mut Vec<Item>) {
                 "LT" => Some(u256_from_bool(x < y)),
                 "GT" => Some(u256_from_bool(x > y)),
                 "EQ" => Some(u256_from_bool(x == y)),
+                "BYTE" => Some(if x < 32.into() {
+                    y.byte(x.as_usize()).into()
+                } else {
+                    U256::zero()
+                }),
                 _ => None,
             }
             .map(|res| vec![Push(Literal(res))])

From 8aad0b07465f6043dc1d9a609b3d3ab2cc410e83 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Tue, 2 Aug 2022 15:57:06 -0700
Subject: [PATCH 12/14] Feedback

---
 evm/src/cpu/kernel/assembler.rs | 2 +-
 evm/src/cpu/kernel/optimizer.rs | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs
index 0a1232e1..4e3381b5 100644
--- a/evm/src/cpu/kernel/assembler.rs
+++ b/evm/src/cpu/kernel/assembler.rs
@@ -538,6 +538,6 @@ mod tests {
         constants: HashMap<String, U256>,
     ) -> Kernel {
         let parsed_files = files.iter().map(|f| parse(f)).collect_vec();
-        assemble(parsed_files, constants, false)
+        assemble(parsed_files, constants, true)
     }
 }
diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index a334f2d8..6fe9d496 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -89,6 +89,7 @@ fn no_op_jumps(code: &mut Vec<Item>) {
 }
 
 /// Remove swaps: `[PUSH x, PUSH y, SWAP1] -> [PUSH y, PUSH x]`.
+// Could be generalized to recognize more than two pushes.
 fn remove_swapped_pushes(code: &mut Vec<Item>) {
     replace_windows(code, |window| {
         if let [Push(x), Push(y), StandardOp(swap1)] = window
@@ -113,6 +114,7 @@ fn remove_swaps_commutative(code: &mut Vec<Item>) {
 }
 
 /// Remove push-pop type patterns, such as: `[DUP1, POP]`.
+// Could be extended to other non-side-effecting operations, e.g. [DUP1, ADD, POP] -> [POP].
 fn remove_ignored_values(code: &mut Vec<Item>) {
     replace_windows(code, |[a, b]| {
         if let StandardOp(pop) = b && &pop == "POP" {

From 9b5b77d3e9d0b6d830c2277e297ba6df23e5e116 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Wed, 3 Aug 2022 09:57:40 -0700
Subject: [PATCH 13/14] Check if suggested code is actually better

---
 evm/src/cpu/kernel/cost_estimator.rs | 37 ++++++++++++++++++++++++++++
 evm/src/cpu/kernel/mod.rs            |  1 +
 evm/src/cpu/kernel/optimizer.rs      | 25 ++++++++++++++++---
 3 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 evm/src/cpu/kernel/cost_estimator.rs

diff --git a/evm/src/cpu/kernel/cost_estimator.rs b/evm/src/cpu/kernel/cost_estimator.rs
new file mode 100644
index 00000000..3dfcf63a
--- /dev/null
+++ b/evm/src/cpu/kernel/cost_estimator.rs
@@ -0,0 +1,37 @@
+use crate::cpu::kernel::assembler::BYTES_PER_OFFSET;
+use crate::cpu::kernel::ast::Item;
+use crate::cpu::kernel::ast::Item::*;
+use crate::cpu::kernel::ast::PushTarget::*;
+use crate::cpu::kernel::utils::u256_to_trimmed_be_bytes;
+
+pub(crate) fn is_code_improved(before: &[Item], after: &[Item]) -> bool {
+    cost_estimate(after) < cost_estimate(before)
+}
+
+fn cost_estimate(code: &[Item]) -> u32 {
+    code.iter().map(cost_estimate_item).sum()
+}
+
+fn cost_estimate_item(item: &Item) -> u32 {
+    match item {
+        MacroDef(_, _, _) => 0,
+        GlobalLabelDeclaration(_) => 0,
+        LocalLabelDeclaration(_) => 0,
+        Push(Literal(n)) => cost_estimate_push(u256_to_trimmed_be_bytes(n).len()),
+        Push(Label(_)) => cost_estimate_push(BYTES_PER_OFFSET as usize),
+        ProverInput(_) => 1,
+        StandardOp(op) => cost_estimate_standard_op(op.as_str()),
+        _ => panic!("Unexpected item: {:?}", item),
+    }
+}
+
+fn cost_estimate_standard_op(_op: &str) -> u32 {
+    // For now we just treat any standard operation as having the same cost. This is pretty naive,
+    // but should work fine with our current set of simple optimization rules.
+    1
+}
+
+fn cost_estimate_push(num_bytes: usize) -> u32 {
+    // TODO: Once PUSH is actually implemented, check if this needs to be revised.
+    num_bytes as u32
+}
diff --git a/evm/src/cpu/kernel/mod.rs b/evm/src/cpu/kernel/mod.rs
index d87c1e13..59caff76 100644
--- a/evm/src/cpu/kernel/mod.rs
+++ b/evm/src/cpu/kernel/mod.rs
@@ -1,6 +1,7 @@
 pub mod aggregator;
 pub mod assembler;
 mod ast;
+mod cost_estimator;
 pub(crate) mod keccak_util;
 mod opcodes;
 mod optimizer;
diff --git a/evm/src/cpu/kernel/optimizer.rs b/evm/src/cpu/kernel/optimizer.rs
index 6fe9d496..2a1db6d3 100644
--- a/evm/src/cpu/kernel/optimizer.rs
+++ b/evm/src/cpu/kernel/optimizer.rs
@@ -5,6 +5,7 @@ use PushTarget::Literal;
 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::cost_estimator::is_code_improved;
 use crate::cpu::kernel::utils::{replace_windows, u256_from_bool};
 
 pub(crate) fn optimize_asm(code: &mut Vec<Item>) {
@@ -30,7 +31,7 @@ fn optimize_asm_once(code: &mut Vec<Item>) {
 /// Constant propagation.
 fn constant_propagation(code: &mut Vec<Item>) {
     // Constant propagation for unary ops: `[PUSH x, UNARYOP] -> [PUSH UNARYOP(x)]`
-    replace_windows(code, |window| {
+    replace_windows_if_better(code, |window| {
         if let [Push(Literal(x)), StandardOp(op)] = window {
             match op.as_str() {
                 "ISZERO" => Some(vec![Push(Literal(u256_from_bool(x.is_zero())))]),
@@ -43,7 +44,7 @@ fn constant_propagation(code: &mut Vec<Item>) {
     });
 
     // Constant propagation for binary ops: `[PUSH y, PUSH x, BINOP] -> [PUSH BINOP(x, y)]`
-    replace_windows(code, |window| {
+    replace_windows_if_better(code, |window| {
         if let [Push(Literal(y)), Push(Literal(x)), StandardOp(op)] = window {
             match op.as_str() {
                 "ADD" => Some(x.overflowing_add(y).0),
@@ -129,6 +130,17 @@ fn remove_ignored_values(code: &mut Vec<Item>) {
     });
 }
 
+/// Like `replace_windows`, but specifically for code, and only makes replacements if our cost
+/// estimator thinks that the new code is more efficient.
+fn replace_windows_if_better<const W: usize, F>(code: &mut Vec<Item>, maybe_replace: F)
+where
+    F: Fn([Item; W]) -> Option<Vec<Item>>,
+{
+    replace_windows(code, |window| {
+        maybe_replace(window.clone()).filter(|suggestion| is_code_improved(&window, suggestion))
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -153,13 +165,18 @@ mod tests {
 
     #[test]
     fn test_constant_propagation_sub_underflowing() {
-        let mut code = vec![
+        let original = vec![
             Push(Literal(U256::one())),
             Push(Literal(U256::zero())),
             StandardOp("SUB".into()),
         ];
+        let mut code = original.clone();
         constant_propagation(&mut code);
-        assert_eq!(code, vec![Push(Literal(U256::max_value()))]);
+        // Constant propagation could replace the code with [PUSH U256::MAX], but that's actually
+        // more expensive, so the code shouldn't be changed.
+        // (The code could also be replaced with [PUSH 0; NOT], which would be an improvement, but
+        // our optimizer isn't smart enough yet.)
+        assert_eq!(code, original);
     }
 
     #[test]

From dfd715fafb73ded564c11f559755943947070952 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Wed, 3 Aug 2022 13:52:52 -0700
Subject: [PATCH 14/14] Fix case where a valid constant propagation a broke
 test

---
 evm/src/cpu/kernel/assembler.rs | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs
index 4e3381b5..636251a3 100644
--- a/evm/src/cpu/kernel/assembler.rs
+++ b/evm/src/cpu/kernel/assembler.rs
@@ -446,10 +446,11 @@ mod tests {
 
     #[test]
     fn macro_with_vars() {
-        let kernel = parse_and_assemble(&[
+        let files = &[
             "%macro add(x, y) PUSH $x PUSH $y ADD %endmacro",
             "%add(2, 3)",
-        ]);
+        ];
+        let kernel = parse_and_assemble_ext(files, HashMap::new(), false);
         let push1 = get_push_opcode(1);
         let add = get_opcode("ADD");
         assert_eq!(kernel.code, vec![push1, 2, push1, 3, add]);
@@ -487,7 +488,7 @@ mod tests {
         let mut constants = HashMap::new();
         constants.insert("DEAD_BEEF".into(), 0xDEADBEEFu64.into());
 
-        let kernel = parse_and_assemble_with_constants(code, constants);
+        let kernel = parse_and_assemble_ext(code, constants, true);
         let push4 = get_push_opcode(4);
         assert_eq!(kernel.code, vec![push4, 0xDE, 0xAD, 0xBE, 0xEF]);
     }
@@ -518,7 +519,7 @@ mod tests {
 
         let mut consts = HashMap::new();
         consts.insert("LIFE".into(), 42.into());
-        parse_and_assemble_with_constants(&["%stack (a, b) -> (b, @LIFE)"], consts);
+        parse_and_assemble_ext(&["%stack (a, b) -> (b, @LIFE)"], consts, true);
         // We won't check the code since there are two equally efficient implementations.
 
         let kernel = parse_and_assemble(&["start: %stack (a, b) -> (start)"]);
@@ -530,14 +531,15 @@ mod tests {
     }
 
     fn parse_and_assemble(files: &[&str]) -> Kernel {
-        parse_and_assemble_with_constants(files, HashMap::new())
+        parse_and_assemble_ext(files, HashMap::new(), true)
     }
 
-    fn parse_and_assemble_with_constants(
+    fn parse_and_assemble_ext(
         files: &[&str],
         constants: HashMap<String, U256>,
+        optimize: bool,
     ) -> Kernel {
         let parsed_files = files.iter().map(|f| parse(f)).collect_vec();
-        assemble(parsed_files, constants, true)
+        assemble(parsed_files, constants, optimize)
     }
 }