Merge pull request #125 from ralexstokes/even-more-gi

Some more generalized index functionality
ralexstokes · Feb 19, 2024 · 6329986 · 6329986
2 parents e671bd5 + 76e42e2
commit 6329986
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 10 deletions.
diff --git a/ssz-rs/src/array.rs b/ssz-rs/src/array.rs
@@ -106,6 +106,9 @@ where
         if let Some((next, rest)) = path.split_first() {
             match next {
                 PathElement::Index(i) => {
+                    if *i >= N {
+                        return Err(MerkleizationError::InvalidPathElement(next.clone()))
+                    }
                     let chunk_position = i * T::item_length() / 32;
                     let child =
                         parent * get_power_of_two_ceil(Self::chunk_count()) + chunk_position;

diff --git a/ssz-rs/src/bitlist.rs b/ssz-rs/src/bitlist.rs
@@ -184,11 +184,31 @@ impl<const N: usize> Merkleized for Bitlist<N> {
 }
 
 impl<const N: usize> Indexed for Bitlist<N> {
-    fn generalized_index(_path: Path) -> Result<GeneralizedIndex, MerkleizationError>
-    where
-        Self: Sized,
-    {
-        unimplemented!()
+    fn chunk_count() -> usize {
+        Self::chunk_count()
+    }
+
+    fn compute_generalized_index(
+        parent: GeneralizedIndex,
+        path: Path,
+    ) -> Result<GeneralizedIndex, MerkleizationError> {
+        if let Some((next, rest)) = path.split_first() {
+            match next {
+                PathElement::Index(i) => {
+                    if *i >= N {
+                        return Err(MerkleizationError::InvalidPathElement(next.clone()))
+                    }
+                    let chunk_position = i / 256;
+                    let child = parent * get_power_of_two_ceil(<Self as Indexed>::chunk_count()) +
+                        chunk_position;
+                    // NOTE: use `bool` as effective type of element
+                    bool::compute_generalized_index(child, rest)
+                }
+                elem => Err(MerkleizationError::InvalidPathElement(elem.clone())),
+            }
+        } else {
+            Ok(parent)
+        }
     }
 }
 

diff --git a/ssz-rs/src/bitvector.rs b/ssz-rs/src/bitvector.rs
@@ -172,11 +172,31 @@ impl<const N: usize> Merkleized for Bitvector<N> {
 }
 
 impl<const N: usize> Indexed for Bitvector<N> {
-    fn generalized_index(_path: Path) -> Result<GeneralizedIndex, MerkleizationError>
-    where
-        Self: Sized,
-    {
-        unimplemented!()
+    fn chunk_count() -> usize {
+        Self::chunk_count()
+    }
+
+    fn compute_generalized_index(
+        parent: GeneralizedIndex,
+        path: Path,
+    ) -> Result<GeneralizedIndex, MerkleizationError> {
+        if let Some((next, rest)) = path.split_first() {
+            match next {
+                PathElement::Index(i) => {
+                    if *i >= N {
+                        return Err(MerkleizationError::InvalidPathElement(next.clone()))
+                    }
+                    let chunk_position = i / 256;
+                    let child = parent * get_power_of_two_ceil(<Self as Indexed>::chunk_count()) +
+                        chunk_position;
+                    // NOTE: use `bool` as effective type of element
+                    bool::compute_generalized_index(child, rest)
+                }
+                elem => Err(MerkleizationError::InvalidPathElement(elem.clone())),
+            }
+        } else {
+            Ok(parent)
+        }
     }
 }
 

diff --git a/ssz-rs/src/list.rs b/ssz-rs/src/list.rs
@@ -267,6 +267,9 @@ where
         if let Some((next, rest)) = path.split_first() {
             match next {
                 PathElement::Index(i) => {
+                    if *i >= N {
+                        return Err(MerkleizationError::InvalidPathElement(next.clone()))
+                    }
                     let chunk_position = i * T::item_length() / 32;
                     let child =
                         parent * 2 * get_power_of_two_ceil(<Self as Indexed>::chunk_count()) +
@@ -427,4 +430,13 @@ mod tests {
         let bad_input_str = "[1, 2, 3, 4, 5]";
         let _: L = serde_json::from_str(bad_input_str).unwrap();
     }
+
+    #[test]
+    #[should_panic]
+    fn test_illegal_generalized_index() {
+        type L = List<u8, 4>;
+
+        let path = &[5.into()];
+        let _ = L::generalized_index(path).unwrap();
+    }
 }
diff --git a/ssz-rs/src/vector.rs b/ssz-rs/src/vector.rs
@@ -263,6 +263,9 @@ where
         if let Some((next, rest)) = path.split_first() {
             match next {
                 PathElement::Index(i) => {
+                    if *i >= N {
+                        return Err(MerkleizationError::InvalidPathElement(next.clone()))
+                    }
                     let chunk_position = i * T::item_length() / 32;
                     let child =
                         parent * get_power_of_two_ceil(Self::chunk_count()) + chunk_position;
@@ -443,4 +446,13 @@ mod tests {
         let bad_input_str = "[]";
         let _: V = serde_json::from_str(bad_input_str).unwrap();
     }
+
+    #[test]
+    #[should_panic]
+    fn test_illegal_generalized_index() {
+        type V = Vector<u8, 4>;
+
+        let path = &[5.into()];
+        let _ = V::generalized_index(path).unwrap();
+    }
 }