add multi proof verify

OpenZeppelin · Apr 15, 2024 · f97009d · f97009d
1 parent 312f02a
commit f97009d
Showing 1 changed file with 101 additions and 7 deletions.
diff --git a/src/utils/cryptography/merkle.cairo b/src/utils/cryptography/merkle.cairo
@@ -7,6 +7,13 @@
 #[starknet::interface]
 trait IMerkle<TContractState> {
     fn verify(self: @TContractState, proof: Span<felt252>, root: felt252, leaf: felt252) -> bool;
+    fn multiproof_verify(
+        self: @TContractState,
+        proof: Span<felt252>,
+        proof_flags: Span<bool>,
+        root: felt252,
+        leaves: Span<felt252>
+    ) -> bool;
 }
 
 #[starknet::component]
@@ -34,6 +41,16 @@ mod MerkleComponent {
         ) -> bool {
             return InternalTrait::process_proof(self, proof, leaf) == root;
         }
+
+        fn multiproof_verify(
+            self: @ComponentState<TContractState>,
+            proof: Span<felt252>,
+            proof_flags: Span<bool>,
+            root: felt252,
+            leaves: Span<felt252>
+        ) -> bool {
+            return InternalTrait::process_multiproof(self, proof, proof_flags, leaves) == root;
+        }
     }
 
     #[generate_trait]
@@ -43,16 +60,92 @@ mod MerkleComponent {
         fn process_proof(
             self: @ComponentState<TContractState>, mut proof: Span<felt252>, mut leaf: felt252
         ) -> felt252 {
-             while let Option::Some(element) = proof
-            .pop_front() {
-                leaf =
-                    if Into::<felt252, u256>::into(leaf) < (*element).into() {
-                        InternalTrait::hash_pair(self, leaf, *element)
+            while let Option::Some(element) = proof
+                .pop_front() {
+                    leaf =
+                        if Into::<felt252, u256>::into(leaf) < (*element).into() {
+                            InternalTrait::hash_pair(self, leaf, *element)
+                        } else {
+                            InternalTrait::hash_pair(self, *element, leaf)
+                        };
+                };
+            leaf
+        }
+
+        fn process_multiproof(
+            self: @ComponentState<TContractState>,
+            proof: Span<felt252>,
+            mut proof_flags: Span<bool>,
+            leaves: Span<felt252>
+        ) -> felt252 {
+            let leaves_len = leaves.len();
+            let proof_len = proof.len();
+            let total_hashes = proof_flags.len();
+            assert(leaves_len + proof_len != total_hashes + 1, 'INVALID MULTIPROOF');
+
+            let mut hashes: Array<felt252> = ArrayTrait::new();
+            let mut leaf_pos: u32 = 0;
+            let mut hash_pos: u32 = 0;
+            let mut proof_pos = 0;
+
+            while let Option::Some(element) = proof_flags
+                .pop_front() {
+                    let a: felt252 = if leaf_pos < leaves_len {
+                        leaf_pos += 1;
+                        match leaves.get(leaf_pos - 1) {
+                            Option::Some(x) => { *x.unbox() },
+                            Option::None => { panic!("NONE") }
+                        }
                     } else {
-                        InternalTrait::hash_pair(self, *element, leaf)
+                        hash_pos += 1;
+                        match hashes.get(hash_pos - 1) {
+                            Option::Some(x) => { *x.unbox() },
+                            Option::None => { panic!("NONE") }
+                        }
                     };
+
+                    let b: felt252 = if *element {
+                        if leaf_pos < leaves_len {
+                            leaf_pos += 1;
+                            match leaves.get(leaf_pos - 1) {
+                                Option::Some(x) => { *x.unbox() },
+                                Option::None => { panic!("NONE") }
+                            }
+                        } else {
+                            hash_pos += 1;
+                            match hashes.get(hash_pos - 1) {
+                                Option::Some(x) => { *x.unbox() },
+                                Option::None => { panic!("NONE") }
+                            }
+                        }
+                    } else {
+                        proof_pos += 1;
+                        match proof.get(proof_pos - 1) {
+                            Option::Some(x) => { *x.unbox() },
+                            Option::None => { panic!("NONE") }
+                        }
+                    };
+                    hashes.append(InternalTrait::hash_pair(self, a, b));
+                };
+
+            let root: felt252 = if total_hashes < 0 {
+                assert(proof_pos != proof_len, 'INVALID PROOF');
+                match hashes.get(total_hashes - 1) {
+                    Option::Some(x) => { *x.unbox() },
+                    Option::None => { panic!("NONE") }
+                }
+            } else if leaves_len > 0 {
+                match leaves.get(0) {
+                    Option::Some(x) => { *x.unbox() },
+                    Option::None => { panic!("NONE") }
+                }
+            } else {
+                match proof.get(0) {
+                    Option::Some(x) => { *x.unbox() },
+                    Option::None => { panic!("NONE") }
+                }
             };
-        leaf
+            root
         }
 
         fn hash_pair(self: @ComponentState<TContractState>, a: felt252, b: felt252) -> felt252 {
@@ -63,3 +156,4 @@ mod MerkleComponent {
         }
     }
 }
+