Skip to content

Commit

Permalink
dump
Browse files Browse the repository at this point in the history
  • Loading branch information
@clabby
clabby committed Apr 22, 2024
1 parent a69c439 commit 34a999a
Show file tree
Hide file tree
Showing 3 changed files with 245 additions and 24 deletions.
72 changes: 67 additions & 5 deletions crates/mpt/src/db/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use core::marker::PhantomData;
use revm::{db::DbAccount, Database, DatabaseCommit, DatabaseRef, InMemoryDB};
use revm_primitives::{hash_map::Entry, Account, AccountInfo, Bytecode, HashMap};

use crate::TrieNode;
use crate::{NodeElement, TrieNode};

/// A Trie DB that caches account state in-memory. When accounts that don't already exist within the cache are queried,
/// the database fetches the preimages of the trie nodes on the path to the account using the `PreimageFetcher`
Expand Down Expand Up @@ -89,14 +89,76 @@ where
todo!()
}

/// Walks down the trie to a leaf value with the given key, if it exists. Preimages for blinded nodes along the
/// path are fetched using the `fetcher` function.
///
/// TODO: Fix nibble relations
fn get_trie(
&self,
key: Nibbles,
item_key: &Bytes,
trie_node: TrieNode,
pos: usize,
mut pos: usize,
fetcher: PF,
) -> Result<TrieNode> {
todo!()
) -> Result<(Bytes, Bytes)> {
match trie_node {
TrieNode::Branch { stack } => {
let next = item_key[pos];
extern crate std;
std::dbg!(next);

// for (i, node) in stack.into_iter().enumerate() {
// match node {
// NodeElement::String(s) => {
// // If the string is a hash, we need to grab the preimage for it and
// // continue recursing.
// let hash: B256 = s.as_ref().try_into().map_err(|e| anyhow!("Conversion error: {e}"))?;
// let trie_node = TrieNode::decode(&mut fetcher(hash)?.as_ref()).map_err(|e| anyhow!(e))?;
//
// // If the value was found in the blinded node, return it.
// if let Ok((key, value)) = self.get_trie(item_key, trie_node, pos, fetcher) {
// return Ok((key, value));
// }
// }
// list @ NodeElement::List(_) => {
// let trie_node = list.try_list_into_node()?;
//
// // If the value was found in the blinded node, return it.
// if let Ok((key, value)) = self.get_trie(item_key, trie_node, pos, fetcher) {
// return Ok((key, value));
// }
// }
// _ => { /* Skip over empty lists and strings; We're looking for leaves */ }
// };
// }

anyhow::bail!("Key does not exist in trie");
}
TrieNode::Leaf { key, value } => {
let shared_nibbles = key[1..].as_ref();
let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref();
if item_key_nibbles == shared_nibbles {
Ok((key, value))
} else {
anyhow::bail!("Key does not exist in trie");
}
}
TrieNode::Extension { prefix, node } => {
let shared_nibbles = prefix[1..].as_ref();
let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref();
if item_key_nibbles == shared_nibbles {
// Increase the offset within the key by the length of the shared nibbles
pos += shared_nibbles.len();

// Follow extension branch
let hash = B256::from_slice(node.as_ref());
let extension_link =
TrieNode::decode(&mut fetcher(hash)?.as_ref()).map_err(|e| anyhow!(e))?;
self.get_trie(item_key, extension_link, pos, fetcher)
} else {
anyhow::bail!("Key does not exist in trie");
}
}
}
}
}

Expand Down
4 changes: 4 additions & 0 deletions crates/mpt/src/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,10 @@ impl TrieNode {
match path[0] >> 4 {
PREFIX_EXTENSION_EVEN | PREFIX_EXTENSION_ODD => {
// extension node
{
extern crate std;
std::dbg!("FOUND EXTENSION NODE", &path, &value);
}
Ok(TrieNode::Extension { prefix: path, node: value })
}
PREFIX_LEAF_EVEN | PREFIX_LEAF_ODD => {
Expand Down
193 changes: 174 additions & 19 deletions crates/mpt/src/test_util.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,23 @@
extern crate std;

use std::dbg;

use alloc::{
collections::BTreeMap,
vec::{self, Vec},
};
use alloy_consensus::{Receipt, ReceiptEnvelope, ReceiptWithBloom, TxEnvelope, TxType};
use alloy_primitives::{b256, keccak256, Bytes, Log, B256};
use alloy_provider::{network::eip2718::Encodable2718, Provider, ProviderBuilder};
use alloy_rlp::{BufMut, Encodable};
use alloy_rlp::{BufMut, Decodable, Encodable};
use alloy_rpc_types::BlockTransactions;
use alloy_trie::{HashBuilder, Nibbles};
use anyhow::{anyhow, Result};
use reqwest::Url;

use crate::{NodeElement, TrieNode};

const RPC_URL: &str = "https://docs-demo.quiknode.pro/";

/// Grabs a live merkleized receipts list within a block header.
Expand Down Expand Up @@ -173,24 +177,175 @@ pub(crate) const fn adjust_index_for_rlp(i: usize, len: usize) -> usize {
fn test_trie() {
use alloc::vec;

let mut hb = HashBuilder::default()
.with_proof_retainer(vec![Nibbles::unpack(&[0x00]), Nibbles::unpack(&[0x01])]);

// hb.add_leaf(Nibbles::unpack(&[0x00, 0xEE]), b"test one");
// hb.add_leaf(Nibbles::unpack(&[0x00, 0xFF]), b"test one");
// hb.add_leaf(Nibbles::unpack(&[0x01, 0xDD]), b"test one");
// hb.add_leaf(Nibbles::unpack(&[0x01, 0xEE]), b"test one");
// hb.add_leaf(Nibbles::unpack(&[0x01, 0xFF]), b"test one");
hb.add_branch(
Nibbles::unpack(&[0x00]),
b256!("f4ae7801fd7296c9cb9f2387149e93079bd7c74158fea76d978947fddbead8b7"),
true,
);
hb.add_branch(
let mut hb = HashBuilder::default().with_proof_retainer(vec![
Nibbles::unpack(&[0x80]),
Nibbles::unpack(&[0x01]),
b256!("91c0bc2b7771df00372f3b3ec799e2586115046fabc2b406c94b4d793ff1669c"),
true,
);
Nibbles::unpack(&[0xFF]),
]);

hb.add_leaf(Nibbles::unpack(&[0x01]), b"test two");
hb.add_leaf(Nibbles::unpack(&[0x88]), b"test one");
hb.add_leaf(Nibbles::unpack(&[0xFF]), b"test three");
// hb.add_branch(
// Nibbles::unpack(&[0x00]),
// b256!("f4ae7801fd7296c9cb9f2387149e93079bd7c74158fea76d978947fddbead8b7"),
// true,
// );
// hb.add_branch(
// Nibbles::unpack(&[0x01]),
// b256!("91c0bc2b7771df00372f3b3ec799e2586115046fabc2b406c94b4d793ff1669c"),
// true,
// );
std::dbg!(hb.root());
std::dbg!(hb.take_proofs());

let proofs = hb.take_proofs();
let preimages = proofs.into_iter().fold(BTreeMap::default(), |mut acc, (_, v)| {
acc.insert(keccak256(v.as_ref()), v);
acc
});
let fetcher = |hash: B256| -> Result<Bytes> { Ok(preimages.get(&hash).cloned().unwrap()) };

let root = TrieNode::decode(&mut fetcher(hb.root()).unwrap().as_ref()).unwrap();
std::dbg!(get_trie(&Nibbles::unpack(&[0x01]), root, 0, fetcher).unwrap());
}

fn adjust_index_for_read(index: usize) -> usize {
match index.cmp(&0x80) {
core::cmp::Ordering::Less => index,
core::cmp::Ordering::Equal => 0x00,
core::cmp::Ordering::Greater => index - 1,
}
}

#[tokio::test]
async fn test_trie_get() {
// Initialize the provider.
let provider =
ProviderBuilder::new().on_http(Url::parse(RPC_URL).expect("invalid rpc url")).unwrap();

let block_number = 19005266;
let block = provider.get_block(block_number.into(), true).await.unwrap().unwrap();

let BlockTransactions::Full(txs) = block.transactions else {
panic!("Did not fetch full block");
};
let consensus_txs = txs
.into_iter()
.map(|tx| TxEnvelope::try_from(tx).map_err(|e| anyhow!(e)))
.collect::<Result<Vec<_>>>()
.unwrap();

// Compute the derivable list
let mut list =
ordered_trie_with_encoder(consensus_txs.as_ref(), |rlp, buf| rlp.encode_2718(buf));
let root = list.root();

// Sanity check transaction root is correct
assert_eq!(block.header.transactions_root, root);

// Construct the mapping of hashed intermediates -> raw intermediates
let proofs = list.take_proofs();
let preimages = proofs.into_iter().fold(BTreeMap::default(), |mut acc, (_, value)| {
acc.insert(keccak256(value.as_ref()), value);
acc
});

let fetcher = |hash: B256| -> Result<Bytes> { Ok(preimages.get(&hash).cloned().unwrap()) };

let root = TrieNode::decode(&mut fetcher(root).unwrap().as_ref()).unwrap();
for i in 1..135 {
let (_, v) = get_trie(
&Nibbles::unpack(alloc::vec![i as u8]),
root.clone(),
0,
fetcher,
)
.unwrap();
let mut rlp_buf = Vec::new();
consensus_txs[adjust_index_for_read(i)].encode_2718(&mut rlp_buf);
assert_eq!(v.as_ref(), rlp_buf.as_slice(), "Failed at index: {}", i);
}
dbg!(block.header.number);
dbg!(consensus_txs[0x0].tx_hash());
}

/// Walks down the trie to a leaf value with the given key, if it exists. Preimages for blinded nodes along the
/// path are fetched using the `fetcher` function.
fn get_trie(
item_key: &Nibbles,
trie_node: TrieNode,
mut pos: usize,
fetcher: impl Fn(B256) -> Result<Bytes> + Copy,
) -> Result<(Bytes, Bytes)> {
match trie_node {
TrieNode::Branch { mut stack } => {
let branch_nibble = item_key[pos];
pos += 1;

match stack
.remove(branch_nibble as usize)
.ok_or(anyhow!("Key does not exist in trie"))?
{
NodeElement::String(s) => {
// If the string is a hash, we need to grab the preimage for it and
// continue recursing.
let hash: B256 =
s.as_ref().try_into().map_err(|e| anyhow!("Conversion error: {e}"))?;
let trie_node =
TrieNode::decode(&mut fetcher(hash)?.as_ref()).map_err(|e| anyhow!(e))?;

// If the value was found in the blinded node, return it.
if let Ok((key, value)) = get_trie(item_key, trie_node, pos, fetcher) {
return Ok((key, value));
}
}
list @ NodeElement::List(_) => {
let trie_node = list.try_list_into_node()?;

// If the value was found in the blinded node, return it.
if let Ok((key, value)) = get_trie(item_key, trie_node, pos, fetcher) {
return Ok((key, value));
}
}
_ => { /* Skip over empty lists and strings; We're looking for leaves */ }
};

anyhow::bail!("Key does not exist in trie");
}
TrieNode::Leaf { key, value } => {
// If the key length is one, it only contains the prefix and no shared nibbles. Return the
// key and value.
if key.len() == 1 {
return Ok((key, value));
}

let key_nibbles = Nibbles::unpack(key.clone());
let shared_nibbles = key_nibbles[1..].as_ref();
let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref();

if item_key_nibbles == shared_nibbles {
Ok((key, value))
} else {
anyhow::bail!("Key does not exist in trie");
}
}
TrieNode::Extension { prefix, node } => {
std::dbg!(&prefix);
let prefix_nibbles = Nibbles::unpack(prefix);
let shared_nibbles = prefix_nibbles[1..].as_ref();
let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref();
if item_key_nibbles == shared_nibbles {
// Increase the offset within the key by the length of the shared nibbles
pos += shared_nibbles.len();

// Follow extension branch
let hash = B256::from_slice(node.as_ref());
let extension_link =
TrieNode::decode(&mut fetcher(hash)?.as_ref()).map_err(|e| anyhow!(e))?;
get_trie(item_key, extension_link, pos, fetcher)
} else {
anyhow::bail!("Key does not exist in trie");
}
}
}
}

0 comments on commit 34a999a

Please sign in to comment.