Some simple optimization rules

Depends on #647.

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)]

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
 

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

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
 }

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

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()),
     }
 }

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]);