add MastNodeInfo method

0xPolygonMiden · Jul 3, 2024 · df621c5 · df621c5
1 parent f9a3a0b
commit df621c5
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 77 deletions.
diff --git a/core/src/mast/serialization/info.rs b/core/src/mast/serialization/info.rs
@@ -1,19 +1,18 @@
 use miden_crypto::hash::rpo::RpoDigest;
 use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
-use crate::mast::{MastNode, MerkleTreeNode};
+use crate::mast::{MastForest, MastNode, MastNodeId, MerkleTreeNode};
 
-use super::DataOffset;
+use super::{basic_block_data_decoder::BasicBlockDataDecoder, DataOffset};
 
 // MAST NODE INFO
 // ===============================================================================================
 
 #[derive(Debug)]
 pub struct MastNodeInfo {
-    // TODOP: Remove pub(super)?
-    pub(super) ty: MastNodeType,
-    pub(super) offset: DataOffset,
-    pub(super) digest: RpoDigest,
+    ty: MastNodeType,
+    offset: DataOffset,
+    digest: RpoDigest,
 }
 
 impl MastNodeInfo {
@@ -32,6 +31,68 @@ impl MastNodeInfo {
             digest: mast_node.digest(),
         }
     }
+
+    pub fn try_into_mast_node(
+        self,
+        mast_forest: &MastForest,
+        basic_block_data_decoder: &mut BasicBlockDataDecoder,
+    ) -> Result<MastNode, DeserializationError> {
+        let mast_node = match self.ty {
+            MastNodeType::Block {
+                len: num_operations_and_decorators,
+            } => {
+                let (operations, decorators) = basic_block_data_decoder
+                    .decode_operations_and_decorators(num_operations_and_decorators)?;
+
+                Ok(MastNode::new_basic_block_with_decorators(operations, decorators))
+            }
+            MastNodeType::Join {
+                left_child_id,
+                right_child_id,
+            } => {
+                let left_child = MastNodeId::from_u32_safe(left_child_id, mast_forest)?;
+                let right_child = MastNodeId::from_u32_safe(right_child_id, mast_forest)?;
+
+                Ok(MastNode::new_join(left_child, right_child, mast_forest))
+            }
+            MastNodeType::Split {
+                if_branch_id,
+                else_branch_id,
+            } => {
+                let if_branch = MastNodeId::from_u32_safe(if_branch_id, mast_forest)?;
+                let else_branch = MastNodeId::from_u32_safe(else_branch_id, mast_forest)?;
+
+                Ok(MastNode::new_split(if_branch, else_branch, mast_forest))
+            }
+            MastNodeType::Loop { body_id } => {
+                let body_id = MastNodeId::from_u32_safe(body_id, mast_forest)?;
+
+                Ok(MastNode::new_loop(body_id, mast_forest))
+            }
+            MastNodeType::Call { callee_id } => {
+                let callee_id = MastNodeId::from_u32_safe(callee_id, mast_forest)?;
+
+                Ok(MastNode::new_call(callee_id, mast_forest))
+            }
+            MastNodeType::SysCall { callee_id } => {
+                let callee_id = MastNodeId::from_u32_safe(callee_id, mast_forest)?;
+
+                Ok(MastNode::new_syscall(callee_id, mast_forest))
+            }
+            MastNodeType::Dyn => Ok(MastNode::new_dynexec()),
+            MastNodeType::External => Ok(MastNode::new_external(self.digest)),
+        }?;
+
+        if mast_node.digest() == self.digest {
+            Ok(mast_node)
+        } else {
+            Err(DeserializationError::InvalidValue(format!(
+                "MastNodeInfo's digest '{}' doesn't match deserialized MastNode's digest '{}'",
+                self.digest,
+                mast_node.digest()
+            )))
+        }
+    }
 }
 
 impl Serializable for MastNodeInfo {

diff --git a/core/src/mast/serialization/mod.rs b/core/src/mast/serialization/mod.rs
@@ -1,14 +1,12 @@
 use alloc::vec::Vec;
 use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
-use crate::mast::MerkleTreeNode;
-
 use super::{MastForest, MastNode, MastNodeId};
 
 mod decorator;
 
 mod info;
-use info::{MastNodeInfo, MastNodeType};
+use info::MastNodeInfo;
 
 mod basic_block_data_builder;
 use basic_block_data_builder::BasicBlockDataBuilder;
@@ -149,11 +147,9 @@ impl Deserializable for MastForest {
             let mut mast_forest = MastForest::new();
 
             for mast_node_info in mast_node_infos {
-                let node = try_info_to_mast_node(
-                    mast_node_info,
-                    &mast_forest,
-                    &mut basic_block_data_decoder,
-                )?;
+                let node = mast_node_info
+                    .try_into_mast_node(&mast_forest, &mut basic_block_data_decoder)?;
+
                 mast_forest.add_node(node);
             }
 
@@ -167,66 +163,3 @@ impl Deserializable for MastForest {
         Ok(mast_forest)
     }
 }
-
-// TODOP: Make `MastNodeInfo` method
-fn try_info_to_mast_node(
-    mast_node_info: MastNodeInfo,
-    mast_forest: &MastForest,
-    basic_block_data_decoder: &mut BasicBlockDataDecoder,
-) -> Result<MastNode, DeserializationError> {
-    let mast_node = match mast_node_info.ty {
-        MastNodeType::Block {
-            len: num_operations_and_decorators,
-        } => {
-            let (operations, decorators) = basic_block_data_decoder
-                .decode_operations_and_decorators(num_operations_and_decorators)?;
-
-            Ok(MastNode::new_basic_block_with_decorators(operations, decorators))
-        }
-        MastNodeType::Join {
-            left_child_id,
-            right_child_id,
-        } => {
-            let left_child = MastNodeId::from_u32_safe(left_child_id, mast_forest)?;
-            let right_child = MastNodeId::from_u32_safe(right_child_id, mast_forest)?;
-
-            Ok(MastNode::new_join(left_child, right_child, mast_forest))
-        }
-        MastNodeType::Split {
-            if_branch_id,
-            else_branch_id,
-        } => {
-            let if_branch = MastNodeId::from_u32_safe(if_branch_id, mast_forest)?;
-            let else_branch = MastNodeId::from_u32_safe(else_branch_id, mast_forest)?;
-
-            Ok(MastNode::new_split(if_branch, else_branch, mast_forest))
-        }
-        MastNodeType::Loop { body_id } => {
-            let body_id = MastNodeId::from_u32_safe(body_id, mast_forest)?;
-
-            Ok(MastNode::new_loop(body_id, mast_forest))
-        }
-        MastNodeType::Call { callee_id } => {
-            let callee_id = MastNodeId::from_u32_safe(callee_id, mast_forest)?;
-
-            Ok(MastNode::new_call(callee_id, mast_forest))
-        }
-        MastNodeType::SysCall { callee_id } => {
-            let callee_id = MastNodeId::from_u32_safe(callee_id, mast_forest)?;
-
-            Ok(MastNode::new_syscall(callee_id, mast_forest))
-        }
-        MastNodeType::Dyn => Ok(MastNode::new_dynexec()),
-        MastNodeType::External => Ok(MastNode::new_external(mast_node_info.digest)),
-    }?;
-
-    if mast_node.digest() == mast_node_info.digest {
-        Ok(mast_node)
-    } else {
-        Err(DeserializationError::InvalidValue(format!(
-            "MastNodeInfo's digest '{}' doesn't match deserialized MastNode's digest '{}'",
-            mast_node_info.digest,
-            mast_node.digest()
-        )))
-    }
-}