1 files changed, 104 insertions, 31 deletions

diff --git a/graph/src/labelled.rs b/graph/src/labelled.rs
index d0a02d0..6341787 100644
--- a/graph/src/labelled.rs
+++ b/graph/src/labelled.rs
@@ -142,10 +142,11 @@ impl<T: GraphLabel> Graph for DLGraph<T> {
         match self.nodes.get(source) {
             Some(source_node) => {
                 if self.nodes.get(target).is_none() {
-                    return Err(Error::IndexOutOfBounds(target, self.nodes.len()));
+                    Err(Error::IndexOutOfBounds(target, self.nodes.len()))
+                } else {
+                    Ok(source_node.by_target.contains_key(&target)
+                        && !source_node.by_target.get(&target).unwrap().is_empty())
                 }
-
-                Ok(source_node.by_target.contains_key(&target))
             }
             None => Err(Error::IndexOutOfBounds(source, self.nodes.len())),
         }
@@ -198,7 +199,7 @@ impl<T: GraphLabel> Graph for DLGraph<T> {
 /// A delegation of iterators.
 ///
 /// This is used to avoid a boxed pointer to an iterator.
-#[derive(Default, Debug)]
+#[derive(Default, Debug, Clone)]
 pub struct LabelIndexIter<'a> {
     iter: Option<std::iter::Copied<Iter<'a, usize>>>,
 }
@@ -279,25 +280,81 @@ impl<'a, T> Iterator for LabelIter<'a, T> {
     }
 }
 
+/// This is used to avoid a boxed pointer to an iterator.
+#[derive(Debug)]
+pub struct EdgeLabelIter<'a, T> {
+    iter: Option<Iter<'a, T>>,
+}
+
+impl<'a, T> Default for EdgeLabelIter<'a, T> {
+    #[inline]
+    fn default() -> Self {
+        Self { iter: None }
+    }
+}
+
+impl<'a, T: Copy + Clone> ExactSizeIterator for EdgeLabelIter<'a, T> {
+    #[inline]
+    fn len(&self) -> usize {
+        if let Some(iter) = &self.iter {
+            iter.len()
+        } else {
+            0
+        }
+    }
+}
+
+impl<'a, T: Copy + Clone> EdgeLabelIter<'a, T> {
+    fn new(iter: Iter<'a, T>) -> Self {
+        Self { iter: Some(iter) }
+    }
+}
+
+impl<'a, T: Copy + Clone> Iterator for EdgeLabelIter<'a, T> {
+    type Item = T;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(iter) = &mut self.iter {
+            iter.next().copied()
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if let Some(iter) = &self.iter {
+            iter.size_hint()
+        } else {
+            (0, Some(0))
+        }
+    }
+}
+
 impl<T: GraphLabel> LabelGraph<T> for DLGraph<T> {
     type Iter<'a> = LabelIndexIter<'a> where T: 'a;
 
     type LabelIter<'a> = LabelIter<'a,T> where T: 'a;
 
-    fn edge_label(&self, source: usize, target: usize) -> Result<Vec<T>, Error> {
+    type EdgeLabelIter<'a> = EdgeLabelIter<'a,T>
+    where
+        Self: 'a,
+        T: 'a + Copy + Clone;
+
+    fn edge_label(&self, source: usize, target: usize) -> Result<Self::EdgeLabelIter<'_>, Error> {
         if self.has_edge(source, target)? {
-            Ok(self
-                .nodes
-                .get(source)
-                .unwrap()
-                .by_target
-                .get(&target)
-                .unwrap()
-                .iter()
-                .copied()
-                .collect())
+            Ok(EdgeLabelIter::new(
+                self.nodes
+                    .get(source)
+                    .unwrap()
+                    .by_target
+                    .get(&target)
+                    .unwrap()
+                    .iter(),
+            ))
         } else {
-            Ok(Vec::new())
+            Ok(Default::default())
         }
     }
 
@@ -335,11 +392,11 @@ impl<T: GraphLabel> LabelGraph<T> for DLGraph<T> {
                     return Err(Error::IndexOutOfBounds(target, nodes_len));
                 }
 
-                if let Some(labels) = node.by_target.get(&target) {
-                    Ok(labels.contains(label))
-                } else {
-                    Ok(false)
-                }
+                Ok(node
+                    .by_target
+                    .get(&target)
+                    .map(|labels| labels.contains(label))
+                    .unwrap_or(false))
             }
             None => Err(Error::IndexOutOfBounds(node_id, nodes_len)),
         }
@@ -443,6 +500,12 @@ impl<T: GraphLabel> Builder for DLGBuilder<T> {
         self.nodes.len() - 1
     }
 
+    #[inline]
+    fn add_vertices(&mut self, n: usize) {
+        self.nodes
+            .extend(std::iter::repeat_with(Default::default).take(n));
+    }
+
     fn add_edge(&mut self, source: usize, target: usize, label: Self::Label) -> Result<(), Error> {
         let nodes_len = self.nodes.len();
 
@@ -541,14 +604,29 @@ impl<T: GraphLabel> Builder for DLGBuilder<T> {
                 // If after removal no labels remain for the target,
                 // we remove the edge entirely, to avoid false empty
                 // edges.
-                if source_node.flat.is_empty() {
-                    source_node.by_target.remove(&target);
 
-                    for label in to_remove.iter() {
-                        source_node.by_label.remove(label);
+                if let Some(set) = source_node.by_target.get(&target) {
+                    if set.is_empty() {
+                        source_node.by_target.remove(&target);
                     }
                 }
 
+                for label in to_remove.iter() {
+                    if let Some(set) = source_node.by_label.get(label) {
+                        if set.is_empty() {
+                            source_node.by_label.remove(label);
+                        }
+                    }
+                }
+
+                // if source_node.flat.is_empty() {
+                //     source_node.by_target.remove(&target);
+
+                //     for label in to_remove.iter() {
+                //         source_node.by_label.remove(label);
+                //     }
+                // }
+
                 let new_children_set: Set<(T, usize)> = source_node.flat.iter().copied().collect();
 
                 // When source_node is in use, we cannot borrow self
@@ -639,10 +717,7 @@ mod label_test {
         );
 
         // testing edge_label
-        assert_eq!(
-            graph.edge_label(5, 2)?.into_iter().collect::<Set<_>>(),
-            set!(3)
-        );
+        assert_eq!(graph.edge_label(5, 2)?.collect::<Set<_>>(), set!(3));
         assert!(matches!(
             graph.edge_label(6, 2),
             Err(Error::IndexOutOfBounds(6, 6))
@@ -683,8 +758,6 @@ mod label_test {
     }
 }
 
-// TODO: Test print_viz
-
 #[cfg(test)]
 mod test_dlgraph_builder {
     use super::*;