1 files changed, 157 insertions, 14 deletions

diff --git a/graph/src/labelled/binary.rs b/graph/src/labelled/binary.rs
index 67d86f9..bfd8109 100644
--- a/graph/src/labelled/binary.rs
+++ b/graph/src/labelled/binary.rs
@@ -2,7 +2,7 @@
 //! labels and each node knows about its parents.
 
 use super::*;
-use crate::{Parent, ParentsGraph};
+use crate::{Parent, ParentsGraph, RedirectGraph};
 
 use std::{
     collections::{hash_map::Iter as MapIter, HashMap as Map},
@@ -142,6 +142,11 @@ impl<T: GraphLabel> Graph for PLGraph<T> {
     }
 
     #[inline]
+    fn edges_len(&self) -> Option<usize> {
+        Some(self.nodes.iter().map(|node| node.children.len()).sum())
+    }
+
+    #[inline]
     fn children_of(&self, node_id: usize) -> Result<Self::Iter<'_>, Error> {
         if let Some(node) = self.nodes.get(node_id) {
             Ok(node.children.iter().copied())
@@ -243,6 +248,76 @@ impl<T: GraphLabel> ParentsGraph for PLGraph<T> {
     }
 }
 
+impl<T: GraphLabel> RedirectGraph for PLGraph<T> {
+    fn redirect(
+        &mut self,
+        node_id: usize,
+        child_index: usize,
+        new_child: usize,
+    ) -> Result<(), Error> {
+        let nodes_len = self.nodes.len();
+
+        if !self.has_node(new_child) {
+            return Err(Error::IndexOutOfBounds(new_child, nodes_len));
+        }
+
+        if let Some(node) = self.nodes.get_mut(node_id) {
+            if node.children.len() <= child_index {
+                return Err(Error::IndexOutOfBounds(child_index, node.children.len()));
+            }
+
+            // Check if `new_child` is already pointed to by the node.
+            if let Some(index) = node.children.indices.get(&new_child) {
+                // This should not happen in our use case, but we
+                // should still do somthing: as the edges cannot
+                // duplicate, we simply remove the original edge,
+                // unless index = child_index, in which case the old
+                // child is equal to the new child, and we have
+                // nothing to do.
+                if *index != child_index {
+                    node.children.remove(new_child);
+                }
+            } else {
+                // The index has been checked above, so it is safe to
+                // call `unwrap` here.
+                let old_child = std::mem::replace(
+                    node.children.children.get_mut(child_index).unwrap(),
+                    new_child,
+                );
+
+                node.children.indices.remove(&old_child);
+
+                node.children.indices.insert(new_child, child_index);
+
+                // Don't forget to remove `node` from the parents of
+                // the old child.
+
+                if let Some(old_child_node) = self.nodes.get_mut(old_child) {
+                    old_child_node.parents.remove(&node_id);
+                } else {
+                    // The node contained an invalid child.
+                    return Err(Error::IndexOutOfBounds(old_child, nodes_len));
+                }
+
+                // Don't forget to add node as a parent to the new
+                // child.
+
+                // new_child has been checked at the beginning of the
+                // function, so it is safe to call `unwrap`.
+                self.nodes
+                    .get_mut(new_child)
+                    .unwrap()
+                    .parents
+                    .insert(node_id, child_index);
+            }
+
+            Ok(())
+        } else {
+            Err(Error::IndexOutOfBounds(node_id, nodes_len))
+        }
+    }
+}
+
 impl<T: GraphLabel> LabelGraph<T> for PLGraph<T> {
     type Iter<'a> = std::iter::Empty<usize>
     where
@@ -299,6 +374,20 @@ pub struct PLGBuilderMut<'a, T: GraphLabel> {
     graph: &'a mut PLGraph<T>,
 }
 
+impl<'a, T: GraphLabel> std::ops::Deref for PLGBuilderMut<'a, T> {
+    type Target = PLGraph<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.graph
+    }
+}
+
+impl<'a, T: GraphLabel> std::ops::DerefMut for PLGBuilderMut<'a, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.graph
+    }
+}
+
 impl<'a, T: GraphLabel> BuilderMut for PLGBuilderMut<'a, T> {
     type Label = T;
 
@@ -467,6 +556,8 @@ mod binary_test {
         builder.add_vertex(4);
         builder.add_vertex(5);
 
+        // These labels are not used on edges: they are just place
+        // holders.
         builder.add_edge(1, 0, 0)?;
         builder.add_edge(2, 0, 0)?;
         builder.add_edge(2, 1, 0)?;
@@ -481,43 +572,73 @@ mod binary_test {
         // testing adding a duplicatedly labelled node
         assert_eq!(builder.add_vertex(0), 0);
 
-        let graph = graph;
-
         // ensuring the correct length
-        assert_eq!(graph.nodes_len(), 6);
+        assert_eq!(builder.nodes_len(), 6);
 
         // testing children_of
-        assert_eq!(graph.children_of(5)?.collect::<Set<_>>(), set!(1, 3, 2));
+        assert_eq!(builder.children_of(5)?.collect::<Set<_>>(), set!(1, 3, 2));
 
         // testing parents_of
 
         assert_eq!(
-            graph.parents_of(0)?.collect::<Set<_>>(),
+            builder.parents_of(0)?.collect::<Set<_>>(),
             set!(Parent::new(1, 0), Parent::new(2, 0), Parent::new(3, 0))
         );
 
         assert_eq!(
-            graph.parents_of(1)?.collect::<Set<_>>(),
+            builder.parents_of(1)?.collect::<Set<_>>(),
             set!(Parent::new(2, 1), Parent::new(4, 0), Parent::new(5, 2))
         );
 
-        assert_eq!(graph.parents_of(5)?.len(), 0);
+        assert_eq!(builder.parents_of(5)?.len(), 0);
 
         // testing degree
-        assert_eq!(graph.degree(4)?, 2);
+        assert_eq!(builder.degree(4)?, 2);
 
         // testing is_empty_node
-        assert!(graph.is_empty_node(0)?);
-        assert!(!graph.is_empty_node(1)?);
+        assert!(builder.is_empty_node(0)?);
+        assert!(!builder.is_empty_node(1)?);
 
         // testing has_edge
-        assert!(graph.has_edge(3, 2)?);
-        assert!(!graph.has_edge(3, 1)?);
+        assert!(builder.has_edge(3, 2)?);
+        assert!(!builder.has_edge(3, 1)?);
         assert!(matches!(
-            graph.has_edge(3, 6),
+            builder.has_edge(3, 6),
             Err(Error::IndexOutOfBounds(6, 6))
         ));
 
+        // testing redirect
+        builder.redirect(5, 2, 0)?;
+        assert_eq!(builder.children_of(5)?.collect::<Set<_>>(), set!(0, 3, 2));
+
+        assert_eq!(
+            builder.parents_of(0)?.collect::<Set<_>>(),
+            set!(
+                Parent::new(1, 0),
+                Parent::new(2, 0),
+                Parent::new(3, 0),
+                Parent::new(5, 2)
+            )
+        );
+
+        builder.redirect(5, 0, 1)?;
+
+        assert_eq!(builder.children_of(5)?.collect::<Set<_>>(), set!(1, 0, 3));
+
+        assert_eq!(
+            builder.parents_of(1)?.collect::<Set<_>>(),
+            set!(Parent::new(2, 1), Parent::new(4, 0), Parent::new(5, 0))
+        );
+
+        builder.redirect(5, 0, 1)?; // should be no-op
+
+        assert_eq!(builder.children_of(5)?.collect::<Set<_>>(), set!(1, 0, 3));
+
+        assert_eq!(
+            builder.parents_of(1)?.collect::<Set<_>>(),
+            set!(Parent::new(2, 1), Parent::new(4, 0), Parent::new(5, 0))
+        );
+
         Ok(())
     }
 }
@@ -636,6 +757,28 @@ mod test_plgraph_builder {
 
         // println!();
 
+        // source out of bounds
+        assert!(matches!(
+            builder.redirect(5, 0, 0),
+            Err(Error::IndexOutOfBounds(5, 5))
+        ));
+
+        // child_index out of bounds
+        assert!(matches!(
+            builder.redirect(4, 0, 0),
+            Err(Error::IndexOutOfBounds(0, 0))
+        ));
+
+        // new_child out of bounds
+        assert!(matches!(
+            builder.redirect(4, 0, 10),
+            Err(Error::IndexOutOfBounds(10, 5))
+        ));
+
+        // println!("Correct errors when redirecting");
+
+        // println!();
+
         let graph = graph;
 
         println!("final graph: {graph:?}");