sub for sets

Summary: as per title

Reviewed By: stepancheg

Differential Revision: D63471534

fbshipit-source-id: 31da09aa94f614d102938643ffa686b564205c8e

2 files changed, 53 insertions, 2 deletions

diff --git a/starlark/src/values/types/set/value.rs b/starlark/src/values/types/set/value.rs
index c2c23d8a..29046099 100644
--- a/starlark/src/values/types/set/value.rs
+++ b/starlark/src/values/types/set/value.rs
@@ -317,6 +317,34 @@ where
         }
         Ok(heap.alloc(data))
     }
+
+    // Set difference
+    //TODO(romanp) implement difference on small_set level and reuse it here and in difference function
+    fn sub(&self, rhs: Value<'v>, heap: &'v Heap) -> crate::Result<Value<'v>> {
+        let rhs = SetRef::from_value(rhs)
+            .map_or_else(|| ValueError::unsupported_with(self, "-", rhs), Ok)?;
+
+        if self.0.content().is_empty() {
+            return Ok(heap.alloc(SetData {
+                content: SmallSet::new(),
+            }));
+        }
+
+        if rhs.content.is_empty() {
+            return Ok(heap.alloc(SetData {
+                content: self.0.content().clone(),
+            }));
+        }
+
+        let mut data = SetData::default();
+
+        for elem in self.0.content().iter_hashed() {
+            if !rhs.contains_hashed(elem.copied()) {
+                data.add_hashed(elem.copied());
+            }
+        }
+        Ok(heap.alloc(data))
+    }
 }
 
 impl<'v, T: SetLike<'v>> Serialize for SetGen<T> {
@@ -421,4 +449,27 @@ mod tests {
             "Operation `^` not supported for types `set` and `list`",
         );
     }
+
+    #[test]
+    fn test_sub() {
+        assert::eq("set([1, 2, 3]) - set([2])", "set([1, 3])")
+    }
+
+    #[test]
+    fn test_sub_empty_lhs() {
+        assert::eq("set([]) - set([2])", "set([])")
+    }
+
+    #[test]
+    fn test_sub_empty_rhs() {
+        assert::eq("set([1, 2]) - set([])", "set([2, 1])")
+    }
+
+    #[test]
+    fn test_sub_fail_iter() {
+        assert::fail(
+            "set([1, 2, 3]) - []",
+            "Operation `-` not supported for types `set` and `list`",
+        );
+    }
 }
diff --git a/starlark/testcases/eval/go/set.star b/starlark/testcases/eval/go/set.star
index b030bd9d..28a11e4e 100644
--- a/starlark/testcases/eval/go/set.star
+++ b/starlark/testcases/eval/go/set.star
@@ -173,8 +173,8 @@ 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]))
+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]))
 
 # # issuperset: set >= set or set.issuperset(iterable)
 # asserts.true(set([1,2,3]).issuperset([1,2]))