From 9e59aece8cdfda3528e9f7f2f4ebe5a70e87d589 Mon Sep 17 00:00:00 2001
From: Bobbin Threadbare <bobbinth@protonmail.com>
Date: Sat, 20 Jul 2024 17:35:32 -0700
Subject: [PATCH] fix: remove implicit serialization of MastNodeId

---
 core/src/mast/mod.rs               | 24 ++++++++++--------------
 core/src/mast/serialization/mod.rs | 10 +++++-----
 2 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/core/src/mast/mod.rs b/core/src/mast/mod.rs
index feb1b841b3..20d1ed9309 100644
--- a/core/src/mast/mod.rs
+++ b/core/src/mast/mod.rs
@@ -8,7 +8,7 @@ pub use node::{
     BasicBlockNode, CallNode, DynNode, ExternalNode, JoinNode, LoopNode, MastNode, OpBatch,
     OperationOrDecorator, SplitNode, OP_BATCH_SIZE, OP_GROUP_SIZE,
 };
-use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
+use winter_utils::DeserializationError;
 
 mod serialization;
 
@@ -152,25 +152,21 @@ impl MastNodeId {
     }
 }
 
-impl fmt::Display for MastNodeId {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "MastNodeId({})", self.0)
+impl From<MastNodeId> for u32 {
+    fn from(value: MastNodeId) -> Self {
+        value.0
     }
 }
 
-impl Serializable for MastNodeId {
-    fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        self.0.write_into(target)
+impl From<&MastNodeId> for u32 {
+    fn from(value: &MastNodeId) -> Self {
+        value.0
     }
 }
 
-impl Deserializable for MastNodeId {
-    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let inner = source.read_u32()?;
-
-        // TODO: fix
-
-        Ok(Self(inner))
+impl fmt::Display for MastNodeId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "MastNodeId({})", self.0)
     }
 }
 
diff --git a/core/src/mast/serialization/mod.rs b/core/src/mast/serialization/mod.rs
index e23cdbd2fb..40b32c0923 100644
--- a/core/src/mast/serialization/mod.rs
+++ b/core/src/mast/serialization/mod.rs
@@ -54,7 +54,8 @@ impl Serializable for MastForest {
         target.write_usize(self.nodes.len());
 
         // roots
-        self.roots.write_into(target);
+        let roots: Vec<u32> = self.roots.iter().map(u32::from).collect();
+        roots.write_into(target);
 
         // Prepare MAST node infos, but don't store them yet. We store them at the end to make
         // deserialization more efficient.
@@ -102,11 +103,8 @@ impl Deserializable for MastForest {
         }
 
         let node_count = source.read_usize()?;
-
-        let roots: Vec<MastNodeId> = Deserializable::read_from(source)?;
-
+        let roots: Vec<u32> = Deserializable::read_from(source)?;
         let strings: Vec<DataOffset> = Deserializable::read_from(source)?;
-
         let data: Vec<u8> = Deserializable::read_from(source)?;
 
         let basic_block_data_decoder = BasicBlockDataDecoder::new(&data, &strings);
@@ -128,6 +126,8 @@ impl Deserializable for MastForest {
             }
 
             for root in roots {
+                // make sure the root is valid in the context of the MAST forest
+                let root = MastNodeId::from_u32_safe(root, &mast_forest)?;
                 mast_forest.make_root(root);
             }