diff --git a/src/data_structures.cairo b/src/data_structures.cairo index 53542f7..3aff953 100644 --- a/src/data_structures.cairo +++ b/src/data_structures.cairo @@ -1,4 +1,5 @@ mod eth_mpt; +mod stark_mpt; mod mmr; #[cfg(test)] diff --git a/src/data_structures/stark_mpt.cairo b/src/data_structures/stark_mpt.cairo new file mode 100644 index 0000000..2652c98 --- /dev/null +++ b/src/data_structures/stark_mpt.cairo @@ -0,0 +1,119 @@ +use array::SpanTrait; +use traits::{Into, TryInto}; +use option::OptionTrait; +use hash::LegacyHash; +use cairo_lib::utils::types::bitarr::{BitArr, BitArrTryIntoFelt252}; + +struct StarkMPT { + root: felt252 +} + +impl StarkMPTDefault of Default { + fn default() -> StarkMPT { + StarkMPTTrait::new(0) + } +} + +#[derive(Drop)] +enum StarkMPTNode { + // left, right + Binary: (felt252, felt252), + // child, path + Edge: (felt252, BitArr), +} + +#[generate_trait] +impl StarkMPTNodeImpl of StarkMPTNodeTrait { + fn hash(self: @StarkMPTNode) -> felt252 { + match self { + StarkMPTNode::Binary((left, right)) => { + LegacyHash::hash(*left, *right) + }, + StarkMPTNode::Edge((child, path)) => { + let path_felt252: felt252 = (*path).try_into().unwrap(); + LegacyHash::hash(*child, path_felt252) + (*path).len().into() + } + } + } +} + +#[derive(Drop)] +enum Direction { + Left: (), + Right: () +} + +impl BoolIntoDirection of Into { + fn into(self: bool) -> Direction { + if self { + Direction::Right + } else { + Direction::Left + } + } +} + +#[generate_trait] +impl StarkMPTImpl of StarkMPTTrait { + fn new(root: felt252) -> StarkMPT { + StarkMPT { root: root } + } + + fn verify(self: @StarkMPT, key: BitArr, proof: Span) -> Result { + if key.len() != 251 { + return Result::Err('Ill-formed key'); + } + + let mut expected_hash = *self.root; + let mut remaining_path = key; + + let mut i: usize = 0; + loop { + if i == proof.len() { + break Result::Ok(expected_hash); + } + + let node = proof.at(i); + if node.hash() != expected_hash { + break Result::Err('Invalid proof'); + } + + match node { + StarkMPTNode::Binary((left, right)) => { + let direction: Direction = (*remaining_path.pop_front().unwrap()).into(); + + expected_hash = match direction { + Direction::Left => *left, + Direction::Right => *right, + }; + }, + StarkMPTNode::Edge((child, path)) => { + let path_len = (*path).len(); + if path_len > remaining_path.len() { + break Result::Err('Invalid path'); + } + + let mut j: usize = 0; + let valid_path = loop { + if j == path_len { + break true; + } + + if *remaining_path.at(j) != *(*path).at(j) { + break false; + } + + j += 1; + }; + if !valid_path { + break Result::Err('Invalid path'); + } + + expected_hash = *child; + remaining_path = remaining_path.slice(0, path_len); + } + }; + i += 1; + } + } +} diff --git a/src/data_structures/tests.cairo b/src/data_structures/tests.cairo index e66f65d..6770708 100644 --- a/src/data_structures/tests.cairo +++ b/src/data_structures/tests.cairo @@ -1 +1,2 @@ mod test_eth_mpt; +mod test_stark_mpt; diff --git a/src/data_structures/tests/test_stark_mpt.cairo b/src/data_structures/tests/test_stark_mpt.cairo new file mode 100644 index 0000000..cda35bc --- /dev/null +++ b/src/data_structures/tests/test_stark_mpt.cairo @@ -0,0 +1,25 @@ +use cairo_lib::data_structures::stark_mpt::{StarkMPT, StarkMPTTrait, StarkMPTNode, StarkMPTNodeTrait}; +use cairo_lib::utils::types::bitarr::BitArr; +use array::ArrayTrait; + +#[test] +#[available_gas(9999999999)] +fn test_hash_binary_node() { + let binary = StarkMPTNode::Binary((9, 17)); + let expected_hash = 3448800753491155842114129004100047983009754105484160479464353352489980084140; + + assert(binary.hash() == expected_hash, 'Hash does not match'); +} + +#[test] +#[available_gas(9999999999)] +fn test_hash_edge_node() { + let path = array![ + true, false, false, true, false, false, true, true, + true, false, true, false, true, false, false, true + ].span(); + let edge = StarkMPTNode::Edge((291872, path)); + let expected_hash = 800493211047958006469592108402751484180834315522274838721026790014228804959; + + assert(edge.hash() == expected_hash, 'Hash does not match'); +} diff --git a/src/utils/types.cairo b/src/utils/types.cairo index ea4cdf4..31596df 100644 --- a/src/utils/types.cairo +++ b/src/utils/types.cairo @@ -1,5 +1,6 @@ mod bytes; mod byte; +mod bitarr; #[cfg(test)] mod tests; diff --git a/src/utils/types/bitarr.cairo b/src/utils/types/bitarr.cairo new file mode 100644 index 0000000..33c4080 --- /dev/null +++ b/src/utils/types/bitarr.cairo @@ -0,0 +1,28 @@ +use array::SpanTrait; +use traits::TryInto; + +type BitArr = Span; + +impl BitArrTryIntoFelt252 of TryInto { + fn try_into(self: BitArr) -> Option { + if self.len() > 252 { + return Option::None(()); + } + let mut res = 0; + let mut i: usize = 0; + loop { + if i == self.len() { + break Option::Some(res); + } + + let mut bit = 0; + if *self.at(i) { + bit = 1; + } + + res = res * 2 + bit; + + i += 1; + } + } +} diff --git a/src/utils/types/tests.cairo b/src/utils/types/tests.cairo index 399643b..52300e9 100644 --- a/src/utils/types/tests.cairo +++ b/src/utils/types/tests.cairo @@ -1,2 +1,3 @@ mod test_bytes; mod test_byte; +mod test_bitarr; diff --git a/src/utils/types/tests/test_bitarr.cairo b/src/utils/types/tests/test_bitarr.cairo new file mode 100644 index 0000000..aad4834 --- /dev/null +++ b/src/utils/types/tests/test_bitarr.cairo @@ -0,0 +1,43 @@ +use cairo_lib::utils::types::bitarr::{BitArr, BitArrTryIntoFelt252}; +use traits::TryInto; +use option::OptionTrait; +use array::ArrayTrait; + +#[test] +#[available_gas(999999999)] +fn test_bitarr_try_into_felt252() { + let val_0 = array![].span(); + assert(val_0.try_into().unwrap() == 0, '0'); + + let val_1 = array![true].span(); + assert(val_1.try_into().unwrap() == 1, '1'); + + let val_2 = array![true, false].span(); + assert(val_2.try_into().unwrap() == 2, '2'); + + let val_3 = array![true, false, true].span(); + assert(val_3.try_into().unwrap() == 5, '5'); + + let val_4 = array![true, false, true, false].span(); + assert(val_4.try_into().unwrap() == 10, '10'); + + let val_5 = array![true, false, true, false, true].span(); + assert(val_5.try_into().unwrap() == 21, '21'); +} + +#[test] +#[available_gas(999999999)] +fn test_bitarr_try_into_felt252_long() { + let mut arr = ArrayTrait::new(); + let mut i: usize = 0; + loop { + if i == 254 { + break; + } + arr.append(true); + i += 1; + }; + + let val: Option = arr.span().try_into(); + assert(val.is_none(), 'none'); +}