11use ethereum_types:: H256 ;
22use ethrex_rlp:: encode:: RLPEncode ;
33
4- use crate :: { Node , NodeHash , NodeRLP , Trie , error:: TrieError } ;
4+ use crate :: { Nibbles , Node , NodeRLP , Trie , error:: TrieError } ;
55use std:: {
66 collections:: BTreeMap ,
77 sync:: { Arc , Mutex } ,
88} ;
99
10+ // Nibbles -> encoded node
11+ pub type NodeMap = Arc < Mutex < BTreeMap < Vec < u8 > , Vec < u8 > > > > ;
12+
1013pub trait TrieDB : Send + Sync {
11- fn get ( & self , key : NodeHash ) -> Result < Option < Vec < u8 > > , TrieError > ;
12- fn put_batch ( & self , key_values : Vec < ( NodeHash , Vec < u8 > ) > ) -> Result < ( ) , TrieError > ;
14+ fn get ( & self , key : Nibbles ) -> Result < Option < Vec < u8 > > , TrieError > ;
15+ fn put_batch ( & self , key_values : Vec < ( Nibbles , Vec < u8 > ) > ) -> Result < ( ) , TrieError > ;
1316 // TODO: replace putbatch with this function.
14- fn put_batch_no_alloc ( & self , key_values : & [ ( NodeHash , Node ) ] ) -> Result < ( ) , TrieError > {
17+ fn put_batch_no_alloc ( & self , key_values : & [ ( Nibbles , Node ) ] ) -> Result < ( ) , TrieError > {
1518 self . put_batch (
1619 key_values
1720 . iter ( )
18- . map ( |node| ( node. 0 , node. 1 . encode_to_vec ( ) ) )
21+ . map ( |node| ( node. 0 . clone ( ) , node. 1 . encode_to_vec ( ) ) )
1922 . collect ( ) ,
2023 )
2124 }
22- fn put ( & self , key : NodeHash , value : Vec < u8 > ) -> Result < ( ) , TrieError > {
25+ fn put ( & self , key : Nibbles , value : Vec < u8 > ) -> Result < ( ) , TrieError > {
2326 self . put_batch ( vec ! [ ( key, value) ] )
2427 }
2528}
2629
2730/// InMemory implementation for the TrieDB trait, with get and put operations.
2831#[ derive( Default ) ]
2932pub struct InMemoryTrieDB {
30- pub inner : Arc < Mutex < BTreeMap < NodeHash , Vec < u8 > > > > ,
33+ inner : NodeMap ,
34+ prefix : Option < Nibbles > ,
3135}
3236
3337impl InMemoryTrieDB {
34- pub const fn new ( map : Arc < Mutex < BTreeMap < NodeHash , Vec < u8 > > > > ) -> Self {
35- Self { inner : map }
38+ pub const fn new ( map : NodeMap ) -> Self {
39+ Self {
40+ inner : map,
41+ prefix : None ,
42+ }
43+ }
44+
45+ pub const fn new_with_prefix ( map : NodeMap , prefix : Nibbles ) -> Self {
46+ Self {
47+ inner : map,
48+ prefix : Some ( prefix) ,
49+ }
3650 }
51+
3752 pub fn new_empty ( ) -> Self {
3853 Self {
3954 inner : Default :: default ( ) ,
55+ prefix : None ,
4056 }
4157 }
4258
@@ -45,33 +61,56 @@ impl InMemoryTrieDB {
4561 state_nodes : & BTreeMap < H256 , NodeRLP > ,
4662 ) -> Result < Self , TrieError > {
4763 let mut embedded_root = Trie :: get_embedded_root ( state_nodes, root_hash) ?;
48- let mut hashed_nodes: Vec < ( NodeHash , Vec < u8 > ) > = vec ! [ ] ;
49- embedded_root. commit ( & mut hashed_nodes) ;
64+ let mut hashed_nodes = vec ! [ ] ;
65+ embedded_root. commit ( Nibbles :: default ( ) , & mut hashed_nodes) ;
5066
51- let hashed_nodes = hashed_nodes. into_iter ( ) . collect ( ) ;
67+ let hashed_nodes = hashed_nodes
68+ . into_iter ( )
69+ . map ( |( k, v) | ( k. into_vec ( ) , v) )
70+ . collect ( ) ;
5271
5372 let in_memory_trie = Arc :: new ( Mutex :: new ( hashed_nodes) ) ;
5473 Ok ( Self :: new ( in_memory_trie) )
5574 }
75+
76+ fn apply_prefix ( & self , path : Nibbles ) -> Nibbles {
77+ match & self . prefix {
78+ Some ( prefix) => prefix. concat ( & path) ,
79+ None => path,
80+ }
81+ }
5682}
5783
5884impl TrieDB for InMemoryTrieDB {
59- fn get ( & self , key : NodeHash ) -> Result < Option < Vec < u8 > > , TrieError > {
85+ fn get ( & self , key : Nibbles ) -> Result < Option < Vec < u8 > > , TrieError > {
6086 Ok ( self
6187 . inner
6288 . lock ( )
6389 . map_err ( |_| TrieError :: LockError ) ?
64- . get ( & key)
90+ . get ( self . apply_prefix ( key) . as_ref ( ) )
6591 . cloned ( ) )
6692 }
6793
68- fn put_batch ( & self , key_values : Vec < ( NodeHash , Vec < u8 > ) > ) -> Result < ( ) , TrieError > {
94+ fn put_batch ( & self , key_values : Vec < ( Nibbles , Vec < u8 > ) > ) -> Result < ( ) , TrieError > {
6995 let mut db = self . inner . lock ( ) . map_err ( |_| TrieError :: LockError ) ?;
7096
7197 for ( key, value) in key_values {
72- db. insert ( key, value) ;
98+ let prefixed_key = self . apply_prefix ( key) ;
99+ db. insert ( prefixed_key. into_vec ( ) , value) ;
73100 }
74101
75102 Ok ( ( ) )
76103 }
77104}
105+
106+ pub fn nibbles_to_fixed_size ( nibbles : Nibbles ) -> [ u8 ; 33 ] {
107+ let node_hash_ref = nibbles. to_bytes ( ) ;
108+ let original_len = node_hash_ref. len ( ) ;
109+
110+ let mut buffer = [ 0u8 ; 33 ] ;
111+
112+ // Encode the node as [node_path..., original_len]
113+ buffer[ 32 ] = nibbles. len ( ) as u8 ;
114+ buffer[ ..original_len] . copy_from_slice ( & node_hash_ref) ;
115+ buffer
116+ }
0 commit comments