Fix various bugs in matrix graph (petgraph#427)

* Fix various bugs in matrix graph

- Extension function for directed graphs did not copy the edges
- Off by one in with_capacity
- Check if the lazy adjacency matrix contains node before accesing

* Exact allocation in MatrixGraph::with_capacity

Avoids unnecessary quadratic memory waste.
Dynamic expansions continue rounding to the next power of two.

* Avoid unnecesary bonds checks in update_edge

Speeds up edge inserting benchmarks by ~35%

* Use array swaps when possible when extending

Overlapping slices still use the iterative method.

* Fix bug in iterative copy, add test

src/matrix_graph.rs

@@ -235,13 +235,23 @@ impl<N, E, Ty: EdgeType, Null: Nullable<Wrapped = E>, Ix: IndexType>
         };
 
         debug_assert!(node_capacity <= <Ix as IndexType>::max().index());
-        m.extend_capacity_for_node(NodeIndex::new(node_capacity));
+        if node_capacity > 0 {
+            m.extend_capacity_for_node(NodeIndex::new(node_capacity - 1), true);
+        }
 
         m
     }
 
     #[inline]
-    fn to_edge_position(&self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) -> usize {
+    fn to_edge_position(&self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) -> Option<usize> {
+        if cmp::max(a.index(), b.index()) >= self.node_capacity {
+            return None;
+        }
+        Some(self.to_edge_position_unchecked(a, b))
+    }
+
+    #[inline]
+    fn to_edge_position_unchecked(&self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) -> usize {
         to_linearized_matrix_position::<Ty>(a.index(), b.index(), self.node_capacity)
     }
 
@@ -295,31 +305,36 @@ impl<N, E, Ty: EdgeType, Null: Nullable<Wrapped = E>, Ix: IndexType>
     pub fn remove_node(&mut self, a: NodeIndex<Ix>) -> N {
         for id in self.nodes.iter_ids() {
             let position = self.to_edge_position(a, NodeIndex::new(id));
-            self.node_adjacencies[position] = Default::default();
+            if let Some(pos) = position {
+                self.node_adjacencies[pos] = Default::default();
+            }
 
             if Ty::is_directed() {
                 let position = self.to_edge_position(NodeIndex::new(id), a);
-                self.node_adjacencies[position] = Default::default();
+                if let Some(pos) = position {
+                    self.node_adjacencies[pos] = Default::default();
+                }
             }
         }
 
         self.nodes.remove(a.index())
     }
 
     #[inline]
-    fn extend_capacity_for_node(&mut self, min_node: NodeIndex<Ix>) {
+    fn extend_capacity_for_node(&mut self, min_node: NodeIndex<Ix>, exact: bool) {
         self.node_capacity = extend_linearized_matrix::<Ty, _>(
             &mut self.node_adjacencies,
             self.node_capacity,
-            min_node.index(),
+            min_node.index() + 1,
+            exact,
         );
     }
 
     #[inline]
     fn extend_capacity_for_edge(&mut self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) {
         let min_node = cmp::max(a, b);
         if min_node.index() >= self.node_capacity {
-            self.extend_capacity_for_node(min_node);
+            self.extend_capacity_for_node(min_node, false);
         }
     }
 
@@ -333,8 +348,7 @@ impl<N, E, Ty: EdgeType, Null: Nullable<Wrapped = E>, Ix: IndexType>
     /// **Panics** if any of the nodes don't exist.
     pub fn update_edge(&mut self, a: NodeIndex<Ix>, b: NodeIndex<Ix>, weight: E) -> Option<E> {
         self.extend_capacity_for_edge(a, b);
-
-        let p = self.to_edge_position(a, b);
+        let p = self.to_edge_position_unchecked(a, b);
         let old_weight = mem::replace(&mut self.node_adjacencies[p], Null::new(weight));
         if old_weight.is_null() {
             self.nb_edges += 1;
@@ -365,7 +379,9 @@ impl<N, E, Ty: EdgeType, Null: Nullable<Wrapped = E>, Ix: IndexType>
     /// **Panics** if any of the nodes don't exist.
     /// **Panics** if no edge exists between `a` and `b`.
     pub fn remove_edge(&mut self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) -> E {
-        let p = self.to_edge_position(a, b);
+        let p = self
+            .to_edge_position(a, b)
+            .expect("No edge found between the nodes.");
         let old_weight = mem::replace(&mut self.node_adjacencies[p], Default::default())
             .into()
             .unwrap();
@@ -378,8 +394,10 @@ impl<N, E, Ty: EdgeType, Null: Nullable<Wrapped = E>, Ix: IndexType>
     ///
     /// **Panics** if any of the nodes don't exist.
     pub fn has_edge(&self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) -> bool {
-        let p = self.to_edge_position(a, b);
-        !self.node_adjacencies[p].is_null()
+        if let Some(p) = self.to_edge_position(a, b) {
+            return !self.node_adjacencies[p].is_null();
+        }
+        false
     }
 
     /// Access the weight for node `a`.
@@ -406,8 +424,12 @@ impl<N, E, Ty: EdgeType, Null: Nullable<Wrapped = E>, Ix: IndexType>
     ///
     /// **Panics** if no edge exists between `a` and `b`.
     pub fn edge_weight(&self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) -> &E {
-        let p = self.to_edge_position(a, b);
-        self.node_adjacencies[p].as_ref().unwrap()
+        let p = self
+            .to_edge_position(a, b)
+            .expect("No edge found between the nodes.");
+        self.node_adjacencies[p]
+            .as_ref()
+            .expect("No edge found between the nodes.")
     }
 
     /// Access the weight for edge `e`, mutably.
@@ -416,8 +438,12 @@ impl<N, E, Ty: EdgeType, Null: Nullable<Wrapped = E>, Ix: IndexType>
     ///
     /// **Panics** if no edge exists between `a` and `b`.
     pub fn edge_weight_mut(&mut self, a: NodeIndex<Ix>, b: NodeIndex<Ix>) -> &mut E {
-        let p = self.to_edge_position(a, b);
-        self.node_adjacencies[p].as_mut().unwrap()
+        let p = self
+            .to_edge_position(a, b)
+            .expect("No edge found between the nodes.");
+        self.node_adjacencies[p]
+            .as_mut()
+            .expect("No edge found between the nodes.")
     }
 
     /// Return an iterator of all nodes with an edge starting from `a`.
@@ -789,12 +815,16 @@ fn to_linearized_matrix_position<Ty: EdgeType>(row: usize, column: usize, width:
 fn extend_linearized_matrix<Ty: EdgeType, T: Default>(
     node_adjacencies: &mut Vec<T>,
     old_node_capacity: usize,
-    min_node_capacity: usize,
+    new_capacity: usize,
+    exact: bool,
 ) -> usize {
+    if old_node_capacity >= new_capacity {
+        return old_node_capacity;
+    }
     if Ty::is_directed() {
-        extend_flat_square_matrix(node_adjacencies, old_node_capacity, min_node_capacity)
+        extend_flat_square_matrix(node_adjacencies, old_node_capacity, new_capacity, exact)
     } else {
-        extend_lower_triangular_matrix(node_adjacencies, min_node_capacity)
+        extend_lower_triangular_matrix(node_adjacencies, new_capacity)
     }
 }
 
@@ -807,30 +837,42 @@ fn to_flat_square_matrix_position(row: usize, column: usize, width: usize) -> us
 fn extend_flat_square_matrix<T: Default>(
     node_adjacencies: &mut Vec<T>,
     old_node_capacity: usize,
-    min_node_capacity: usize,
+    new_node_capacity: usize,
+    exact: bool,
 ) -> usize {
-    let min_node_capacity = (min_node_capacity + 1).next_power_of_two();
+    // Grow the capacity by exponential steps to avoid repeated allocations.
+    // Disabled for the with_capacity constructor.
+    let new_node_capacity = if exact {
+        new_node_capacity
+    } else {
+        const MIN_CAPACITY: usize = 4;
+        cmp::max(new_node_capacity.next_power_of_two(), MIN_CAPACITY)
+    };
 
     // Optimization: when resizing the matrix this way we skip the first few grows to make
     // small matrices a bit faster to work with.
-    const MIN_CAPACITY: usize = 4;
-    let new_node_capacity = cmp::max(min_node_capacity, MIN_CAPACITY);
-
-    let mut new_node_adjacencies = vec![];
-    ensure_len(&mut new_node_adjacencies, new_node_capacity.pow(2));
 
-    for c in 0..old_node_capacity {
+    ensure_len(node_adjacencies, new_node_capacity.pow(2));
+    for c in (1..old_node_capacity).rev() {
         let pos = c * old_node_capacity;
         let new_pos = c * new_node_capacity;
 
-        let mut old = &mut node_adjacencies[pos..pos + old_node_capacity];
-        let mut new = &mut new_node_adjacencies[new_pos..new_pos + old_node_capacity];
+        // Move the slices directly if they do not overlap with their new position
+        if pos + old_node_capacity <= new_pos {
+            let old = node_adjacencies[pos..pos + old_node_capacity].as_mut_ptr();
+            let new = node_adjacencies[new_pos..new_pos + old_node_capacity].as_mut_ptr();
 
-        mem::swap(&mut old, &mut new);
+            // SAFE: new starts at least `old_node_capacity` positions after old.
+            unsafe {
+                std::ptr::swap_nonoverlapping(old, new, old_node_capacity);
+            }
+        } else {
+            for i in (0..old_node_capacity).rev() {
+                node_adjacencies.as_mut_slice().swap(pos + i, new_pos + i);
+            }
+        }
     }
 
-    mem::swap(node_adjacencies, &mut new_node_adjacencies);
-
     new_node_capacity
 }
 
@@ -847,11 +889,12 @@ fn to_lower_triangular_matrix_position(row: usize, column: usize) -> usize {
 #[inline]
 fn extend_lower_triangular_matrix<T: Default>(
     node_adjacencies: &mut Vec<T>,
-    new_node_capacity: usize,
+    new_capacity: usize,
 ) -> usize {
-    let max_pos = to_lower_triangular_matrix_position(new_node_capacity, new_node_capacity);
+    let max_node = new_capacity - 1;
+    let max_pos = to_lower_triangular_matrix_position(max_node, max_node);
     ensure_len(node_adjacencies, max_pos + 1);
-    new_node_capacity + 1
+    new_capacity
 }
 
 /// Grow a Vec by appending the type's default value until the `size` is reached.
@@ -1284,6 +1327,45 @@ mod tests {
         assert_eq!(g.edge_count(), 2);
     }
 
+    #[test]
+    /// Adds an edge that triggers a second extension of the matrix.
+    /// From #425
+    fn test_add_edge_with_extension() {
+        let mut g = DiMatrix::<u8, ()>::new();
+        let _n0 = g.add_node(0);
+        let n1 = g.add_node(1);
+        let n2 = g.add_node(2);
+        let n3 = g.add_node(3);
+        let n4 = g.add_node(4);
+        let _n5 = g.add_node(5);
+        g.add_edge(n2, n1, ());
+        g.add_edge(n2, n3, ());
+        g.add_edge(n2, n4, ());
+        assert_eq!(g.node_count(), 6);
+        assert_eq!(g.edge_count(), 3);
+        assert!(g.has_edge(n2, n1));
+        assert!(g.has_edge(n2, n3));
+        assert!(g.has_edge(n2, n4));
+    }
+
+    #[test]
+    fn test_matrix_resize() {
+        let mut g = DiMatrix::<u8, ()>::with_capacity(3);
+        let n0 = g.add_node(0);
+        let n1 = g.add_node(1);
+        let n2 = g.add_node(2);
+        let n3 = g.add_node(3);
+        g.add_edge(n1, n0, ());
+        g.add_edge(n1, n1, ());
+        // Triggers a resize from capacity 3 to 4
+        g.add_edge(n2, n3, ());
+        assert_eq!(g.node_count(), 4);
+        assert_eq!(g.edge_count(), 3);
+        assert!(g.has_edge(n1, n0));
+        assert!(g.has_edge(n1, n1));
+        assert!(g.has_edge(n2, n3));
+    }
+
     #[test]
     fn test_add_edge_with_weights() {
         let mut g = MatrixGraph::new();