Merge remote-tracking branch 'smt_trie/feat/type2' into cleanup/smt_t…

…rie_to_develop

smt_trie/Cargo.toml

@@ -0,0 +1,39 @@
+[package]
+name = "smt_trie"
+description = "Types and utility functions for building/working with Polygon Hermez Sparse Merkle Trees."
+version = "0.1.0"
+authors = ["William Borgeaud <[email protected]>"]
+readme = "README.md"
+categories = ["cryptography"]
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+homepage.workspace = true
+keywords.workspace = true
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+bytes = { workspace = true }
+enum-as-inner = { workspace = true }
+ethereum-types = { workspace = true }
+hex = { workspace = true }
+hex-literal = { workspace = true }
+keccak-hash = { workspace = true }
+log = { workspace = true }
+num-traits = { workspace = true }
+parking_lot = { workspace = true, features = ["serde"] }
+plonky2 = { workspace = true }
+rand = { workspace = true }
+rlp = { workspace = true }
+serde = { workspace = true, features = ["derive", "rc"] }
+thiserror = { workspace = true }
+uint = { workspace = true }
+
+
+[dev-dependencies]
+eth_trie = "0.4.0"
+pretty_env_logger = "0.5.0"
+rlp-derive = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }

smt_trie/README.md

@@ -0,0 +1,2 @@
+Types and functions to work with the Hermez/Polygon zkEVM sparse Merkle tree (SMT) format.
+See https://github.com/0xPolygonHermez/zkevm-commonjs for reference implementation.

smt_trie/src/bits.rs

@@ -0,0 +1,103 @@
+use std::ops::Add;
+
+use ethereum_types::{BigEndianHash, H256, U256};
+use serde::{Deserialize, Serialize};
+
+pub type Bit = bool;
+
+#[derive(
+    Copy, Clone, Deserialize, Default, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Debug,
+)]
+pub struct Bits {
+    /// The number of bits in this sequence.
+    pub count: usize,
+    /// A packed encoding of these bits. Only the first (least significant)
+    /// `count` bits are used. The rest are unused and should be zero.
+    pub packed: U256,
+}
+
+impl From<U256> for Bits {
+    fn from(packed: U256) -> Self {
+        Bits { count: 256, packed }
+    }
+}
+
+impl From<H256> for Bits {
+    fn from(packed: H256) -> Self {
+        Bits {
+            count: 256,
+            packed: packed.into_uint(),
+        }
+    }
+}
+
+impl Add for Bits {
+    type Output = Self;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        assert!(self.count + rhs.count <= 256, "Overflow");
+        Self {
+            count: self.count + rhs.count,
+            packed: self.packed * (U256::one() << rhs.count) + rhs.packed,
+        }
+    }
+}
+
+impl Bits {
+    pub fn empty() -> Self {
+        Bits {
+            count: 0,
+            packed: U256::zero(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.count == 0
+    }
+
+    pub fn pop_next_bit(&mut self) -> Bit {
+        assert!(!self.is_empty(), "Cannot pop from empty bits");
+        let b = !(self.packed & U256::one()).is_zero();
+        self.packed >>= 1;
+        self.count -= 1;
+        b
+    }
+
+    pub fn get_bit(&self, i: usize) -> Bit {
+        assert!(i < self.count, "Index out of bounds");
+        !(self.packed & (U256::one() << (self.count - 1 - i))).is_zero()
+    }
+
+    pub fn push_bit(&mut self, bit: Bit) {
+        self.packed = self.packed * 2 + U256::from(bit as u64);
+        self.count += 1;
+    }
+
+    pub fn add_bit(&self, bit: Bit) -> Self {
+        let mut x = *self;
+        x.push_bit(bit);
+        x
+    }
+
+    pub fn common_prefix(&self, k: &Bits) -> (Self, Option<(Bit, Bit)>) {
+        let mut a = *self;
+        let mut b = *k;
+        while a.count > b.count {
+            a.pop_next_bit();
+        }
+        while a.count < b.count {
+            b.pop_next_bit();
+        }
+        if a == b {
+            return (a, None);
+        }
+        let mut a_bit = a.pop_next_bit();
+        let mut b_bit = b.pop_next_bit();
+        while a != b {
+            a_bit = a.pop_next_bit();
+            b_bit = b.pop_next_bit();
+        }
+        assert_ne!(a_bit, b_bit, "Sanity check.");
+        (a, Some((a_bit, b_bit)))
+    }
+}

smt_trie/src/code.rs

@@ -0,0 +1,85 @@
+/// Functions to hash contract bytecode using Poseidon.
+/// See `hashContractBytecode()` in https://github.com/0xPolygonHermez/zkevm-commonjs/blob/main/src/smt-utils.js for reference implementation.
+use ethereum_types::U256;
+use plonky2::field::types::Field;
+use plonky2::hash::poseidon::{self, Poseidon};
+
+use crate::smt::{HashOut, F};
+use crate::utils::hashout2u;
+
+pub fn hash_contract_bytecode(mut code: Vec<u8>) -> HashOut {
+    poseidon_pad_byte_vec(&mut code);
+
+    poseidon_hash_padded_byte_vec(code)
+}
+
+pub fn poseidon_hash_padded_byte_vec(bytes: Vec<u8>) -> HashOut {
+    let mut capacity = [F::ZERO; poseidon::SPONGE_CAPACITY];
+    let mut arr = [F::ZERO; poseidon::SPONGE_WIDTH];
+    for blocks in bytes.chunks_exact(poseidon::SPONGE_RATE * 7) {
+        arr[..poseidon::SPONGE_RATE].copy_from_slice(
+            &blocks
+                .chunks_exact(7)
+                .map(|block| {
+                    let mut bytes = [0u8; poseidon::SPONGE_RATE];
+                    bytes[..7].copy_from_slice(block);
+                    F::from_canonical_u64(u64::from_le_bytes(bytes))
+                })
+                .collect::<Vec<F>>(),
+        );
+        arr[poseidon::SPONGE_RATE..poseidon::SPONGE_WIDTH].copy_from_slice(&capacity);
+        capacity = F::poseidon(arr)[0..poseidon::SPONGE_CAPACITY]
+            .try_into()
+            .unwrap();
+    }
+    HashOut { elements: capacity }
+}
+
+pub fn poseidon_pad_byte_vec(bytes: &mut Vec<u8>) {
+    bytes.push(0x01);
+    while bytes.len() % 56 != 0 {
+        bytes.push(0x00);
+    }
+    *bytes.last_mut().unwrap() |= 0x80;
+}
+
+pub fn hash_bytecode_u256(code: Vec<u8>) -> U256 {
+    hashout2u(hash_contract_bytecode(code))
+}
+
+#[cfg(test)]
+mod tests {
+    use hex_literal::hex;
+
+    use super::*;
+
+    #[test]
+    fn test_empty_code() {
+        assert_eq!(
+            hash_contract_bytecode(vec![]).elements,
+            [
+                10052403398432742521,
+                15195891732843337299,
+                2019258788108304834,
+                4300613462594703212,
+            ]
+            .map(F::from_canonical_u64)
+        );
+    }
+
+    #[test]
+    fn test_some_code() {
+        let code = hex!("60806040526004361061003f5760003560e01c80632b68b9c6146100445780633fa4f2451461005b5780635cfb28e714610086578063718da7ee14610090575b600080fd5b34801561005057600080fd5b506100596100b9565b005b34801561006757600080fd5b506100706100f2565b60405161007d9190610195565b60405180910390f35b61008e6100f8565b005b34801561009c57600080fd5b506100b760048036038101906100b29190610159565b610101565b005b60008054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16ff5b60015481565b34600181905550565b806000806101000a81548173ffffffffffffffffffffffffffffffffffffffff021916908373ffffffffffffffffffffffffffffffffffffffff16021790555050565b600081359050610153816101f1565b92915050565b60006020828403121561016f5761016e6101ec565b5b600061017d84828501610144565b91505092915050565b61018f816101e2565b82525050565b60006020820190506101aa6000830184610186565b92915050565b60006101bb826101c2565b9050919050565b600073ffffffffffffffffffffffffffffffffffffffff82169050919050565b6000819050919050565b600080fd5b6101fa816101b0565b811461020557600080fd5b5056fea26469706673582212207ae6e5d5feddef608b24cca98990c37cf78f8b377163a7c4951a429d90d6120464736f6c63430008070033");
+
+        assert_eq!(
+            hash_contract_bytecode(code.to_vec()).elements,
+            [
+                13311281292453978464,
+                8384462470517067887,
+                14733964407220681187,
+                13541155386998871195
+            ]
+            .map(F::from_canonical_u64)
+        );
+    }
+}

smt_trie/src/db.rs

@@ -0,0 +1,23 @@
+use std::collections::HashMap;
+
+use crate::smt::{Key, Node};
+
+pub trait Db: Default {
+    fn get_node(&self, key: &Key) -> Option<&Node>;
+    fn set_node(&mut self, key: Key, value: Node);
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct MemoryDb {
+    pub db: HashMap<Key, Node>,
+}
+
+impl Db for MemoryDb {
+    fn get_node(&self, key: &Key) -> Option<&Node> {
+        self.db.get(key)
+    }
+
+    fn set_node(&mut self, key: Key, value: Node) {
+        self.db.insert(key, value);
+    }
+}

smt_trie/src/keys.rs

@@ -0,0 +1,99 @@
+#![allow(clippy::needless_range_loop)]
+
+/// This module contains functions to generate keys for the SMT.
+/// See https://github.com/0xPolygonHermez/zkevm-commonjs/blob/main/src/smt-utils.js for reference implementation.
+use ethereum_types::{Address, U256};
+use plonky2::{field::types::Field, hash::poseidon::Poseidon};
+
+use crate::smt::{Key, F};
+
+const HASH_ZEROS: [u64; 4] = [
+    4330397376401421145,
+    14124799381142128323,
+    8742572140681234676,
+    14345658006221440202,
+];
+
+const SMT_KEY_BALANCE: u64 = 0;
+const SMT_KEY_NONCE: u64 = 1;
+const SMT_KEY_CODE: u64 = 2;
+const SMT_KEY_STORAGE: u64 = 3;
+const SMT_KEY_LENGTH: u64 = 4;
+
+pub fn key_balance(addr: Address) -> Key {
+    let mut arr = [F::ZERO; 12];
+    for i in 0..5 {
+        arr[i] = F::from_canonical_u32(u32::from_be_bytes(
+            addr.0[16 - 4 * i..16 - 4 * i + 4].try_into().unwrap(),
+        ));
+    }
+
+    arr[6] = F::from_canonical_u64(SMT_KEY_BALANCE);
+    arr[8..12].copy_from_slice(&HASH_ZEROS.map(F::from_canonical_u64));
+
+    Key(F::poseidon(arr)[0..4].try_into().unwrap())
+}
+
+pub fn key_nonce(addr: Address) -> Key {
+    let mut arr = [F::ZERO; 12];
+    for i in 0..5 {
+        arr[i] = F::from_canonical_u32(u32::from_be_bytes(
+            addr.0[16 - 4 * i..16 - 4 * i + 4].try_into().unwrap(),
+        ));
+    }
+
+    arr[6] = F::from_canonical_u64(SMT_KEY_NONCE);
+    arr[8..12].copy_from_slice(&HASH_ZEROS.map(F::from_canonical_u64));
+
+    Key(F::poseidon(arr)[0..4].try_into().unwrap())
+}
+
+pub fn key_code(addr: Address) -> Key {
+    let mut arr = [F::ZERO; 12];
+    for i in 0..5 {
+        arr[i] = F::from_canonical_u32(u32::from_be_bytes(
+            addr.0[16 - 4 * i..16 - 4 * i + 4].try_into().unwrap(),
+        ));
+    }
+
+    arr[6] = F::from_canonical_u64(SMT_KEY_CODE);
+    arr[8..12].copy_from_slice(&HASH_ZEROS.map(F::from_canonical_u64));
+
+    Key(F::poseidon(arr)[0..4].try_into().unwrap())
+}
+
+pub fn key_storage(addr: Address, slot: U256) -> Key {
+    let mut arr = [F::ZERO; 12];
+    for i in 0..5 {
+        arr[i] = F::from_canonical_u32(u32::from_be_bytes(
+            addr.0[16 - 4 * i..16 - 4 * i + 4].try_into().unwrap(),
+        ));
+    }
+
+    arr[6] = F::from_canonical_u64(SMT_KEY_STORAGE);
+    let capacity: [F; 4] = {
+        let mut arr = [F::ZERO; 12];
+        for i in 0..4 {
+            arr[2 * i] = F::from_canonical_u32(slot.0[i] as u32);
+            arr[2 * i + 1] = F::from_canonical_u32((slot.0[i] >> 32) as u32);
+        }
+        F::poseidon(arr)[0..4].try_into().unwrap()
+    };
+    arr[8..12].copy_from_slice(&capacity);
+
+    Key(F::poseidon(arr)[0..4].try_into().unwrap())
+}
+
+pub fn key_code_length(addr: Address) -> Key {
+    let mut arr = [F::ZERO; 12];
+    for i in 0..5 {
+        arr[i] = F::from_canonical_u32(u32::from_be_bytes(
+            addr.0[16 - 4 * i..16 - 4 * i + 4].try_into().unwrap(),
+        ));
+    }
+
+    arr[6] = F::from_canonical_u64(SMT_KEY_LENGTH);
+    arr[8..12].copy_from_slice(&HASH_ZEROS.map(F::from_canonical_u64));
+
+    Key(F::poseidon(arr)[0..4].try_into().unwrap())
+}

smt_trie/src/lib.rs

@@ -0,0 +1,8 @@
+pub mod bits;
+pub mod code;
+pub mod db;
+pub mod keys;
+pub mod smt;
+#[cfg(test)]
+mod smt_test;
+pub mod utils;