feat: imple verkle trie

verkle/Cargo.toml

@@ -6,14 +6,20 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+alloy-primitives = { version = "0.7.0", features = ["serde", "ssz"] }
+anyhow = "1.0.82"
 ark-ec = "0.4.2"
 ark-ed-on-bls12-381-bandersnatch = "0.4.0"
 ark-ff = "0.4.2"
 ark-serialize = "0.4.2"
 banderwagon = { git = "https://github.com/crate-crypto/rust-verkle.git", rev = "7688f0aedfb147d3d391abfe8495e46c46d72ce0" }
+db = { path = "../db" }
 derive_more = "0.99.17"
+ethereum_ssz = "0.5.3"
+ethereum_ssz_derive = "0.5.3"
 once_cell = "1.19.0"
 sha2 = "0.10.8"
+ssz_types = "0.6.0"
 
 [dev-dependencies]
 hex = "0.4.3"

verkle/src/committer.rs

@@ -0,0 +1,50 @@
+use banderwagon::{msm::MSMPrecompWnaf, Element, Fr, Zero};
+use once_cell::sync::Lazy;
+
+use crate::{constants::VERKLE_NODE_WIDTH, crs::CRS};
+
+pub static DEFAULT_COMMITER: Lazy<Committer> = Lazy::new(Committer::new);
+
+pub struct Committer {
+    precomp: MSMPrecompWnaf,
+}
+
+impl Committer {
+    fn new() -> Self {
+        Self {
+            precomp: MSMPrecompWnaf::new(CRS.as_slice(), 12),
+        }
+    }
+
+    // Commit to a lagrange polynomial, evaluations.len() must equal the size of the CRS at the moment
+    pub fn commit_lagrange(&self, evaluations: &[Fr]) -> Element {
+        // Preliminary benchmarks indicate that the parallel version is faster
+        // for vectors of length 64 or more
+        if evaluations.len() >= 64 {
+            self.precomp.mul_par(evaluations)
+        } else {
+            self.precomp.mul(evaluations)
+        }
+    }
+
+    pub fn scalar_mul(&self, index: usize, value: Fr) -> Element {
+        self.precomp.mul_index(value, index)
+    }
+
+    pub fn commit_sparse(&self, evaluations: Vec<(usize, Fr)>) -> Element {
+        // TODO consider if 64 is good value
+        if evaluations.len() >= 64 {
+            let mut dense = [Fr::zero(); VERKLE_NODE_WIDTH];
+            for (index, value) in evaluations {
+                dense[index] = value;
+            }
+            self.commit_lagrange(&dense)
+        } else {
+            let mut result = Element::zero();
+            for (index, value) in evaluations {
+                result += self.scalar_mul(index, value)
+            }
+            result
+        }
+    }
+}

verkle/src/constants.rs

@@ -0,0 +1 @@
+pub const VERKLE_NODE_WIDTH: usize = 256;

verkle/src/crs.rs

@@ -3,19 +3,19 @@ use derive_more::{AsRef, Deref, Index};
 use once_cell::sync::Lazy;
 use sha2::{Digest, Sha256};
 
-const BASIS_LENGTH: usize = 256;
+use crate::constants::VERKLE_NODE_WIDTH;
 
 const PEDERSEN_SEED: &[u8] = b"eth_verkle_oct_2021";
 
-pub static CRS: Lazy<Basis> = Lazy::new(Basis::new);
+pub static CRS: Lazy<Bases> = Lazy::new(Bases::new);
 
 #[derive(AsRef, Deref, Index)]
-pub struct Basis([Element; BASIS_LENGTH]);
+pub struct Bases([Element; VERKLE_NODE_WIDTH]);
 
-impl Basis {
+impl Bases {
     fn new() -> Self {
         let mut generated_elements = 0;
-        let mut elements = [Element::zero(); BASIS_LENGTH];
+        let mut elements = [Element::zero(); VERKLE_NODE_WIDTH];
 
         for i in 0u64.. {
             if generated_elements == elements.len() {

verkle/src/lib.rs

@@ -1 +1,34 @@
+use alloy_primitives::{B256, U256};
+use derive_more::{Deref, Index};
+
+use banderwagon::Fr;
+pub use stem::Stem as TrieStem;
+
+mod committer;
+mod constants;
 pub mod crs;
+pub mod nodes;
+mod stem;
+pub mod trie;
+mod utils;
+
+pub type TrieValue = U256;
+
+type Db = dyn db::Db<Fr, Vec<u8>>;
+
+#[derive(PartialEq, Eq, Clone, Index, Deref)]
+pub struct TrieKey(B256);
+
+impl TrieKey {
+    pub fn length() -> usize {
+        B256::len_bytes()
+    }
+
+    pub fn stem(&self) -> stem::Stem {
+        self.into()
+    }
+
+    pub fn last(&self) -> u8 {
+        self[self.len() - 1]
+    }
+}

verkle/src/nodes/branch.rs

@@ -0,0 +1,127 @@
+use std::{array, collections::BTreeMap, ops::DerefMut};
+
+use alloy_primitives::B256;
+use anyhow::Result;
+use banderwagon::{Element, Fr};
+use ssz::{Decode, Encode};
+
+use crate::{
+    committer::DEFAULT_COMMITER,
+    constants::VERKLE_NODE_WIDTH,
+    utils::{b256_to_fr, fr_to_b256},
+    Db, TrieKey, TrieValue,
+};
+
+use super::{node::NodeTrait, Node};
+
+pub struct BranchNode {
+    values: Box<[Node; VERKLE_NODE_WIDTH]>,
+    cp: Element,
+}
+
+impl BranchNode {
+    pub fn new() -> Self {
+        Self {
+            values: array::from_fn(|_| Node::Empty).into(),
+            cp: Element::zero(),
+        }
+    }
+
+    pub fn set(&mut self, index: usize, node: Node) {
+        let node_at_index = &mut self.values[index];
+        let pre_commitment = node_at_index.commit();
+        *node_at_index = node;
+        let post_commitment = node_at_index.commit();
+        self.cp += DEFAULT_COMMITER.scalar_mul(index, post_commitment - pre_commitment);
+    }
+
+    pub(super) fn get_mut(&mut self, index: usize) -> &mut Node {
+        &mut self.values[index]
+    }
+
+    pub fn insert(&mut self, depth: usize, key: TrieKey, value: TrieValue, db: &Db) -> Result<()> {
+        let index = key[depth] as usize;
+        let node = &mut self.values[index];
+        let pre_commitment = node.commit();
+        node.insert(depth + 1, key, value, db)?;
+        let post_commitment = node.commit();
+        self.cp += DEFAULT_COMMITER.scalar_mul(index, post_commitment - pre_commitment);
+        Ok(())
+    }
+
+    pub fn write_and_commit(&mut self, db: &mut Db) -> Result<Fr> {
+        for node in self.values.deref_mut() {
+            node.write_and_commit(db)?;
+        }
+        Ok(self.commit())
+    }
+}
+
+impl Default for BranchNode {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl NodeTrait for BranchNode {
+    fn commit(&self) -> Fr {
+        self.cp.map_to_scalar_field()
+    }
+}
+
+impl Encode for BranchNode {
+    fn is_ssz_fixed_len() -> bool {
+        false
+    }
+
+    fn ssz_append(&self, buf: &mut Vec<u8>) {
+        let commitments: BTreeMap<u8, B256> = self
+            .values
+            .iter()
+            .enumerate()
+            .filter_map(|(index, node)| {
+                if node.is_empty() {
+                    None
+                } else {
+                    Some((index as u8, fr_to_b256(&node.commit())))
+                }
+            })
+            .collect();
+        commitments.ssz_append(buf);
+    }
+
+    fn ssz_bytes_len(&self) -> usize {
+        // TODO: optimize this
+        // let number_of_non_empty_items = self.values.iter().filter(|node| !node.is_empty()).count();
+        // let size_per_item = <(u8, B256) as Decode>::ssz_fixed_len() ;
+        // number_of_non_empty_items * size_per_item
+        self.as_ssz_bytes().len()
+    }
+}
+
+impl Decode for BranchNode {
+    fn is_ssz_fixed_len() -> bool {
+        false
+    }
+
+    fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, ssz::DecodeError> {
+        let commitments = BTreeMap::<u8, B256>::from_ssz_bytes(bytes)?;
+        let commitments: BTreeMap<usize, Fr> = commitments
+            .iter()
+            .map(|(index, commitment)| (*index as usize, b256_to_fr(commitment)))
+            .collect();
+
+        let values = array::from_fn(|i| {
+            commitments
+                .get(&i)
+                .map(|c| Node::Commitment(*c))
+                .unwrap_or_default()
+        });
+        let cp = DEFAULT_COMMITER.commit_sparse(commitments.into_iter().collect());
+
+        Ok(Self {
+            values: values.into(),
+            cp,
+        })
+    }
+}