difference

Summary: set([1,2,3,4]) - set([1,2]) == set([3,4])

Reviewed By: stepancheg

Differential Revision: D60527927

fbshipit-source-id: 3f97e2362711a80f6c0d6d650442eddd875da827

2 files changed, 70 insertions, 6 deletions

diff --git a/starlark/src/values/types/set/methods.rs b/starlark/src/values/types/set/methods.rs
index 6fe4d784..01837b0c 100644
--- a/starlark/src/values/types/set/methods.rs
+++ b/starlark/src/values/types/set/methods.rs
@@ -261,6 +261,46 @@ pub(crate) fn set_methods(builder: &mut MethodsBuilder) {
             None => Err(value_error!("pop from an empty set")),
         }
     }
+
+    /// Returns a new set with elements unique the set when compared to the specified iterable.
+    /// ```
+    /// # starlark::assert::is_true(r#"
+    /// x = set([1, 2, 3])
+    /// y = [3, 4, 5]
+    /// x.difference(y) == set([1, 2])
+    /// # "#);
+    /// ```
+    fn difference<'v>(
+        this: SetRef<'v>,
+        #[starlark(require=pos)] other: ValueOfUnchecked<'v, StarlarkIter<Value<'v>>>,
+        heap: &'v Heap,
+    ) -> starlark::Result<SetData<'v>> {
+        //TODO(romanp) check if other is set
+        let other_it = other.get().iterate(heap)?;
+
+        if this.content.is_empty() {
+            return Ok(SetData::default());
+        }
+
+        let mut other_set: SmallSet<Value<'v>> = SmallSet::default();
+        for elem in other_it {
+            other_set.insert_hashed(elem.get_hashed()?);
+        }
+
+        if other_set.is_empty() {
+            return Ok(SetData {
+                content: this.content.clone(),
+            });
+        }
+
+        let mut data = SetData::default();
+        for elem in this.content.iter_hashed() {
+            if !other_set.contains_hashed(elem) {
+                data.add_hashed(elem.copied());
+            }
+        }
+        Ok(data)
+    }
 }
 #[cfg(test)]
 mod tests {
@@ -436,4 +476,29 @@ mod tests {
     fn test_pop_empty() {
         assert::fail("x = set([]); x.pop()", "pop from an empty set");
     }
+
+    #[test]
+    fn test_difference() {
+        assert::eq("set([1, 2, 3]).difference(set([2]))", "set([1, 3])")
+    }
+
+    #[test]
+    fn test_difference_iter() {
+        assert::eq("set([1, 2, 3]).difference([3, 2])", "set([1])")
+    }
+
+    #[test]
+    fn test_difference_order() {
+        assert::eq("list(set([3, 2, 1]).difference([2]))", "[3, 1]")
+    }
+
+    #[test]
+    fn test_difference_empty_lhs() {
+        assert::eq("set([]).difference(set([2]))", "set([])")
+    }
+
+    #[test]
+    fn test_difference_empty_rhs() {
+        assert::eq("set([1, 2]).difference(set([]))", "set([2, 1])")
+    }
 }
diff --git a/starlark/testcases/eval/go/set.star b/starlark/testcases/eval/go/set.star
index f8a3b5eb..b030bd9d 100644
--- a/starlark/testcases/eval/go/set.star
+++ b/starlark/testcases/eval/go/set.star
@@ -145,7 +145,6 @@ freeze(discard_set)
 asserts.eq(discard_set.discard(3), None)  # no mutation of frozen set because key doesn't exist
 asserts.fails(lambda: discard_set.discard(1), "discard: cannot delete from frozen hash table")
 
-
 # pop
 pop_set = set([1,2,3])
 asserts.eq(pop_set.pop(), 1)
@@ -168,11 +167,11 @@ other_clear_set = set([1,2,3])
 freeze(other_clear_set)
 asserts.fails(lambda: other_clear_set.clear(), "clear: cannot clear frozen hash table")
 
-# # difference: set - set or set.difference(iterable)
-# asserts.eq(set([1,2,3,4]).difference([1,2,3,4]), set([]))
-# asserts.eq(set([1,2,3,4]).difference([1,2]), set([3,4]))
-# asserts.eq(set([1,2,3,4]).difference([]), set([1,2,3,4]))
-# asserts.eq(set([1,2,3,4]).difference(set([1,2,3])), set([4]))
+# difference: set - set or set.difference(iterable)
+asserts.eq(set([1,2,3,4]).difference([1,2,3,4]), set([]))
+asserts.eq(set([1,2,3,4]).difference([1,2]), set([3,4]))
+asserts.eq(set([1,2,3,4]).difference([]), set([1,2,3,4]))
+asserts.eq(set([1,2,3,4]).difference(set([1,2,3])), set([4]))
 
 # asserts.eq(set([1,2,3,4]) - set([1,2,3,4]), set())
 # asserts.eq(set([1,2,3,4]) - set([1,2]), set([3,4]))