1 files changed, 49 insertions, 21 deletions

diff --git a/forest/src/default.rs b/forest/src/default.rs
index 9295f1a..b96439f 100644
--- a/forest/src/default.rs
+++ b/forest/src/default.rs
@@ -24,6 +24,8 @@ pub enum Error {
     /// Encounter an invalid error in converting from an error of
     /// graphs.
     InvalidGraphError(GError),
+    /// Encounter an error when converting forest labels.
+    LabelConversion(ForestLabelError),
 }
 
 impl Display for Error {
@@ -35,6 +37,7 @@ impl Display for Error {
             Error::DuplicatedNode(n) => write!(f, "node {n} is duplicated"),
             Error::InvalidGraphError(ge) => write!(f, "invalid error: {ge}"),
             Error::NodeNoLabel(n) => write!(f, "node {n} has no labels, but it should have one"),
+            Error::LabelConversion(le) => write!(f, "fail to convert labels: {le}"),
         }
     }
 }
@@ -51,6 +54,12 @@ impl From<GError> for Error {
     }
 }
 
+impl From<ForestLabelError> for Error {
+    fn from(le: ForestLabelError) -> Self {
+        Self::LabelConversion(le)
+    }
+}
+
 /// A default implementation of forest.
 #[derive(Debug, Clone)]
 pub struct DefaultForest<T: GraphLabel> {
@@ -193,7 +202,7 @@ impl<T: GraphLabel> LabelGraph<T> for DefaultForest<T> {
     }
 }
 
-impl<T: GraphLabel> Forest<T> for DefaultForest<T> {
+impl<T: GraphLabel> Forest<T> for DefaultForest<ForestLabel<T>> {
     type Error = Error;
 
     fn root(&self) -> Option<usize> {
@@ -205,7 +214,7 @@ impl<T: GraphLabel> Forest<T> for DefaultForest<T> {
 
         let mut builder = PLGBuilderMut::from_graph_mut(&mut graph);
 
-        let root = Some(builder.add_vertex(label));
+        let root = Some(builder.add_vertex(ForestLabel::from(label)));
 
         Self { graph, root }
     }
@@ -299,12 +308,12 @@ impl<T: GraphLabel> Forest<T> for DefaultForest<T> {
 
         let nodes_len = fragment.nodes_len();
 
-        // If the fragment root has a duplicate label, the forest will
-        // not grow, so we use the label to find the adjoined node
-        // index.
-
-        // the nodes hava already been added to the forest, so it is
-        // safe to call unwrap.
+        /// If the fragment root has a duplicate label, the forest
+        /// will not grow, so we use the label to find the adjoined
+        /// node index.
+        ///
+        /// The nodes hava already been added to the forest, so it is
+        /// safe to call unwrap.
         macro_rules! conversion (
             ($node:expr) => {
                 {
@@ -360,20 +369,37 @@ impl<T: GraphLabel> Forest<T> for DefaultForest<T> {
         Ok(())
     }
 
-    fn clone_node<F>(&mut self, node_id: usize, clone_transform: F) -> Result<usize, Self::Error>
-    where
-        F: Fn(T) -> T,
-    {
-        let mut builder = PLGBuilderMut::from_graph_mut(&mut self.graph);
+    fn clone_node(&mut self, node_id: usize) -> Result<usize, Self::Error> {
+        let builder = PLGBuilderMut::from_graph_mut(&mut self.graph);
 
         let old_label = builder
             .vertex_label(node_id)?
             .ok_or(Error::NodeNoLabel(node_id))?;
 
-        let new_label = clone_transform(old_label);
+        if old_label.is_packed() {
+            return Err(ForestLabelError::ClonePack.into());
+        }
+
+        if old_label.is_cloned().is_some() {
+            let mut parents = self.parents_of(node_id)?;
+
+            // Make sure it only has one parent, which is the
+            // representative node.
+            assert_eq!(parents.len(), 1);
+
+            // We checked its length is 1, so it is safe to unwrap
+            // here.
+            return Ok(parents.next().unwrap().node());
+        }
+
+        let rep_label = old_label.pack()?;
+
+        let new_label = old_label.clone(self)?;
+
+        let mut builder = PLGBuilderMut::from_graph_mut(&mut self.graph);
 
         // Make a new node
-        let new_index = builder.add_vertex(new_label);
+        let new_index = builder.add_vertex(rep_label);
 
         // Re-direct parents to the new node.
         //
@@ -393,6 +419,10 @@ impl<T: GraphLabel> Forest<T> for DefaultForest<T> {
 
         builder.add_edge(new_index, node_id, old_label)?;
 
+        // Modify the label of the old node.
+
+        builder.set_label(node_id, new_label)?;
+
         Ok(new_index)
     }
 }
@@ -403,10 +433,10 @@ mod forest_test {
 
     macro_rules! leaf (
         ($label:expr, $type:tt) =>{
-            DefaultForest::<$type>::new_leaf($label)
+            DefaultForest::<ForestLabel<$type>>::new_leaf($label)
         };
         ($label:expr) => {
-            DefaultForest::<usize>::new_leaf($label)
+            DefaultForest::<ForestLabel<usize>>::new_leaf($label)
         }
     );
 
@@ -479,11 +509,9 @@ mod forest_test {
         // add a duplicate label
         forest.plant(3, leaf!(5), false)?;
 
-        let len = forest.nodes_len();
-
-        let clone_transform = |label| label + len;
+        let _len = forest.nodes_len();
 
-        forest.clone_node(5, clone_transform)?;
+        forest.clone_node(5)?;
 
         assert_eq!(forest.nodes_len(), 7);