Starknet MPT

src/data_structures.cairo

@@ -1,4 +1,5 @@
 mod eth_mpt;
+mod stark_mpt;
 mod mmr;
 
 #[cfg(test)]

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<StarkMPT> {
+    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<bool, Direction> {
+    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<StarkMPTNode>) -> Result<felt252, felt252> {
+        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;
+        }
+    }
+}

src/data_structures/tests.cairo

@@ -1 +1,2 @@
 mod test_eth_mpt;
+mod test_stark_mpt;

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');
+}

src/utils/types.cairo

@@ -1,5 +1,6 @@
 mod bytes;
 mod byte;
+mod bitarr;
 
 #[cfg(test)]
 mod tests;

src/utils/types/bitarr.cairo

@@ -0,0 +1,28 @@
+use array::SpanTrait;
+use traits::TryInto;
+
+type BitArr = Span<bool>;
+
+impl BitArrTryIntoFelt252 of TryInto<BitArr, felt252> {
+    fn try_into(self: BitArr) -> Option<felt252> {
+        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;
+        }
+    }
+}

src/utils/types/tests.cairo

@@ -1,2 +1,3 @@
 mod test_bytes;
 mod test_byte;
+mod test_bitarr;

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<felt252> = arr.span().try_into();
+    assert(val.is_none(), 'none');
+}