pop

Summary: set([1,2,3]).pop() ==  1

Reviewed By: stepancheg

Differential Revision: D60514255

fbshipit-source-id: 055a3a84cc00ee858c038cdbc67b23e53dff3615

starlark/src/tests/go.rs

@@ -105,6 +105,7 @@ fn test_go() {
                 "cannot insert into frozen hash table", // We don't actually have freeze
                 "cannot clear frozen hash table",
                 "discard: cannot delete from frozen hash table",
+                "pop: cannot delete from frozen hash table",
             ],
         ),
         &[],

starlark/src/values/types/set/methods.rs

@@ -231,6 +231,36 @@ pub(crate) fn set_methods(builder: &mut MethodsBuilder) {
         set.remove_hashed(hashed.as_ref());
         Ok(NoneType)
     }
+
+    /// Removes and returns the first element of a set.
+    ///
+    /// `S.pop()` removes and returns the first element of the set S.
+    ///
+    /// `pop` fails if the set is empty, or if the set is frozen or has active iterators.
+    /// Time complexity of this operation is *O(N)* where *N* is the number of entries in the set.
+    ///
+    /// ```
+    /// # starlark::assert::is_true(r#"
+    /// x = set([1, 2, 3])
+    /// # (
+    /// x.pop() == 1
+    /// # and
+    /// x.pop() == 2
+    /// # and
+    /// x == set([3])
+    /// # )"#);
+    /// ```
+    fn pop<'v>(this: Value<'v>) -> starlark::Result<Value<'v>> {
+        let mut set = SetMut::from_value(this)?;
+        let first = set.iter_hashed().next();
+        match first {
+            Some(x) => {
+                set.remove_hashed(x.as_ref());
+                Ok(x.into_key())
+            }
+            None => Err(value_error!("pop from an empty set")),
+        }
+    }
 }
 #[cfg(test)]
 mod tests {
@@ -396,4 +426,14 @@ mod tests {
     fn test_discard_multiple_times() {
         assert::eq("x = set([0, 1]); x.discard(0); x.discard(0); x", "set([1])");
     }
+
+    #[test]
+    fn test_pop() {
+        assert::is_true("x = set([1, 0]); (x.pop() == 1 and x.pop() == 0 and x == set())");
+    }
+
+    #[test]
+    fn test_pop_empty() {
+        assert::fail("x = set([]); x.pop()", "pop from an empty set");
+    }
 }

starlark/testcases/eval/go/set.star

@@ -146,16 +146,16 @@ asserts.eq(discard_set.discard(3), None)  # no mutation of frozen set because ke
 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)
-# asserts.eq(pop_set.pop(), 2)
-# asserts.eq(pop_set.pop(), 3)
-# asserts.fails(lambda: pop_set.pop(), "pop: empty set")
-# pop_set.add(1)
-# pop_set.add(2)
-# freeze(pop_set)
-# asserts.fails(lambda: pop_set.pop(), "pop: cannot delete from frozen hash table")
+# pop
+pop_set = set([1,2,3])
+asserts.eq(pop_set.pop(), 1)
+asserts.eq(pop_set.pop(), 2)
+asserts.eq(pop_set.pop(), 3)
+asserts.fails(lambda: pop_set.pop(), "pop: empty set")
+pop_set.add(1)
+pop_set.add(2)
+freeze(pop_set)
+asserts.fails(lambda: pop_set.pop(), "pop: cannot delete from frozen hash table")
 
 # # clear
 # clear_set = set([1,2,3])