Skip to content

Commit

Permalink
Merge pull request #17 from HerodotusDev/optimization/lazy-rlp
Browse files Browse the repository at this point in the history 
Lazy RLP
  • Loading branch information
@tiagofneto
tiagofneto authored Oct 17, 2023
2 parents aed6b51 + 1c2a5d3 commit 9ca7589
Show file tree
Hide file tree
Showing 4 changed files with 476 additions and 75 deletions.
109 changes: 36 additions & 73 deletions src/data_structures/eth_mpt.cairo
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use cairo_lib::hashing::keccak::keccak_cairo_words64;
use cairo_lib::encoding::rlp::{RLPItem, rlp_decode};
use cairo_lib::encoding::rlp::{RLPItem, rlp_decode, rlp_decode_list_lazy};
use cairo_lib::utils::types::byte::{Byte, ByteTrait};
use cairo_lib::utils::bitwise::{right_shift, left_shift};
use cairo_lib::utils::types::words64::{Words64, Words64Trait, Words64TryIntoU256LE};
Expand All @@ -19,11 +19,13 @@ impl MPTDefault of Default<MPT> {
}

// @notice Represents a node in the MPT
#[derive(Drop)]
#[derive(Drop, PartialEq)]
enum MPTNode {
// @param 16 hashes of children
// @param Value of the node
Branch: (Span<Words64>, Words64),
// @param hash of the next node
LazyBranch: u256,
// @param shared_nibbles
// @param next_node
// @param nibbles_skip Number of nibbles to skip in shared nibbles
Expand Down Expand Up @@ -55,8 +57,10 @@ impl MPTImpl of MPTTrait {
let mut proof_index: usize = 0;
let mut key_pow2: u256 = pow(2, (key_len.into() - 1) * 4);

let proof_len = proof.len();

loop {
if proof_index == proof.len() {
if proof_index == proof_len {
break Result::Err('Proof reached end');
}

Expand All @@ -65,10 +69,22 @@ impl MPTImpl of MPTTrait {
let hash = MPTTrait::hash_rlp_node(node);
assert(hash == current_hash, 'Element not matching');

let decoded = match MPTTrait::decode_rlp_node(node) {
Result::Ok(decoded) => decoded,
Result::Err(e) => {
break Result::Err(e);
// If it's not the last node and more than 9 words, it must be a branch node
let decoded = if proof_index != proof_len - 1 && node.len() > 9 {
let current_nibble = (key / key_pow2) & 0xf;
// Unwrap impossible to fail, as we are masking with 0xf, meaning the result is always a nibble
match MPTTrait::lazy_rlp_decode_branch_node(node, current_nibble.try_into().unwrap()) {
Result::Ok(decoded) => decoded,
Result::Err(e) => {
break Result::Err(e);
}
}
} else {
match MPTTrait::decode_rlp_node(node) {
Result::Ok(decoded) => decoded,
Result::Err(e) => {
break Result::Err(e);
}
}
};
match decoded {
Expand Down Expand Up @@ -96,6 +112,10 @@ impl MPTImpl of MPTTrait {
};
key_pow2 = key_pow2 / 16;
},
MPTNode::LazyBranch(next_node) => {
current_hash = next_node;
key_pow2 = key_pow2 / 16;
},
MPTNode::Extension((
shared_nibbles, next_node, nibbles_skip
)) => {
Expand Down Expand Up @@ -252,76 +272,19 @@ impl MPTImpl of MPTTrait {
}
}


fn lazy_rlp_decode_branch_node(rlp: Words64, current_nibble: u8) -> Result<MPTNode, felt252> {
let hash_words = rlp_decode_list_lazy(rlp, array![current_nibble.into()].span())?;
match (*hash_words.at(0)).try_into() {
Option::Some(h) => Result::Ok(MPTNode::LazyBranch(h)),
Option::None(_) => Result::Err('Invalid hash')
}
}

// @notice keccak256 hashes an RLP encoded node
// @param rlp RLP encoded node
// @return keccak256 hash of the node
fn hash_rlp_node(rlp: Words64) -> u256 {
keccak_cairo_words64(rlp)
}
}

impl MPTNodePartialEq of PartialEq<MPTNode> {
fn eq(lhs: @MPTNode, rhs: @MPTNode) -> bool {
match lhs {
MPTNode::Branch((
lhs_nibbles, lhs_value
)) => {
match rhs {
MPTNode::Branch((
rhs_nibbles, rhs_value
)) => {
if (*lhs_nibbles).len() != (*rhs_nibbles).len() {
return false;
}
let mut i: usize = 0;
loop {
if i >= (*lhs_nibbles).len() {
break lhs_value == rhs_value;
}
if (*lhs_nibbles).at(i) != (*rhs_nibbles).at(i) {
break false;
}
i += 1;
}
},
MPTNode::Extension(_) => false,
MPTNode::Leaf(_) => false
}
},
MPTNode::Extension((
lhs_shared_nibbles, lhs_next_node, lhs_nibbles_skip
)) => {
match rhs {
MPTNode::Branch(_) => false,
MPTNode::Extension((
rhs_shared_nibbles, rhs_next_node, rhs_nibbles_skip
)) => {
lhs_shared_nibbles == rhs_shared_nibbles
&& lhs_next_node == rhs_next_node
&& lhs_nibbles_skip == rhs_nibbles_skip
},
MPTNode::Leaf(_) => false
}
},
MPTNode::Leaf((
lhs_key_end, lhs_value, lhs_nibbles_skip
)) => {
match rhs {
MPTNode::Branch(_) => false,
MPTNode::Extension(_) => false,
MPTNode::Leaf((
rhs_key_end, rhs_value, rhs_nibbles_skip
)) => {
lhs_key_end == rhs_key_end
&& lhs_value == rhs_value
&& lhs_nibbles_skip == rhs_nibbles_skip
}
}
}
}
}

fn ne(lhs: @MPTNode, rhs: @MPTNode) -> bool {
!(lhs == rhs)
}
}
83 changes: 83 additions & 0 deletions src/data_structures/tests/test_eth_mpt.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,9 @@ fn test_decode_rlp_node_branch() {
i += 1;
};
},
MPTNode::LazyBranch(_) => {
panic_with_felt252('Branch node differs');
},
MPTNode::Extension(_) => {
panic_with_felt252('Branch node differs');
},
Expand All @@ -119,6 +122,86 @@ fn test_decode_rlp_node_branch() {
}
}

#[test]
#[available_gas(9999999999)]
fn test_lazy_rlp_node_branch() {
let rlp_node = array![
0x09cf7077a01102f9,
0xa962df351b7a06b5,
0xadecaece75818924,
0x0c4044a8b4cd681f,
0x85a31ea0f44ac173,
0x045c6d4661b25ad0,
0x1a9fc1344568fe87,
0x35361adc184b5c4b,
0x4c2ca0b471500260,
0x1846d1d34035ce04,
0x8366e5a5533c3072,
0x0c80a8368d4f30c1,
0xa9a0eecd3ffaf56a,
0xc4d37d4bc58d77dc,
0xb0fe61d139e72282,
0x3717d5dcb2ceeec0,
0xa05138a6378e5bf0,
0xdd62df56554d5fa9,
0x9b56ae97049962c2,
0x9307207bdafd8ecd,
0xd71897db4cded3f8,
0x2238146d06d439a0,
0x74a843e9c94aaf6e,
0xb91dd8b05fc2a9a9,
0x03e2b336138c1d86,
0x6ab4637ccc7aa04c,
0x25a141a0c9b318a4,
0x7a396b316173cb6b,
0x13bb1b4967885ada,
0x25818a3515a03001,
0xc736fe137193c42e,
0x3497a1fb11b74680,
0x5f78007a1829bb91,
0xd3429168a0ae52f8,
0xdfce8b1ca7faab16,
0x254e10b2db1d2049,
0x1f2256e8c490dc0a,
0x5036dca058964a53,
0xa714a3a8fd342599,
0xb59dc7a83baeb0db,
0xc060242ace690c55,
0xb020a0a3c1c4ad07,
0xe19e05b055663b68,
0xc1cb6b504b4ed003,
0x11b1dab792630039,
0x8ea0e7420366c278,
0xd91c0f63fb45ebed,
0xcb17225718eb3697,
0x03e21bb715f3d5c6,
0xa014269bd9e83cb0,
0x6f985af63da32379,
0x69b9c2e4e6f9e7d5,
0x3999be4e94086b73,
0xf309e62f6114864a,
0x71201ad0d73465a0,
0xce46b9552afba44a,
0xa22aadff2d22c364,
0xb12ac97334928ad1,
0xb8fe8bc2f9bfa0fd,
0xb0c3c818b6a92dbf,
0x4714bdc0b10ce86f,
0xe229ff6121c4f738,
0x3c6961147fa02f50,
0x5ea3bb1b02a54e70,
0x8f459e43f602c572,
0x8fea4837d02e2498,
0x805fb3e2
];

let expected = 0xE7420366C27811B1DAB792630039C1CB6B504B4ED003E19E05B055663B68B020;

let decoded = MPTTrait::lazy_rlp_decode_branch_node(rlp_node.span(), 0xa).unwrap();
let expected_node = MPTNode::LazyBranch(expected);
assert(decoded == expected_node, 'Lazy branch node diffes');
}

#[test]
#[available_gas(9999999999)]
fn test_decode_rlp_node_leaf_odd() {
Expand Down
92 changes: 91 additions & 1 deletion src/encoding/rlp.cairo
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use cairo_lib::utils::types::words64::{Words64, Words64Trait, reverse_endianness_u64};
use cairo_lib::utils::types::words64::{Words64, Words64Trait, reverse_endianness_u64, pow2};
use cairo_lib::utils::types::byte::Byte;
use cairo_lib::utils::array::span_contains;

// @notice Enum with all possible RLP types
// For more info: https://ethereum.org/en/developers/docs/data-structures-and-encoding/rlp/
Expand Down Expand Up @@ -141,6 +142,95 @@ fn rlp_decode_list(ref input: Words64, len: usize) -> Result<Span<Words64>, felt
}
}

fn rlp_decode_list_lazy(input: Words64, lazy: Span<usize>) -> Result<Span<Words64>, felt252> {
let mut output = ArrayTrait::new();
let mut lazy_index = 0;

let list_prefix: u32 = (*input.at(0) & 0xff).try_into().unwrap();
let list_type = RLPTypeTrait::from_byte(list_prefix.try_into().unwrap()).unwrap();
let (mut current_input_index, len) = match list_type {
RLPType::String(()) => { return Result::Err('Not a list'); },
RLPType::StringShort(()) => { return Result::Err('Not a list'); },
RLPType::StringLong(()) => { return Result::Err('Not a list'); },
RLPType::ListShort(()) => (1, list_prefix - 0xc0),
RLPType::ListLong(()) => {
let len_len = list_prefix - 0xf7;
let len_span = input.slice_le(6, len_len);
// Enough to store 4.29 GB (fits in u32)
assert(len_span.len() == 1 && *len_span.at(0) <= 0xffffffff, 'Len of len too big');

// len fits in 32 bits, confirmed by previous assertion
let len = reverse_endianness_u64(*len_span.at(0), Option::Some(len_len.into()))
.try_into()
.unwrap();
(1 + len_len, len)
}
};

loop {
if output.len() == lazy.len() {
break Result::Ok(output.span());
}

if current_input_index >= len {
break Result::Err('Too many items to decode');
}

let current_word = current_input_index / 8;
let current_word_offset = 7 - (current_input_index % 8);

let pow2_shift = pow2((7 - current_word_offset) * 8);
let prefix = (*input.at(current_word) / pow2_shift) & 0xff;

let rlp_type = RLPTypeTrait::from_byte(prefix.try_into().unwrap()).unwrap();
let (item_start_skip, item_len) = match rlp_type {
RLPType::String(()) => {
(0, 1)
},
RLPType::StringShort(()) => {
let len = prefix - 0x80;
(1, len)
},
RLPType::StringLong(()) => {
let len_len = prefix - 0xb7;

let current_word = (current_input_index + 1) / 8;
let current_word_offset = 7 - ((current_input_index + 1) % 8);

let len_span = input.slice_le(current_word * 8 + current_word_offset, len_len.try_into().unwrap());
// Enough to store 4.29 GB (fits in u32)
assert(len_span.len() == 1 && *len_span.at(0) <= 0xffffffff, 'Len of len too big');

// len fits in 32 bits, confirmed by previous assertion
let len: u32 = reverse_endianness_u64(*len_span.at(0), Option::Some(len_len.try_into().unwrap()))
.try_into()
.unwrap();

(1 + len_len, len.into())
},
RLPType::ListShort(()) => {
panic_with_felt252('Recursive list not supported')
},
RLPType::ListLong(()) => {
panic_with_felt252('Recursive list not supported')
}
};

current_input_index += item_start_skip.try_into().unwrap();
if span_contains(lazy, lazy_index) {
let current_word = current_input_index / 8;
let current_word_offset = 7 - (current_input_index % 8);
let start = current_word * 8 + current_word_offset;
let decoded = input.slice_le(start, item_len.try_into().unwrap());
output.append(decoded);
}

current_input_index += item_len.try_into().unwrap();

lazy_index += 1;
}
}

impl RLPItemPartialEq of PartialEq<RLPItem> {
fn eq(lhs: @RLPItem, rhs: @RLPItem) -> bool {
match lhs {
Expand Down
Loading

0 comments on commit 9ca7589

Please sign in to comment.