From 78a4b92e8324dafb92936399d3731c35255028dd Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Wed, 14 Sep 2022 11:20:24 -0700
Subject: [PATCH] in %stack, treat identifiers as 1-length blocks

---
 evm/src/cpu/kernel/ast.rs                     |  5 +--
 evm/src/cpu/kernel/parser.rs                  |  4 +-
 .../cpu/kernel/stack/stack_manipulation.rs    | 39 +++++++------------
 evm/src/cpu/kernel/tests/sha2.rs              |  4 +-
 4 files changed, 19 insertions(+), 33 deletions(-)

diff --git a/evm/src/cpu/kernel/ast.rs b/evm/src/cpu/kernel/ast.rs
index 3728aa35..6180b1c8 100644
--- a/evm/src/cpu/kernel/ast.rs
+++ b/evm/src/cpu/kernel/ast.rs
@@ -38,10 +38,7 @@ pub(crate) enum Item {
 
 /// The left hand side of a %stack stack-manipulation macro.
 #[derive(Eq, PartialEq, Clone, Debug)]
-pub(crate) enum StackPlaceholder {
-    Identifier(String),
-    Block(String, usize),
-}
+pub(crate) struct StackPlaceholder(pub String, pub usize);
 
 /// The right hand side of a %stack stack-manipulation macro.
 #[derive(Eq, PartialEq, Clone, Debug)]
diff --git a/evm/src/cpu/kernel/parser.rs b/evm/src/cpu/kernel/parser.rs
index fd762eae..b7a8124b 100644
--- a/evm/src/cpu/kernel/parser.rs
+++ b/evm/src/cpu/kernel/parser.rs
@@ -119,12 +119,12 @@ fn parse_stack_placeholder(target: Pair<Rule>) -> StackPlaceholder {
     assert_eq!(target.as_rule(), Rule::stack_placeholder);
     let inner = target.into_inner().next().unwrap();
     match inner.as_rule() {
-        Rule::identifier => StackPlaceholder::Identifier(inner.as_str().into()),
+        Rule::identifier => StackPlaceholder(inner.as_str().into(), 1),
         Rule::stack_block => {
             let mut block = inner.into_inner();
             let identifier = block.next().unwrap().as_str();
             let length = block.next().unwrap().as_str().parse().unwrap();
-            StackPlaceholder::Block(identifier.to_string(), length)
+            StackPlaceholder(identifier.to_string(), length)
         }
         _ => panic!("Unexpected {:?}", inner.as_rule()),
     }
diff --git a/evm/src/cpu/kernel/stack/stack_manipulation.rs b/evm/src/cpu/kernel/stack/stack_manipulation.rs
index 7264dde9..36e4b83a 100644
--- a/evm/src/cpu/kernel/stack/stack_manipulation.rs
+++ b/evm/src/cpu/kernel/stack/stack_manipulation.rs
@@ -31,20 +31,14 @@ fn expand(names: Vec<StackPlaceholder>, replacements: Vec<StackReplacement>) ->
     let mut src = names
         .iter()
         .cloned()
-        .flat_map(|item| match item {
-            StackPlaceholder::Identifier(name) => {
-                stack_blocks.insert(name.clone(), 1);
-                vec![StackItem::NamedItem(name)]
-            }
-            StackPlaceholder::Block(name, n) => {
-                stack_blocks.insert(name.clone(), n);
-                (0..n)
-                    .map(|i| {
-                        let literal_name = format!("@{}.{}", name, i);
-                        StackItem::NamedItem(literal_name)
-                    })
-                    .collect_vec()
-            }
+        .flat_map(|StackPlaceholder(name, n)| {
+            stack_blocks.insert(name.clone(), n);
+            (0..n)
+                .map(|i| {
+                    let literal_name = format!("@{}.{}", name, i);
+                    StackItem::NamedItem(literal_name)
+                })
+                .collect_vec()
         })
         .collect_vec();
 
@@ -56,17 +50,12 @@ fn expand(names: Vec<StackPlaceholder>, replacements: Vec<StackReplacement>) ->
                 // May be either a named item or a label. Named items have precedence.
                 if stack_blocks.contains_key(&name) {
                     let n = *stack_blocks.get(&name).unwrap();
-                    if n == 1 {
-                        // A name, not an actual block.
-                        vec![StackItem::NamedItem(name)]
-                    } else {
-                        (0..n)
-                            .map(|i| {
-                                let literal_name = format!("@{}.{}", name, i);
-                                StackItem::NamedItem(literal_name)
-                            })
-                            .collect_vec()
-                    }
+                    (0..n)
+                        .map(|i| {
+                            let literal_name = format!("@{}.{}", name, i);
+                            StackItem::NamedItem(literal_name)
+                        })
+                        .collect_vec()
                 } else {
                     vec![StackItem::PushTarget(PushTarget::Label(name))]
                 }
diff --git a/evm/src/cpu/kernel/tests/sha2.rs b/evm/src/cpu/kernel/tests/sha2.rs
index 97edb992..94b3ac61 100644
--- a/evm/src/cpu/kernel/tests/sha2.rs
+++ b/evm/src/cpu/kernel/tests/sha2.rs
@@ -20,7 +20,7 @@ fn test_sha2() -> Result<()> {
 
     let mut hasher = Sha256::new();
     hasher.update(message.clone());
-    let expected = format!("{:02X}", hasher.finalize());
+    let expected = format!("{:X}", hasher.finalize());
 
     let bytes: Vec<U256> = message.iter().map(|&x| U256::from(x as u32)).collect();
 
@@ -32,7 +32,7 @@ fn test_sha2() -> Result<()> {
     let after_sha2 = run(&kernel.code, sha2, initial_stack, &kernel.prover_inputs)?;
     let stack_after_sha2 = after_sha2.stack();
     let result = stack_after_sha2[1];
-    let actual = format!("{:02X}", result);
+    let actual = format!("{:X}", result);
 
     assert_eq!(expected, actual);