Generate proofs for gindices

src/fixed_vector.rs

@@ -210,6 +210,14 @@ where
     }
 }
 
+impl<T, N: Unsigned> tree_hash::prototype::MerkleProof for FixedVector<T, N>
+where 
+    T: tree_hash::TreeHash
+{
+    fn compute_proof_for_gindex(&self, gindex: usize) -> Result<Vec<Hash256>, tree_hash::prototype::Error>{
+        crate::tree_hash::generate_proof_for_vec::<T, N>(&self.vec, gindex)
+    }
+}
 impl<T, N: Unsigned> ssz::Encode for FixedVector<T, N>
 where
     T: ssz::Encode,
@@ -569,4 +577,74 @@ mod test {
         let result: Result<FixedVector<u64, U4>, _> = serde_json::from_value(json);
         assert!(result.is_ok());
     }
+
+    #[test]
+    fn merkle_proof_basic() {
+        use tree_hash::prototype::MerkleProof;
+        use typenum::U4;
+
+        let vec: FixedVector<u64, U4> = FixedVector::new(vec![1, 2, 3, 4]).unwrap();
+
+        let proof = vec.compute_proof_for_gindex(1);
+        assert!(proof.is_ok());
+        if let Ok(proof) = proof {
+            assert_eq!(proof.len(), 0);
+        }
+
+        let proof = vec.compute_proof_for_gindex(2);
+        assert!(proof.is_ok());
+        if let Ok(proof) = proof {
+            assert!(!proof.is_empty());
+        }
+
+        let proof = vec.compute_proof_for_gindex(0);
+        assert!(proof.is_err());
+    }
+
+    #[test]
+    fn merkle_proof_complex_types() {
+        use tree_hash::prototype::MerkleProof;
+        use typenum::U2;
+
+        let a1 = A { a: 1, b: 2 };
+        let a2 = A { a: 3, b: 4 };
+        let vec: FixedVector<A, U2> = FixedVector::new(vec![a1, a2]).unwrap();
+
+        let proof = vec.compute_proof_for_gindex(2);
+        assert!(proof.is_ok());
+        if let Ok(proof) = proof {
+            assert!(!proof.is_empty());
+
+            for hash in proof {
+                assert_eq!(hash.len(), 32);
+            }
+        }
+    }
+
+    #[test]
+    fn merkle_proof_tree_depth() {
+        use tree_hash::prototype::MerkleProof;
+        use typenum::U8;
+
+        let vec: FixedVector<u64, U8> = FixedVector::new(vec![1, 2, 3, 4, 5, 6, 7, 8]).unwrap();
+
+        let gindices = vec![1, 2, 3, 4, 5, 6, 7, 8, 15, 16];
+
+        for gindex in gindices {
+            let proof = vec.compute_proof_for_gindex(gindex);
+
+            if gindex == 0 {
+                assert!(proof.is_err());
+            } else if gindex == 1 {
+                if let Ok(proof) = proof {
+                    assert_eq!(proof.len(), 0);
+                }
+            } else {
+                if let Ok(proof) = proof {
+                    let expected_depth = 64 - gindex.leading_zeros() as usize - 1;
+                    assert_eq!(proof.len(), expected_depth);
+                }
+            }
+        }
+    }
 }

src/tree_hash.rs

@@ -6,6 +6,122 @@ use typenum::{
     Unsigned,
 };
 
+pub fn generate_proof_for_vec<T, N>(vec: &[T], gindex: usize) -> Result<Vec<Hash256>, tree_hash::prototype::Error>
+where
+    T: TreeHash,
+    N: Unsigned,
+{
+    let target_size = N::to_usize();
+
+    if gindex == 0 {
+        return Err(tree_hash::prototype::Error::Oops);
+    }
+
+    if target_size == 0 {
+        return Ok(vec![]);
+    }
+
+    generate_proof::<T, N>(vec, gindex)
+}
+
+fn generate_proof<T, N>(vec: &[T], gindex: usize) -> Result<Vec<Hash256>, tree_hash::prototype::Error>
+where
+    T: TreeHash,
+    N: Unsigned,
+{
+    let target_size = N::to_usize();
+    let mut proof = Vec::new();
+    let mut current_gindex = gindex;
+
+    let (_, effective_size) = match T::tree_hash_type() {
+        TreeHashType::Basic => {
+            let chunk_count = (target_size + T::tree_hash_packing_factor() - 1) / T::tree_hash_packing_factor();
+            let padded_count = chunk_count.next_power_of_two();
+            (64 - padded_count.leading_zeros() as usize, padded_count)
+        }
+        _ => {
+            let padded_size = target_size.next_power_of_two();
+            (64 - padded_size.leading_zeros() as usize, padded_size)
+        }
+    };
+
+    while current_gindex > 1 {
+        let is_right_child = current_gindex % 2 == 1;
+        let sibling_gindex = if is_right_child {
+            current_gindex - 1
+        } else {
+            current_gindex + 1
+        };
+
+        let sibling_hash = compute_node_hash_at_gindex::<T, N>(vec, sibling_gindex, effective_size)?;
+        proof.push(sibling_hash);
+
+        current_gindex /= 2;
+    }
+
+    Ok(proof)
+}
+
+fn compute_node_hash_at_gindex<T, N>(
+    vec: &[T], 
+    gindex: usize, 
+    effective_size: usize
+) -> Result<Hash256, tree_hash::prototype::Error>
+where
+    T: TreeHash,
+    N: Unsigned,
+{
+    let target_size = N::to_usize();
+
+    match T::tree_hash_type() {
+        TreeHashType::Basic => {
+            let chunk_count = (target_size + T::tree_hash_packing_factor() - 1) / T::tree_hash_packing_factor();
+
+            if gindex >= effective_size {
+                let chunk_index = gindex - effective_size;
+                if chunk_index < chunk_count {
+                    let start_idx = chunk_index * T::tree_hash_packing_factor();
+                    let end_idx = std::cmp::min(start_idx + T::tree_hash_packing_factor(), vec.len());
+
+                    let mut hasher = MerkleHasher::with_leaves(1);
+                    for j in start_idx..end_idx {
+                        hasher.write(&vec[j].tree_hash_packed_encoding()).map_err(|_| tree_hash::prototype::Error::Oops)?;
+                    }
+                    hasher.finish().map_err(|_| tree_hash::prototype::Error::Oops)
+                } else {
+                    Ok(Hash256::new([0; 32]))
+                }
+            } else {
+                let left_child = gindex * 2;
+                let right_child = gindex * 2 + 1;
+
+                let left_hash = compute_node_hash_at_gindex::<T, N>(vec, left_child, effective_size)?;
+                let right_hash = compute_node_hash_at_gindex::<T, N>(vec, right_child, effective_size)?;
+
+                Ok(tree_hash::merkle_root(&[left_hash.as_slice(), right_hash.as_slice()].concat(), 0))
+            }
+        }
+        _ => {
+            if gindex >= effective_size {
+                let leaf_index = gindex - effective_size;
+                if leaf_index < vec.len() {
+                    Ok(vec[leaf_index].tree_hash_root())
+                } else {
+                    Ok(Hash256::new([0; 32]))
+                }
+            } else {
+                let left_child = gindex * 2;
+                let right_child = gindex * 2 + 1;
+
+                let left_hash = compute_node_hash_at_gindex::<T, N>(vec, left_child, effective_size)?;
+                let right_hash = compute_node_hash_at_gindex::<T, N>(vec, right_child, effective_size)?;
+
+                Ok(tree_hash::merkle_root(&[left_hash.as_slice(), right_hash.as_slice()].concat(), 0))
+            }
+        }
+    }
+}
+
 /// A helper function providing common functionality between the `TreeHash` implementations for
 /// `FixedVector` and `VariableList`.
 pub fn vec_tree_hash_root<T, N>(vec: &[T]) -> Hash256

src/variable_list.rs

@@ -228,6 +228,27 @@ where
     }
 }
 
+impl<T, N: Unsigned> tree_hash::prototype::MerkleProof for VariableList<T, N>
+where 
+    T: tree_hash::TreeHash
+{
+    fn compute_proof_for_gindex(&self, gindex: usize) -> Result<Vec<Hash256>, tree_hash::prototype::Error>{
+        if gindex < 2 {
+            return Err(tree_hash::prototype::Error::Oops);
+        }
+
+        let adjusted_gindex = if gindex == 2 {
+            1 
+        } else if gindex > 2 {
+            gindex - 2 
+        } else {
+            return Err(tree_hash::prototype::Error::Oops);
+        };
+
+        crate::tree_hash::generate_proof_for_vec::<T, N>(&self.vec, adjusted_gindex)
+    }
+}
+
 impl<T, N: Unsigned> ssz::Encode for VariableList<T, N>
 where
     T: ssz::Encode,
@@ -616,4 +637,97 @@ mod test {
         let result: Result<VariableList<u64, U4>, _> = serde_json::from_value(json);
         assert!(result.is_ok());
     }
+
+    #[test]
+    fn merkle_proof_basic() {
+        use tree_hash::prototype::MerkleProof;
+        use typenum::U4;
+
+        let list: VariableList<u64, U4> = VariableList::new(vec![1, 2, 3]).unwrap();
+
+        let proof = list.compute_proof_for_gindex(0);
+        assert!(proof.is_err());
+
+        let proof = list.compute_proof_for_gindex(1);
+        assert!(proof.is_err());
+
+        let proof = list.compute_proof_for_gindex(2);
+        assert!(proof.is_ok());
+        if let Ok(proof) = proof {
+            assert_eq!(proof.len(), 0);
+        }
+
+        let proof = list.compute_proof_for_gindex(4);
+        assert!(proof.is_ok());
+        if let Ok(proof) = proof {
+            assert!(!proof.is_empty());
+        }
+    }
+
+
+    #[test]
+    fn merkle_proof_complex_types() {
+        use tree_hash::prototype::MerkleProof;
+        use typenum::U4;
+
+        // Create a list of composite types
+        let a1 = A { a: 1, b: 2 };
+        let a2 = A { a: 3, b: 4 };
+        let a3 = A { a: 5, b: 6 };
+        let list: VariableList<A, U4> = VariableList::new(vec![a1, a2, a3]).unwrap();
+
+        // Test proof generation for complex types
+        let proof = list.compute_proof_for_gindex(4);
+        assert!(proof.is_ok());
+        if let Ok(proof) = proof {
+            // Verify proof structure
+            assert!(!proof.is_empty());
+
+            // Test that all proof elements are valid Hash256
+            for hash in proof {
+                assert_eq!(hash.len(), 32);
+            }
+        }
+    }
+
+
+
+    #[test]
+    fn merkle_proof_tree_structure() {
+        use tree_hash::prototype::MerkleProof;
+        use typenum::U8;
+
+        let list: VariableList<u64, U8> = VariableList::new(vec![1, 2, 3, 4, 5]).unwrap();
+
+        let test_gindices = vec![2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16];
+
+        for gindex in test_gindices {
+            let proof = list.compute_proof_for_gindex(gindex);
+            assert!(proof.is_ok(), "Failed to generate proof for gindex {}", gindex);
+
+            if let Ok(proof) = proof {
+                let adjusted_gindex = if gindex == 2 {
+                    1
+                } else {
+                    gindex - 2
+                };
+
+                if adjusted_gindex == 1 {
+                    assert_eq!(proof.len(), 0, "Proof should be empty for gindex {} (maps to root)", gindex);
+                } else {
+                    assert!(!proof.is_empty(), "Proof should not be empty for gindex {}", gindex);
+                }
+
+                let adjusted_gindex = if gindex == 2 {
+                    1
+                } else {
+                    gindex - 2
+                };
+                if adjusted_gindex > 1 {
+                    let expected_depth = 64 - adjusted_gindex.leading_zeros() as usize - 1;
+                    assert_eq!(proof.len(), expected_depth, "Incorrect proof length for gindex {}", gindex);
+                }
+            }
+        }
+    }
 }