Check if suggested code is actually better

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
+}

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;

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]