Add checked UnionFind methods (petgraph#730)

PR provides the following `UnionFind` methods, that are free from
panics:

```rust
pub fn try_find(&self, mut x: K) -> Option<K>

pub fn try_find_mut(&mut self, x: K) -> Option<K>

pub fn try_equiv(&self, x: K, y: K) -> Result<bool, K>

pub fn try_union(&mut self, x: K, y: K) -> Result<bool, K>
```

Some of old methods were rewritten using the `try` twins with the
addition of `.unwrap` or `.expect`, for example:
```rust
pub fn find(&self, x: K) -> K {
    self.try_find(x).expect("The index is out of bounds")
}
```
If you consider these changes to be critical & breaking, I will return
back the old `assert`ions.

I'm open to discussion!

Author: starovoid

src/unionfind.rs

@@ -95,19 +95,25 @@ where
     ///
     /// **Panics** if `x` is out of bounds.
     pub fn find(&self, x: K) -> K {
-        assert!(x.index() < self.len());
-        unsafe {
-            let mut x = x;
-            loop {
-                // Use unchecked indexing because we can trust the internal set ids.
-                let xparent = *get_unchecked(&self.parent, x.index());
-                if xparent == x {
-                    break;
-                }
-                x = xparent;
+        self.try_find(x).expect("The index is out of bounds")
+    }
+
+    /// Return the representative for `x` or `None` if `x` is out of bounds.
+    pub fn try_find(&self, mut x: K) -> Option<K> {
+        if x.index() >= self.len() {
+            return None;
+        }
+
+        loop {
+            // Use unchecked indexing because we can trust the internal set ids.
+            let xparent = unsafe { *get_unchecked(&self.parent, x.index()) };
+            if xparent == x {
+                break;
             }
-            x
+            x = xparent;
         }
+
+        Some(x)
     }
 
     /// Return the representative for `x`.
@@ -121,6 +127,17 @@ where
         unsafe { self.find_mut_recursive(x) }
     }
 
+    /// Return the representative for `x` or `None` if `x` is out of bounds.
+    ///
+    /// Write back the found representative, flattening the internal
+    /// datastructure in the process and quicken future lookups.
+    pub fn try_find_mut(&mut self, x: K) -> Option<K> {
+        if x.index() >= self.len() {
+            return None;
+        }
+        Some(unsafe { self.find_mut_recursive(x) })
+    }
+
     unsafe fn find_mut_recursive(&mut self, mut x: K) -> K {
         let mut parent = *get_unchecked(&self.parent, x.index());
         while parent != x {
@@ -134,24 +151,46 @@ where
 
     /// Returns `true` if the given elements belong to the same set, and returns
     /// `false` otherwise.
+    ///
+    /// **Panics** if `x` or `y` is out of bounds.
     pub fn equiv(&self, x: K, y: K) -> bool {
         self.find(x) == self.find(y)
     }
 
+    /// Returns `Ok(true)` if the given elements belong to the same set, and returns
+    /// `Ok(false)` otherwise.
+    ///
+    /// If `x` or `y` are out of bounds, it returns `Err` with the first bad index found.
+    pub fn try_equiv(&self, x: K, y: K) -> Result<bool, K> {
+        let xrep = self.try_find(x).ok_or(x)?;
+        let yrep = self.try_find(y).ok_or(y)?;
+        Ok(xrep == yrep)
+    }
+
     /// Unify the two sets containing `x` and `y`.
     ///
     /// Return `false` if the sets were already the same, `true` if they were unified.
     ///
     /// **Panics** if `x` or `y` is out of bounds.
     pub fn union(&mut self, x: K, y: K) -> bool {
+        self.try_union(x, y).unwrap()
+    }
+
+    /// Unify the two sets containing `x` and `y`.
+    ///
+    /// Return `Ok(false)` if the sets were already the same, `Ok(true)` if they were unified.
+    ///
+    /// If `x` or `y` are out of bounds, it returns `Err` with first found bad index.
+    /// But if `x == y`, the result will be `Ok(false)` even if the indexes go out of bounds.
+    pub fn try_union(&mut self, x: K, y: K) -> Result<bool, K> {
         if x == y {
-            return false;
+            return Ok(false);
         }
-        let xrep = self.find_mut(x);
-        let yrep = self.find_mut(y);
+        let xrep = self.try_find_mut(x).ok_or(x)?;
+        let yrep = self.try_find_mut(y).ok_or(y)?;
 
         if xrep == yrep {
-            return false;
+            return Ok(false);
         }
 
         let xrepu = xrep.index();
@@ -169,7 +208,7 @@ where
                 self.rank[xrepu] += 1;
             }
         }
-        true
+        Ok(true)
     }
 
     /// Return a vector mapping each element to its representative.

tests/unionfind.rs

@@ -33,6 +33,37 @@ fn uf_test() {
     assert_eq!(set.len(), 3);
 }
 
+#[test]
+fn uf_test_checked() {
+    let n = 8;
+    let mut u = UnionFind::new(n);
+    for i in 0..n {
+        assert_eq!(u.try_find(i), Some(i));
+        assert_eq!(u.try_find_mut(i), Some(i));
+        assert_eq!(u.try_union(i, i), Ok(false));
+    }
+
+    assert!(u.try_union(0, 1).is_ok());
+    assert_eq!(u.try_find(0), u.try_find(1));
+    assert!(u.try_find(0).is_some());
+    assert!(u.try_union(1, 3).is_ok());
+    assert!(u.try_union(1, 4).is_ok());
+    assert!(u.try_union(4, 7).is_ok());
+    assert_eq!(u.try_find(0), u.try_find(3));
+    assert_eq!(u.try_find(1), u.try_find(3));
+    assert!(u.try_find(0).is_some());
+    assert!(u.try_find(1).is_some());
+    assert!(u.try_find(0) != u.try_find(2));
+    assert_eq!(u.try_find(7), u.try_find(0));
+    assert!(u.try_union(5, 6).is_ok());
+    assert_eq!(u.try_find(6), u.try_find(5));
+    assert!(u.try_find(6) != u.try_find(7));
+
+    // check that there are now 3 disjoint sets
+    let set = (0..n).map(|i| u.find(i)).collect::<HashSet<_>>();
+    assert_eq!(set.len(), 3);
+}
+
 #[test]
 fn uf_test_with_equiv() {
     let n = 8;
@@ -61,6 +92,34 @@ fn uf_test_with_equiv() {
     assert_eq!(set.len(), 3);
 }
 
+#[test]
+fn uf_test_with_checked_equiv() {
+    let n = 8;
+    let mut u = UnionFind::new(n);
+    for i in 0..n {
+        assert_eq!(u.find(i), i);
+        assert_eq!(u.find_mut(i), i);
+        assert_eq!(u.try_equiv(i, i), Ok(true));
+    }
+
+    u.union(0, 1);
+    assert_eq!(u.try_equiv(0, 1), Ok(true));
+    u.union(1, 3);
+    u.union(1, 4);
+    u.union(4, 7);
+    assert_eq!(u.try_equiv(0, 7), Ok(true));
+    assert_eq!(u.try_equiv(1, 3), Ok(true));
+    assert_eq!(u.try_equiv(0, 2), Ok(false));
+    assert_eq!(u.try_equiv(7, 0), Ok(true));
+    u.union(5, 6);
+    assert_eq!(u.try_equiv(6, 5), Ok(true));
+    assert_eq!(u.try_equiv(6, 7), Ok(false));
+
+    // check that there are now 3 disjoint sets
+    let set = (0..n).map(|i| u.find(i)).collect::<HashSet<_>>();
+    assert_eq!(set.len(), 3);
+}
+
 #[test]
 fn uf_rand() {
     let n = 1 << 14;
@@ -89,6 +148,20 @@ fn uf_u8() {
     }
 }
 
+#[test]
+fn uf_u8_checked() {
+    let n = 256;
+    let mut rng = ChaChaRng::from_rng(thread_rng()).unwrap();
+    let mut u = UnionFind::<u8>::new(n);
+    for _ in 0..(n * 8) {
+        let a = rng.gen();
+        let b = rng.gen();
+        let ar = u.try_find(a).unwrap();
+        let br = u.try_find(b).unwrap();
+        assert_eq!(ar != br, u.try_union(a, b).unwrap());
+    }
+}
+
 #[test]
 fn labeling() {
     let mut u = UnionFind::<u32>::new(48);
@@ -136,3 +209,27 @@ fn uf_incremental() {
         .collect::<HashSet<_>>();
     assert_eq!(set.len(), 3);
 }
+
+#[test]
+fn uf_test_out_of_bounds() {
+    let n = 8;
+    let mut u = UnionFind::new(n);
+    for i in 0..n {
+        u.find(i);
+        u.find_mut(i);
+        u.union(i, i);
+    }
+
+    assert!(u.try_find(50).is_none());
+    assert!(u.try_find_mut(50).is_none());
+
+    assert_eq!(u.try_union(1, 50), Err(50));
+    assert_eq!(u.try_union(50, 1), Err(50));
+    assert_eq!(u.try_union(30, 50), Err(30));
+    assert_eq!(u.try_union(50, 30), Err(50));
+
+    assert_eq!(u.try_equiv(1, 50), Err(50));
+    assert_eq!(u.try_equiv(50, 1), Err(50));
+    assert_eq!(u.try_equiv(30, 50), Err(30));
+    assert_eq!(u.try_equiv(50, 30), Err(50));
+}