Implement BitOps for TreeSet

1 files changed, 125 insertions, 1 deletions

diff --git a/src/libcollections/tree/set.rs b/src/libcollections/tree/set.rs
index 6988b929df6..cf6a534cf19 100644
--- a/src/libcollections/tree/set.rs
+++ b/src/libcollections/tree/set.rs
@@ -21,7 +21,6 @@ use std::hash::{Writer, Hash};
 use tree_map::{TreeMap, Entries, RevEntries, MoveEntries};
 
 // FIXME(conventions): implement bounded iterators
-// FIXME(conventions): implement BitOr, BitAnd, BitXor, and Sub
 // FIXME(conventions): replace rev_iter(_mut) by making iter(_mut) DoubleEnded
 
 /// An implementation of the `Set` trait on top of the `TreeMap` container. The
@@ -666,6 +665,90 @@ impl<'a, T: Ord> Iterator<&'a T> for UnionItems<'a, T> {
     }
 }
 
+#[unstable = "matches collection reform specification, waiting for dust to settle"]
+impl<T: Ord + Clone> BitOr<TreeSet<T>, TreeSet<T>> for TreeSet<T> {
+    /// Returns the union of `self` and `rhs` as a new `TreeSet<T>`.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use std::collections::TreeSet;
+    ///
+    /// let a: TreeSet<int> = vec![1, 2, 3].into_iter().collect();
+    /// let b: TreeSet<int> = vec![3, 4, 5].into_iter().collect();
+    ///
+    /// let set: TreeSet<int> = a | b;
+    /// let v: Vec<int> = set.into_iter().collect();
+    /// assert_eq!(v, vec![1, 2, 3, 4, 5]);
+    /// ```
+    fn bitor(&self, rhs: &TreeSet<T>) -> TreeSet<T> {
+        self.union(rhs).cloned().collect()
+    }
+}
+
+#[unstable = "matches collection reform specification, waiting for dust to settle"]
+impl<T: Ord + Clone> BitAnd<TreeSet<T>, TreeSet<T>> for TreeSet<T> {
+    /// Returns the intersection of `self` and `rhs` as a new `TreeSet<T>`.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use std::collections::TreeSet;
+    ///
+    /// let a: TreeSet<int> = vec![1, 2, 3].into_iter().collect();
+    /// let b: TreeSet<int> = vec![2, 3, 4].into_iter().collect();
+    ///
+    /// let set: TreeSet<int> = a & b;
+    /// let v: Vec<int> = set.into_iter().collect();
+    /// assert_eq!(v, vec![2, 3]);
+    /// ```
+    fn bitand(&self, rhs: &TreeSet<T>) -> TreeSet<T> {
+        self.intersection(rhs).cloned().collect()
+    }
+}
+
+#[unstable = "matches collection reform specification, waiting for dust to settle"]
+impl<T: Ord + Clone> BitXor<TreeSet<T>, TreeSet<T>> for TreeSet<T> {
+    /// Returns the symmetric difference of `self` and `rhs` as a new `TreeSet<T>`.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use std::collections::TreeSet;
+    ///
+    /// let a: TreeSet<int> = vec![1, 2, 3].into_iter().collect();
+    /// let b: TreeSet<int> = vec![3, 4, 5].into_iter().collect();
+    ///
+    /// let set: TreeSet<int> = a ^ b;
+    /// let v: Vec<int> = set.into_iter().collect();
+    /// assert_eq!(v, vec![1, 2, 4, 5]);
+    /// ```
+    fn bitxor(&self, rhs: &TreeSet<T>) -> TreeSet<T> {
+        self.symmetric_difference(rhs).cloned().collect()
+    }
+}
+
+#[unstable = "matches collection reform specification, waiting for dust to settle"]
+impl<T: Ord + Clone> Sub<TreeSet<T>, TreeSet<T>> for TreeSet<T> {
+    /// Returns the difference of `self` and `rhs` as a new `TreeSet<T>`.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use std::collections::TreeSet;
+    ///
+    /// let a: TreeSet<int> = vec![1, 2, 3].into_iter().collect();
+    /// let b: TreeSet<int> = vec![3, 4, 5].into_iter().collect();
+    ///
+    /// let set: TreeSet<int> = a - b;
+    /// let v: Vec<int> = set.into_iter().collect();
+    /// assert_eq!(v, vec![1, 2]);
+    /// ```
+    fn sub(&self, rhs: &TreeSet<T>) -> TreeSet<T> {
+        self.difference(rhs).cloned().collect()
+    }
+}
+
 impl<T: Ord> FromIterator<T> for TreeSet<T> {
     fn from_iter<Iter: Iterator<T>>(iter: Iter) -> TreeSet<T> {
         let mut set = TreeSet::new();
@@ -695,6 +778,7 @@ impl<S: Writer, T: Ord + Hash<S>> Hash<S> for TreeSet<T> {
 mod test {
     use std::prelude::*;
     use std::hash;
+    use vec::Vec;
 
     use super::TreeSet;
 
@@ -933,6 +1017,46 @@ mod test {
     }
 
     #[test]
+    fn test_bit_or() {
+        let a: TreeSet<int> = vec![1, 3, 5, 9, 11, 16, 19, 24].into_iter().collect();
+        let b: TreeSet<int> = vec![-2, 1, 5, 9, 13, 19].into_iter().collect();
+
+        let set: TreeSet<int> = a | b;
+        let v: Vec<int> = set.into_iter().collect();
+        assert_eq!(v, vec![-2, 1, 3, 5, 9, 11, 13, 16, 19, 24]);
+    }
+
+    #[test]
+    fn test_bit_and() {
+        let a: TreeSet<int> = vec![11, 1, 3, 77, 103, 5, -5].into_iter().collect();
+        let b: TreeSet<int> = vec![2, 11, 77, -9, -42, 5, 3].into_iter().collect();
+
+        let set: TreeSet<int> = a & b;
+        let v: Vec<int> = set.into_iter().collect();
+        assert_eq!(v, vec![3, 5, 11, 77]);
+    }
+
+    #[test]
+    fn test_bit_xor() {
+        let a: TreeSet<int> = vec![1, 3, 5, 9, 11].into_iter().collect();
+        let b: TreeSet<int> = vec![-2, 3, 9, 14, 22].into_iter().collect();
+
+        let set: TreeSet<int> = a ^ b;
+        let v: Vec<int> = set.into_iter().collect();
+        assert_eq!(v, vec![-2, 1, 5, 11, 14, 22]);
+    }
+
+    #[test]
+    fn test_sub() {
+        let a: TreeSet<int> = vec![-5, 11, 22, 33, 40, 42].into_iter().collect();
+        let b: TreeSet<int> = vec![-12, -5, 14, 23, 34, 38, 39, 50].into_iter().collect();
+
+        let set: TreeSet<int> = a - b;
+        let v: Vec<int> = set.into_iter().collect();
+        assert_eq!(v, vec![11, 22, 33, 40, 42]);
+    }
+
+    #[test]
     fn test_zip() {
         let mut x = TreeSet::new();
         x.insert(5u);