Store literals as U256 (or u8 for BYTES)

Instead of the original strings. Will make optimizations simpler.

evm/src/cpu/kernel/assembler.rs

@@ -5,10 +5,11 @@ use itertools::izip;
 use log::debug;
 
 use super::ast::PushTarget;
-use crate::cpu::kernel::ast::{Literal, StackReplacement};
+use crate::cpu::kernel::ast::StackReplacement;
 use crate::cpu::kernel::keccak_util::hash_kernel;
 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;
 use crate::cpu::kernel::{
     ast::{File, Item},
     opcodes::{get_opcode, get_push_opcode},
@@ -184,7 +185,7 @@ fn expand_repeats(body: Vec<Item>) -> Vec<Item> {
     let mut expanded = vec![];
     for item in body {
         if let Item::Repeat(count, block) = item {
-            let reps = count.to_u256().as_usize();
+            let reps = count.as_usize();
             for _ in 0..reps {
                 expanded.extend(block.clone());
             }
@@ -197,12 +198,9 @@ fn expand_repeats(body: Vec<Item>) -> Vec<Item> {
 
 fn inline_constants(body: Vec<Item>, constants: &HashMap<String, U256>) -> Vec<Item> {
     let resolve_const = |c| {
-        Literal::Decimal(
-            constants
-                .get(&c)
-                .unwrap_or_else(|| panic!("No such constant: {}", c))
-                .to_string(),
-        )
+        *constants
+            .get(&c)
+            .unwrap_or_else(|| panic!("No such constant: {}", c))
     };
 
     body.into_iter()
@@ -284,7 +282,7 @@ fn assemble_file(
             }
             Item::Push(target) => {
                 let target_bytes: Vec<u8> = match target {
-                    PushTarget::Literal(literal) => literal.to_trimmed_be_bytes(),
+                    PushTarget::Literal(n) => u256_to_trimmed_be_bytes(&n),
                     PushTarget::Label(label) => {
                         let offset = local_labels
                             .get(&label)
@@ -309,7 +307,7 @@ fn assemble_file(
             Item::StandardOp(opcode) => {
                 code.push(get_opcode(&opcode));
             }
-            Item::Bytes(bytes) => code.extend(bytes.iter().map(|b| b.to_u8())),
+            Item::Bytes(bytes) => code.extend(bytes),
         }
     }
 }
@@ -317,7 +315,7 @@ fn assemble_file(
 /// The size of a `PushTarget`, in bytes.
 fn push_target_size(target: &PushTarget) -> u8 {
     match target {
-        PushTarget::Literal(lit) => lit.to_trimmed_be_bytes().len() as u8,
+        PushTarget::Literal(n) => u256_to_trimmed_be_bytes(n).len() as u8,
         PushTarget::Label(_) => BYTES_PER_OFFSET,
         PushTarget::MacroVar(v) => panic!("Variable not in a macro: {}", v),
         PushTarget::Constant(c) => panic!("Constant wasn't inlined: {}", c),
@@ -421,16 +419,7 @@ mod tests {
     #[test]
     fn literal_bytes() {
         let file = File {
-            body: vec![
-                Item::Bytes(vec![
-                    Literal::Hex("12".to_string()),
-                    Literal::Decimal("42".to_string()),
-                ]),
-                Item::Bytes(vec![
-                    Literal::Hex("fe".to_string()),
-                    Literal::Decimal("255".to_string()),
-                ]),
-            ],
+            body: vec![Item::Bytes(vec![0x12, 42]), Item::Bytes(vec![0xFE, 255])],
         };
         let code = assemble(vec![file], HashMap::new()).code;
         assert_eq!(code, vec![0x12, 42, 0xfe, 255]);

evm/src/cpu/kernel/ast.rs

@@ -1,5 +1,4 @@
 use ethereum_types::U256;
-use plonky2_util::ceil_div_usize;
 
 use crate::cpu::kernel::prover_input::ProverInputFn;
 
@@ -15,7 +14,7 @@ pub(crate) enum Item {
     /// Calls a macro: name, args.
     MacroCall(String, Vec<PushTarget>),
     /// Repetition, like `%rep` in NASM.
-    Repeat(Literal, Vec<Item>),
+    Repeat(U256, 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.
@@ -32,14 +31,14 @@ pub(crate) enum Item {
     /// Any opcode besides a PUSH opcode.
     StandardOp(String),
     /// Literal hex data; should contain an even number of hex chars.
-    Bytes(Vec<Literal>),
+    Bytes(Vec<u8>),
 }
 
 #[derive(Clone, Debug)]
 pub(crate) enum StackReplacement {
     /// Can be either a named item or a label.
     Identifier(String),
-    Literal(Literal),
+    Literal(U256),
     MacroVar(String),
     Constant(String),
 }
@@ -47,69 +46,8 @@ pub(crate) enum StackReplacement {
 /// The target of a `PUSH` operation.
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub(crate) enum PushTarget {
-    Literal(Literal),
+    Literal(U256),
     Label(String),
     MacroVar(String),
     Constant(String),
 }
-
-#[derive(Clone, Debug, Eq, PartialEq, Hash)]
-pub(crate) enum Literal {
-    Decimal(String),
-    Hex(String),
-}
-
-impl Literal {
-    pub(crate) fn to_trimmed_be_bytes(&self) -> Vec<u8> {
-        let u256 = self.to_u256();
-        let num_bytes = ceil_div_usize(u256.bits(), 8).max(1);
-        // `byte` is little-endian, so we manually reverse it.
-        (0..num_bytes).rev().map(|i| u256.byte(i)).collect()
-    }
-
-    pub(crate) fn to_u256(&self) -> U256 {
-        let (src, radix) = match self {
-            Literal::Decimal(s) => (s, 10),
-            Literal::Hex(s) => (s, 16),
-        };
-        U256::from_str_radix(src, radix)
-            .unwrap_or_else(|_| panic!("Not a valid u256 literal: {:?}", self))
-    }
-
-    pub(crate) fn to_u8(&self) -> u8 {
-        let (src, radix) = match self {
-            Literal::Decimal(s) => (s, 10),
-            Literal::Hex(s) => (s, 16),
-        };
-        u8::from_str_radix(src, radix)
-            .unwrap_or_else(|_| panic!("Not a valid u8 literal: {:?}", self))
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use crate::cpu::kernel::ast::*;
-
-    #[test]
-    fn literal_to_be_bytes() {
-        assert_eq!(
-            Literal::Decimal("0".into()).to_trimmed_be_bytes(),
-            vec![0x00]
-        );
-
-        assert_eq!(
-            Literal::Decimal("768".into()).to_trimmed_be_bytes(),
-            vec![0x03, 0x00]
-        );
-
-        assert_eq!(
-            Literal::Hex("a1b2".into()).to_trimmed_be_bytes(),
-            vec![0xa1, 0xb2]
-        );
-
-        assert_eq!(
-            Literal::Hex("1b2".into()).to_trimmed_be_bytes(),
-            vec![0x1, 0xb2]
-        );
-    }
-}

evm/src/cpu/kernel/mod.rs

@@ -7,6 +7,7 @@ mod parser;
 pub mod prover_input;
 mod stack_manipulation;
 mod txn_fields;
+mod utils;
 
 #[cfg(test)]
 mod interpreter;

evm/src/cpu/kernel/parser.rs

@@ -1,7 +1,10 @@
+use std::str::FromStr;
+
+use ethereum_types::U256;
 use pest::iterators::Pair;
 use pest::Parser;
 
-use crate::cpu::kernel::ast::{File, Item, Literal, PushTarget, StackReplacement};
+use crate::cpu::kernel::ast::{File, Item, PushTarget, StackReplacement};
 
 /// Parses EVM assembly code.
 #[derive(pest_derive::Parser)]
@@ -31,7 +34,7 @@ fn parse_item(item: Pair<Rule>) -> Item {
         Rule::local_label => {
             Item::LocalLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
         }
-        Rule::bytes_item => Item::Bytes(item.into_inner().map(parse_literal).collect()),
+        Rule::bytes_item => Item::Bytes(item.into_inner().map(parse_literal_u8).collect()),
         Rule::push_instruction => Item::Push(parse_push_target(item.into_inner().next().unwrap())),
         Rule::prover_input_instruction => Item::ProverInput(
             item.into_inner()
@@ -84,7 +87,7 @@ fn parse_macro_call(item: Pair<Rule>) -> Item {
 fn parse_repeat(item: Pair<Rule>) -> Item {
     assert_eq!(item.as_rule(), Rule::repeat);
     let mut inner = item.into_inner().peekable();
-    let count = parse_literal(inner.next().unwrap());
+    let count = parse_literal_u256(inner.next().unwrap());
     Item::Repeat(count, inner.map(parse_item).collect())
 }
 
@@ -113,7 +116,7 @@ fn parse_stack_replacement(target: Pair<Rule>) -> StackReplacement {
     let inner = target.into_inner().next().unwrap();
     match inner.as_rule() {
         Rule::identifier => StackReplacement::Identifier(inner.as_str().into()),
-        Rule::literal => StackReplacement::Literal(parse_literal(inner)),
+        Rule::literal => StackReplacement::Literal(parse_literal_u256(inner)),
         Rule::variable => {
             StackReplacement::MacroVar(inner.into_inner().next().unwrap().as_str().into())
         }
@@ -128,7 +131,7 @@ fn parse_push_target(target: Pair<Rule>) -> PushTarget {
     assert_eq!(target.as_rule(), Rule::push_target);
     let inner = target.into_inner().next().unwrap();
     match inner.as_rule() {
-        Rule::literal => PushTarget::Literal(parse_literal(inner)),
+        Rule::literal => PushTarget::Literal(parse_literal_u256(inner)),
         Rule::identifier => PushTarget::Label(inner.as_str().into()),
         Rule::variable => PushTarget::MacroVar(inner.into_inner().next().unwrap().as_str().into()),
         Rule::constant => PushTarget::Constant(inner.into_inner().next().unwrap().as_str().into()),
@@ -136,11 +139,28 @@ fn parse_push_target(target: Pair<Rule>) -> PushTarget {
     }
 }
 
-fn parse_literal(literal: Pair<Rule>) -> Literal {
+fn parse_literal_u8(literal: Pair<Rule>) -> u8 {
     let literal = literal.into_inner().next().unwrap();
     match literal.as_rule() {
-        Rule::literal_decimal => Literal::Decimal(literal.as_str().into()),
-        Rule::literal_hex => Literal::Hex(parse_hex(literal)),
+        Rule::literal_decimal => {
+            u8::from_str(literal.as_str()).expect("Failed to parse literal decimal byte")
+        }
+        Rule::literal_hex => {
+            u8::from_str_radix(&parse_hex(literal), 16).expect("Failed to parse literal hex byte")
+        }
+        _ => panic!("Unexpected {:?}", literal.as_rule()),
+    }
+}
+
+fn parse_literal_u256(literal: Pair<Rule>) -> U256 {
+    let literal = literal.into_inner().next().unwrap();
+    match literal.as_rule() {
+        Rule::literal_decimal => {
+            U256::from_dec_str(literal.as_str()).expect("Failed to parse literal decimal")
+        }
+        Rule::literal_hex => {
+            U256::from_str_radix(&parse_hex(literal), 16).expect("Failed to parse literal hex")
+        }
         _ => panic!("Unexpected {:?}", literal.as_rule()),
     }
 }

evm/src/cpu/kernel/stack_manipulation.rs

@@ -8,6 +8,7 @@ use crate::cpu::columns::NUM_CPU_COLUMNS;
 use crate::cpu::kernel::assembler::BYTES_PER_OFFSET;
 use crate::cpu::kernel::ast::{Item, PushTarget, StackReplacement};
 use crate::cpu::kernel::stack_manipulation::StackOp::Pop;
+use crate::cpu::kernel::utils::u256_to_trimmed_be_bytes;
 use crate::memory;
 
 pub(crate) fn expand_stack_manipulation(body: Vec<Item>) -> Vec<Item> {
@@ -227,7 +228,7 @@ impl StackOp {
         let (cpu_rows, memory_rows) = match self {
             StackOp::Push(target) => {
                 let bytes = match target {
-                    PushTarget::Literal(n) => n.to_trimmed_be_bytes().len() as u32,
+                    PushTarget::Literal(n) => u256_to_trimmed_be_bytes(n).len() as u32,
                     PushTarget::Label(_) => BYTES_PER_OFFSET as u32,
                     PushTarget::MacroVar(_) | PushTarget::Constant(_) => {
                         panic!("Target should have been expanded already: {:?}", target)

evm/src/cpu/kernel/utils.rs

@@ -0,0 +1,24 @@
+use ethereum_types::U256;
+use plonky2_util::ceil_div_usize;
+
+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.
+    (0..num_bytes).rev().map(|i| u256.byte(i)).collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn literal_to_be_bytes() {
+        assert_eq!(u256_to_trimmed_be_bytes(&0.into()), vec![0x00]);
+
+        assert_eq!(u256_to_trimmed_be_bytes(&768.into()), vec![0x03, 0x00]);
+
+        assert_eq!(u256_to_trimmed_be_bytes(&0xa1b2.into()), vec![0xa1, 0xb2]);
+
+        assert_eq!(u256_to_trimmed_be_bytes(&0x1b2.into()), vec![0x1, 0xb2]);
+    }
+}