Implement encoding for Node.

divviup · Apr 16, 2024 · f39c1fc · f39c1fc
1 parent 7d89075
commit f39c1fc
Showing 1 changed file with 35 additions and 15 deletions.
diff --git a/src/bt.rs b/src/bt.rs
@@ -53,7 +53,7 @@
 #![allow(unused_variables)]
 
 use core::fmt::Debug;
-use std::io::Cursor;
+use std::io::{Cursor, Seek, SeekFrom};
 
 use bitvec::slice::BitSlice;
 
@@ -240,10 +240,9 @@ impl<V> Default for BinaryTree<V> {
     }
 }
 
-impl<V: Encode> Encode for BinaryTree<V> {
+impl<V: Encode> Encode for Node<V> {
     fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
-        let mut stack = Vec::new();
-        stack.push(&self.root);
+        let mut stack = vec![Some(self)];
 
         // Nodes are stored following a pre-order traversal.
         while let Some(elem) = stack.pop() {
@@ -252,8 +251,8 @@ impl<V: Encode> Encode for BinaryTree<V> {
                 Some(node) => {
                     CodecMarker::Inner.encode(bytes)?;
                     node.value.encode(bytes)?;
-                    stack.push(&node.right);
-                    stack.push(&node.left);
+                    stack.push(node.right.as_deref());
+                    stack.push(node.left.as_deref());
                 }
             }
         }
@@ -262,12 +261,10 @@ impl<V: Encode> Encode for BinaryTree<V> {
     }
 }
 
-impl<V: ParameterizedDecode<P>, P> ParameterizedDecode<P> for BinaryTree<V> {
+impl<P, V: ParameterizedDecode<P>> ParameterizedDecode<P> for Node<V> {
     fn decode_with_param(param: &P, bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         let mut root = None;
-
-        let mut stack = Vec::new();
-        stack.push(&mut root);
+        let mut stack = vec![&mut root];
 
         // Decode nodes in a pre-order traversal.
         while let Some(elem) = stack.pop() {
@@ -282,7 +279,10 @@ impl<V: ParameterizedDecode<P>, P> ParameterizedDecode<P> for BinaryTree<V> {
             };
         }
 
-        Ok(Self { root })
+        match root {
+            Some(node) => Ok(*node),
+            None => Err(CodecError::UnexpectedValue),
+        }
     }
 }
 
@@ -295,13 +295,11 @@ impl<V: core::fmt::Display> core::fmt::Display for Node<V> {
             node: Option<&'a Node<V>>,
         }
 
-        let mut stack = Vec::new();
-
-        stack.push(Item {
+        let mut stack = vec![Item {
             name: String::default(),
             level: 0,
             node: Some(self),
-        });
+        }];
 
         while let Some(Item { name, level, node }) = stack.pop() {
             if let Some(Node { value, left, right }) = node {
@@ -328,6 +326,28 @@ impl<V: core::fmt::Display> core::fmt::Display for Node<V> {
     }
 }
 
+impl<V: Encode> Encode for BinaryTree<V> {
+    fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
+        match &self.root {
+            None => CodecMarker::Leaf.encode(bytes),
+            Some(node) => node.encode(bytes),
+        }
+    }
+}
+
+impl<P, V: ParameterizedDecode<P>> ParameterizedDecode<P> for BinaryTree<V> {
+    fn decode_with_param(param: &P, bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
+        match CodecMarker::decode(bytes)? {
+            CodecMarker::Leaf => Ok(Self::default()),
+            CodecMarker::Inner => {
+                bytes.seek(SeekFrom::Start(0))?;
+                let root = Some(Box::new(Node::decode_with_param(param, bytes)?));
+                Ok(Self { root })
+            }
+        }
+    }
+}
+
 #[cfg(feature = "test-util")]
 impl<V> core::fmt::Display for BinaryTree<V>
 where