From 8e08b218d22af980887c230e6df1e1727365c0eb Mon Sep 17 00:00:00 2001
From: BGluth <gluthb@gmail.com>
Date: Thu, 29 Sep 2022 13:45:46 -0600
Subject: [PATCH] Trie roots now use `H256` instead of `U256`

- `H256` preserves any leading `0`s, which could be critical in some
  situations. Also just a slightly more appropriate type for hashes.
---
 evm/src/generation/mod.rs     | 20 +++++++++++++-------
 evm/src/proof.rs              |  8 ++++----
 evm/src/recursive_verifier.rs | 13 ++++++++-----
 evm/src/util.rs               | 13 ++++++++++++-
 4 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/evm/src/generation/mod.rs b/evm/src/generation/mod.rs
index 511aa009..fb46e0fa 100644
--- a/evm/src/generation/mod.rs
+++ b/evm/src/generation/mod.rs
@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use eth_trie_utils::partial_trie::PartialTrie;
-use ethereum_types::{Address, H256};
+use ethereum_types::{Address, BigEndianHash, H256};
 use plonky2::field::extension::Extendable;
 use plonky2::field::polynomial::PolynomialValues;
 use plonky2::field::types::Field;
@@ -86,14 +86,20 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     };
 
     let trie_roots_before = TrieRoots {
-        state_root: read_metadata(GlobalMetadata::StateTrieRootDigestBefore),
-        transactions_root: read_metadata(GlobalMetadata::TransactionsTrieRootDigestBefore),
-        receipts_root: read_metadata(GlobalMetadata::ReceiptsTrieRootDigestBefore),
+        state_root: H256::from_uint(&read_metadata(GlobalMetadata::StateTrieRootDigestBefore)),
+        transactions_root: H256::from_uint(&read_metadata(
+            GlobalMetadata::TransactionsTrieRootDigestBefore,
+        )),
+        receipts_root: H256::from_uint(&read_metadata(
+            GlobalMetadata::ReceiptsTrieRootDigestBefore,
+        )),
     };
     let trie_roots_after = TrieRoots {
-        state_root: read_metadata(GlobalMetadata::StateTrieRootDigestAfter),
-        transactions_root: read_metadata(GlobalMetadata::TransactionsTrieRootDigestAfter),
-        receipts_root: read_metadata(GlobalMetadata::ReceiptsTrieRootDigestAfter),
+        state_root: H256::from_uint(&read_metadata(GlobalMetadata::StateTrieRootDigestAfter)),
+        transactions_root: H256::from_uint(&read_metadata(
+            GlobalMetadata::TransactionsTrieRootDigestAfter,
+        )),
+        receipts_root: H256::from_uint(&read_metadata(GlobalMetadata::ReceiptsTrieRootDigestAfter)),
     };
 
     let GenerationState {
diff --git a/evm/src/proof.rs b/evm/src/proof.rs
index 81614e67..de00abfc 100644
--- a/evm/src/proof.rs
+++ b/evm/src/proof.rs
@@ -1,4 +1,4 @@
-use ethereum_types::{Address, U256};
+use ethereum_types::{Address, H256, U256};
 use itertools::Itertools;
 use maybe_rayon::*;
 use plonky2::field::extension::{Extendable, FieldExtension};
@@ -54,9 +54,9 @@ pub struct PublicValues {
 
 #[derive(Debug, Clone, Default)]
 pub struct TrieRoots {
-    pub state_root: U256,
-    pub transactions_root: U256,
-    pub receipts_root: U256,
+    pub state_root: H256,
+    pub transactions_root: H256,
+    pub receipts_root: H256,
 }
 
 #[derive(Debug, Clone, Default, Deserialize, Serialize)]
diff --git a/evm/src/recursive_verifier.rs b/evm/src/recursive_verifier.rs
index 000efce9..35041a48 100644
--- a/evm/src/recursive_verifier.rs
+++ b/evm/src/recursive_verifier.rs
@@ -11,7 +11,6 @@ use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2::util::reducing::ReducingFactorTarget;
 use plonky2::with_context;
 
-use crate::all_stark::{AllStark, Table};
 use crate::config::StarkConfig;
 use crate::constraint_consumer::RecursiveConstraintConsumer;
 use crate::cpu::cpu_stark::CpuStark;
@@ -27,9 +26,13 @@ use crate::proof::{
     StarkProofChallengesTarget, StarkProofTarget, TrieRoots, TrieRootsTarget,
 };
 use crate::stark::Stark;
-use crate::util::{h160_limbs, u256_limbs};
+use crate::util::h160_limbs;
 use crate::vanishing_poly::eval_vanishing_poly_circuit;
 use crate::vars::StarkEvaluationTargets;
+use crate::{
+    all_stark::{AllStark, Table},
+    util::h256_limbs,
+};
 
 pub fn verify_proof_circuit<
     F: RichField + Extendable<D>,
@@ -504,15 +507,15 @@ pub fn set_trie_roots_target<F, W, const D: usize>(
 {
     witness.set_target_arr(
         trie_roots_target.state_root,
-        u256_limbs(trie_roots.state_root),
+        h256_limbs(trie_roots.state_root),
     );
     witness.set_target_arr(
         trie_roots_target.transactions_root,
-        u256_limbs(trie_roots.transactions_root),
+        h256_limbs(trie_roots.transactions_root),
     );
     witness.set_target_arr(
         trie_roots_target.receipts_root,
-        u256_limbs(trie_roots.receipts_root),
+        h256_limbs(trie_roots.receipts_root),
     );
 }
 
diff --git a/evm/src/util.rs b/evm/src/util.rs
index 12aead46..7f958fd2 100644
--- a/evm/src/util.rs
+++ b/evm/src/util.rs
@@ -1,6 +1,6 @@
 use std::mem::{size_of, transmute_copy, ManuallyDrop};
 
-use ethereum_types::{H160, U256};
+use ethereum_types::{H160, H256, U256};
 use itertools::Itertools;
 use plonky2::field::extension::Extendable;
 use plonky2::field::packed::PackedField;
@@ -59,6 +59,17 @@ pub(crate) fn u256_limbs<F: Field>(u256: U256) -> [F; 8] {
         .unwrap()
 }
 
+/// Returns the 32-bit little-endian limbs of a `H256`.
+pub(crate) fn h256_limbs<F: Field>(h256: H256) -> [F; 8] {
+    h256.0
+        .chunks(4)
+        .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap()))
+        .map(F::from_canonical_u32)
+        .collect_vec()
+        .try_into()
+        .unwrap()
+}
+
 /// Returns the 32-bit limbs of a `U160`.
 pub(crate) fn h160_limbs<F: Field>(h160: H160) -> [F; 5] {
     h160.0