chore: building original skeleton

Cargo.toml

@@ -13,6 +13,7 @@ license = "Apache-2.0"
 license-file = "LICENSE"
 
 [workspace.dependencies]
+bisection = "0.1.0"
 clap = { version = "4.5.4", features = ["cargo", "derive"] }
 derive_more = "0.99.17"
 pretty_assertions = "1.2.1"

crates/committer/Cargo.toml

@@ -14,6 +14,7 @@ pretty_assertions.workspace = true
 rstest.workspace = true
 
 [dependencies]
+bisection.workspace = true
 derive_more.workspace = true
 rayon.workspace = true
 serde.workspace = true

crates/committer/src/patricia_merkle_tree.rs

@@ -1,6 +1,7 @@
 pub mod errors;
 pub mod filled_node;
 pub mod filled_tree;
+pub mod original_skeleton_calc;
 pub mod original_skeleton_node;
 pub mod original_skeleton_tree;
 pub mod serialized_node;

crates/committer/src/patricia_merkle_tree/errors.rs

@@ -1,7 +1,22 @@
-// TODO(Amos, 01/04/2024): Add error types.
+use thiserror::Error;
 
-#[derive(Debug)]
-pub(crate) enum OriginalSkeletonTreeError {}
+use crate::storage::errors::StorageError;
+use crate::storage::storage_trait::StorageValue;
+
+// TODO(Amos, 01/04/2024): Add error types.
+#[allow(dead_code)]
+#[derive(Debug, Error)]
+pub(crate) enum OriginalSkeletonTreeError {
+    #[error(
+        "Failed to deserialize the storage value: {0:?} while building the original skeleton tree."
+    )]
+    Deserialization(StorageValue),
+    #[error(
+        "Unable to read from storage the storage key: {0:?} while building the \
+         original skeleton tree."
+    )]
+    StorageRead(#[from] StorageError),
+}
 
 #[derive(Debug)]
 #[allow(dead_code)]

crates/committer/src/patricia_merkle_tree/filled_node.rs

@@ -1,10 +1,16 @@
+use crate::hash::hash_trait::HashOutput;
 use crate::patricia_merkle_tree::errors::FilledTreeError;
 use crate::patricia_merkle_tree::filled_tree::FilledTreeResult;
+use crate::patricia_merkle_tree::original_skeleton_tree::OriginalSkeletonTreeResult;
 use crate::patricia_merkle_tree::serialized_node::{
     LeafCompiledClassToSerialize, SerializeNode, SERIALIZE_HASH_BYTES,
 };
+use crate::patricia_merkle_tree::serialized_node::{BINARY_BYTES, EDGE_BYTES, EDGE_PATH_BYTES};
 use crate::patricia_merkle_tree::types::{EdgeData, LeafDataTrait};
-use crate::{hash::hash_trait::HashOutput, types::Felt};
+use crate::patricia_merkle_tree::types::{EdgePath, EdgePathLength, PathToBottom};
+use crate::storage::storage_trait::{StorageKey, StorageValue};
+use crate::types::Felt;
+
 // TODO(Nimrod, 1/6/2024): Swap to starknet-types-core types once implemented.
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@@ -41,7 +47,7 @@ pub(crate) struct BinaryData {
 }
 
 #[allow(dead_code)]
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub(crate) enum LeafData {
     StorageValue(Felt),
     CompiledClassHash(ClassHash),
@@ -70,6 +76,52 @@ impl LeafDataTrait for LeafData {
     }
 }
 
+#[allow(dead_code)]
+impl FilledNode<LeafData> {
+    /// Deserializes non-leaf nodes; if a serialized leaf node is given, the hash
+    /// is used but the data is ignored.
+    pub(crate) fn deserialize(
+        key: &StorageKey,
+        value: &StorageValue,
+    ) -> OriginalSkeletonTreeResult<Self> {
+        if value.0.len() == BINARY_BYTES {
+            Ok(Self {
+                hash: HashOutput(Felt::from_bytes_be_slice(&key.0)),
+                data: NodeData::Binary(BinaryData {
+                    left_hash: HashOutput(Felt::from_bytes_be_slice(
+                        &value.0[..SERIALIZE_HASH_BYTES],
+                    )),
+                    right_hash: HashOutput(Felt::from_bytes_be_slice(
+                        &value.0[SERIALIZE_HASH_BYTES..],
+                    )),
+                }),
+            })
+        } else if value.0.len() == EDGE_BYTES {
+            return Ok(Self {
+                hash: HashOutput(Felt::from_bytes_be_slice(&key.0)),
+                data: NodeData::Edge(EdgeData {
+                    bottom_hash: HashOutput(Felt::from_bytes_be_slice(
+                        &value.0[..SERIALIZE_HASH_BYTES],
+                    )),
+                    path_to_bottom: PathToBottom {
+                        path: EdgePath(Felt::from_bytes_be_slice(
+                            &value.0[SERIALIZE_HASH_BYTES..SERIALIZE_HASH_BYTES + EDGE_PATH_BYTES],
+                        )),
+                        length: EdgePathLength(value.0[EDGE_BYTES - 1]),
+                    },
+                }),
+            });
+        } else {
+            // TODO(Nimrod, 5/5/2024): See if deserializing leaves data is needed somewhere.
+            return Ok(Self {
+                hash: HashOutput(Felt::from_bytes_be_slice(&key.0)),
+                // Dummy value which will be ignored.
+                data: NodeData::Leaf(LeafData::StorageValue(Felt::ZERO)),
+            });
+        }
+    }
+}
+
 impl LeafData {
     /// Serializes the leaf data into a byte vector.
     /// The serialization is done as follows:

crates/committer/src/patricia_merkle_tree/original_skeleton_calc.rs

@@ -0,0 +1,201 @@
+use crate::hash::hash_trait::HashOutput;
+use crate::patricia_merkle_tree::filled_node::BinaryData;
+use crate::patricia_merkle_tree::filled_node::FilledNode;
+use crate::patricia_merkle_tree::filled_node::NodeData;
+use crate::patricia_merkle_tree::original_skeleton_tree::OriginalSkeletonTreeResult;
+use crate::patricia_merkle_tree::types::EdgeData;
+use crate::patricia_merkle_tree::types::PathToBottom;
+use crate::patricia_merkle_tree::types::TreeHeight;
+use crate::patricia_merkle_tree::{
+    filled_node::LeafData, original_skeleton_node::OriginalSkeletonNode, types::NodeIndex,
+};
+use crate::storage::errors::StorageError;
+use crate::storage::storage_trait::Storage;
+use crate::storage::storage_trait::StorageKey;
+use crate::storage::storage_trait::StoragePrefix;
+use crate::types::Felt;
+use bisection::{bisect_left, bisect_right};
+use std::collections::HashMap;
+
+#[cfg(test)]
+#[path = "original_skeleton_calc_test.rs"]
+pub mod original_skeleton_calc_test;
+
+#[allow(dead_code)]
+pub(crate) struct OriginalSkeletonTreeImpl {
+    nodes: HashMap<NodeIndex, OriginalSkeletonNode<LeafData>>,
+    // TODO(Nimrod, 30/4/2024): Remove `leaf_modifications` field.
+    leaf_modifications: HashMap<NodeIndex, LeafData>,
+    tree_height: TreeHeight,
+}
+
+struct SubTree<'a> {
+    pub sorted_leaf_indices: &'a [NodeIndex],
+    pub root_index: NodeIndex,
+    pub root_hash: HashOutput,
+}
+
+impl<'a> SubTree<'a> {
+    pub(crate) fn get_height(&self, total_tree_height: &TreeHeight) -> TreeHeight {
+        TreeHeight(total_tree_height.0 - self.root_index.0.bits() + 1)
+    }
+
+    pub(crate) fn split_leaves(
+        &self,
+        total_tree_height: &TreeHeight,
+    ) -> (&'a [NodeIndex], &'a [NodeIndex]) {
+        let height = self.get_height(total_tree_height);
+        let leftmost_index_in_right_subtree = ((self.root_index.times_two_to_the_power(1))
+            + NodeIndex(Felt::ONE))
+        .times_two_to_the_power(height.0 - 1);
+        let mid = bisect_left(self.sorted_leaf_indices, &leftmost_index_in_right_subtree);
+        (
+            &self.sorted_leaf_indices[..mid],
+            &self.sorted_leaf_indices[mid..],
+        )
+    }
+
+    pub(crate) fn is_sibling(&self) -> bool {
+        self.sorted_leaf_indices.is_empty()
+    }
+
+    fn get_bottom_subtree(
+        &self,
+        path_to_bottom: &PathToBottom,
+        total_tree_height: &TreeHeight,
+        bottom_hash: HashOutput,
+    ) -> Self {
+        let bottom_index = path_to_bottom.bottom_index(self.root_index);
+        let bottom_height =
+            self.get_height(total_tree_height) - TreeHeight(path_to_bottom.length.0);
+        let leftmost_in_subtree = bottom_index.times_two_to_the_power(bottom_height.0);
+        let rightmost_in_subtree = leftmost_in_subtree
+            + (NodeIndex(Felt::ONE).times_two_to_the_power(bottom_height.0))
+            - NodeIndex(Felt::ONE);
+        let bottom_leaves =
+            &self.sorted_leaf_indices[bisect_left(self.sorted_leaf_indices, &leftmost_in_subtree)
+                ..bisect_right(self.sorted_leaf_indices, &rightmost_in_subtree)];
+
+        Self {
+            sorted_leaf_indices: bottom_leaves,
+            root_index: bottom_index,
+            root_hash: bottom_hash,
+        }
+    }
+
+    fn get_children_subtrees(
+        &self,
+        left_hash: HashOutput,
+        right_hash: HashOutput,
+        total_tree_height: &TreeHeight,
+    ) -> (Self, Self) {
+        let (left_leaves, right_leaves) = self.split_leaves(total_tree_height);
+        let left_root_index = self.root_index * Felt::TWO;
+        (
+            SubTree {
+                sorted_leaf_indices: left_leaves,
+                root_index: left_root_index,
+                root_hash: left_hash,
+            },
+            SubTree {
+                sorted_leaf_indices: right_leaves,
+                root_index: left_root_index + NodeIndex(Felt::ONE),
+                root_hash: right_hash,
+            },
+        )
+    }
+}
+
+#[allow(dead_code)]
+impl OriginalSkeletonTreeImpl {
+    /// Fetches the Patricia witnesses, required to build the original skeleton tree from storage.
+    /// Given a list of subtrees, traverses towards their leaves and fetches all non-empty and
+    /// sibling nodes. Assumes no subtrees of height 0 (leaves).
+
+    fn fetch_nodes(
+        &mut self,
+        subtrees: Vec<SubTree<'_>>,
+        storage: &impl Storage,
+    ) -> OriginalSkeletonTreeResult<()> {
+        if subtrees.is_empty() {
+            return Ok(());
+        }
+        let mut next_subtrees = Vec::new();
+        let subtrees_roots =
+            OriginalSkeletonTreeImpl::calculate_subtrees_roots(&subtrees, storage)?;
+        for (filled_node, subtree) in subtrees_roots.into_iter().zip(subtrees.iter()) {
+            match filled_node.data {
+                // Binary node.
+                NodeData::Binary(BinaryData {
+                    left_hash,
+                    right_hash,
+                }) => {
+                    if subtree.is_sibling() {
+                        self.nodes.insert(
+                            subtree.root_index,
+                            OriginalSkeletonNode::LeafOrBinarySibling(filled_node.hash),
+                        );
+                        continue;
+                    }
+                    self.nodes
+                        .insert(subtree.root_index, OriginalSkeletonNode::Binary);
+                    let (left_subtree, right_subtree) =
+                        subtree.get_children_subtrees(left_hash, right_hash, &self.tree_height);
+                    next_subtrees.extend(vec![left_subtree, right_subtree]);
+                }
+                // Edge node.
+                NodeData::Edge(EdgeData {
+                    bottom_hash,
+                    path_to_bottom,
+                }) => {
+                    if subtree.is_sibling() {
+                        // Sibling will remain an edge node. No need to open the bottom.
+                        self.nodes.insert(
+                            subtree.root_index,
+                            OriginalSkeletonNode::EdgeSibling(EdgeData {
+                                bottom_hash,
+                                path_to_bottom,
+                            }),
+                        );
+                        continue;
+                    }
+                    // Parse bottom.
+                    let bottom_subtree =
+                        subtree.get_bottom_subtree(&path_to_bottom, &self.tree_height, bottom_hash);
+                    self.nodes.insert(
+                        subtree.root_index,
+                        OriginalSkeletonNode::Edge { path_to_bottom },
+                    );
+                    next_subtrees.push(bottom_subtree);
+                }
+                // Leaf node.
+                NodeData::Leaf(_) => {
+                    if subtree.is_sibling() {
+                        self.nodes.insert(
+                            subtree.root_index,
+                            OriginalSkeletonNode::LeafOrBinarySibling(filled_node.hash),
+                        );
+                    }
+                }
+            }
+        }
+        self.fetch_nodes(next_subtrees, storage)
+    }
+
+    fn calculate_subtrees_roots(
+        subtrees: &[SubTree<'_>],
+        storage: &impl Storage,
+    ) -> OriginalSkeletonTreeResult<Vec<FilledNode<LeafData>>> {
+        let mut subtrees_roots = vec![];
+        for subtree in subtrees.iter() {
+            let key =
+                StorageKey::from(subtree.root_hash.0).with_prefix(StoragePrefix::PatriciaNode);
+            let val = storage.get(&key).ok_or(StorageError::MissingKey(key))?;
+            subtrees_roots.push(FilledNode::deserialize(
+                &StorageKey::from(subtree.root_hash.0),
+                val,
+            )?)
+        }
+        Ok(subtrees_roots)
+    }
+}