Add Dict.keepIf, Dict.dropIf, Set.keepIf, Set.dropIf

crates/compiler/builtins/roc/Dict.roc

@@ -15,6 +15,8 @@ interface Dict
         update,
         walk,
         walkUntil,
+        keepIf,
+        dropIf,
         toList,
         fromList,
         keys,
@@ -316,6 +318,47 @@ walkUntil : Dict k v, state, (state, k, v -> [Continue state, Break state]) -> s
 walkUntil = \@Dict { data }, initialState, transform ->
     List.walkUntil data initialState (\state, (k, v) -> transform state k v)
 
+## Run the given function on each key-value pair of a dictionary, and return
+## a dictionary with just the pairs for which the function returned `Bool.true`.
+## ```
+## expect Dict.empty {}
+##     |> Dict.insert "Alice" 17
+##     |> Dict.insert "Bob" 18
+##     |> Dict.insert "Charlie" 19
+##     |> Dict.keepIf \(_k, v) -> v >= 18
+##     |> Dict.len
+##     |> Bool.isEq 2
+## ```
+keepIf : Dict k v, ((k, v) -> Bool) -> Dict k v
+keepIf = \dict, predicate ->
+    keepIfHelp dict predicate 0 (Dict.len dict)
+
+keepIfHelp : Dict k v, ((k, v) -> Bool), Nat, Nat -> Dict k v
+keepIfHelp = \@Dict dict, predicate, index, length ->
+    if index < length then
+        (key, value) = listGetUnsafe dict.data index
+        if predicate (key, value) then
+            keepIfHelp (@Dict dict) predicate (index + 1) length
+        else
+            keepIfHelp (Dict.remove (@Dict dict) key) predicate index (length - 1)
+    else
+        @Dict dict
+
+## Run the given function on each key-value pair of a dictionary, and return
+## a dictionary with just the pairs for which the function returned `Bool.false`.
+## ```
+## expect Dict.empty {}
+##     |> Dict.insert "Alice" 17
+##     |> Dict.insert "Bob" 18
+##     |> Dict.insert "Charlie" 19
+##     |> Dict.dropIf \(_k, v) -> v >= 18
+##     |> Dict.len
+##     |> Bool.isEq 1
+## ```
+dropIf : Dict k v, ((k, v) -> Bool) -> Dict k v
+dropIf = \dict, predicate ->
+    Dict.keepIf dict (\e -> Bool.not (predicate e))
+
 ## Get the value for a given key. If there is a value for the specified key it
 ## will return [Ok value], otherwise return [Err KeyNotFound].
 ## ```
@@ -1446,3 +1489,29 @@ expect
     |> Dict.insert "Charlie" 19
     |> Dict.walkUntil Bool.false (\_, _, age -> if age >= 18 then Break Bool.true else Continue Bool.false)
     |> Bool.isEq Bool.true
+
+expect
+    d1 = Dict.empty {}
+        |> Dict.insert "Alice" 17
+        |> Dict.insert "Bob" 18
+        |> Dict.insert "Charlie" 19
+        |> Dict.keepIf \(_k, v) -> v >= 18
+
+    d2 = Dict.empty {}
+        |> Dict.insert "Bob" 18
+        |> Dict.insert "Charlie" 19
+
+    d1 == d2
+
+expect
+    d1 = Dict.empty {}
+        |> Dict.insert "Alice" 17
+        |> Dict.insert "Bob" 18
+        |> Dict.insert "Charlie" 19
+        |> Dict.keepIf \(k, _v) -> Str.endsWith k "e"
+
+    d2 = Dict.empty {}
+        |> Dict.insert "Alice" 17
+        |> Dict.insert "Charlie" 19
+
+    d1 == d2

crates/compiler/builtins/roc/Set.roc

@@ -5,6 +5,8 @@ interface Set
         single,
         walk,
         walkUntil,
+        keepIf,
+        dropIf,
         insert,
         len,
         isEmpty,
@@ -344,6 +346,29 @@ walkUntil : Set k, state, (state, k -> [Continue state, Break state]) -> state w
 walkUntil = \@Set dict, state, step ->
     Dict.walkUntil dict state (\s, k, _ -> step s k)
 
+
+## Run the given function on each element in the `Set`, and return
+## a `Set` with just the elements for which the function returned `Bool.true`.
+## ```
+## expect Set.fromList [1,2,3,4,5]
+##     |> Set.keepIf \k -> k >= 3
+##     |> Bool.isEq (Set.fromList [3,4,5])
+## ```
+keepIf : Set k, (k -> Bool) -> Set k
+keepIf = \@Set dict, predicate ->
+    @Set (Dict.keepIf dict (\(k,_v) -> predicate k))
+
+## Run the given function on each element in the `Set`, and return
+## a `Set` with just the elements for which the function returned `Bool.false`.
+## ```
+## expect Set.fromList [1,2,3,4,5]
+##     |> Set.dropIf \k -> k >= 3
+##     |> Bool.isEq (Set.fromList [1,2])
+## ```
+dropIf : Set k, (k -> Bool) -> Set k
+dropIf = \@Set dict, predicate ->
+    @Set (Dict.dropIf dict (\(k,_v) -> predicate k))
+
 expect
     first =
         single "Keep Me"
@@ -443,3 +468,11 @@ expect
         |> insert orderOne
 
     wrapperOne == wrapperTwo
+
+expect Set.fromList [1,2,3,4,5]
+    |> Set.keepIf \k -> k >= 3
+    |> Bool.isEq (Set.fromList [3,4,5])
+
+expect Set.fromList [1,2,3,4,5]
+    |> Set.dropIf \k -> k >= 3
+    |> Bool.isEq (Set.fromList [1,2])

crates/compiler/module/src/symbol.rs

@@ -1484,6 +1484,8 @@ define_builtins! {
         24 DICT_IS_EMPTY: "isEmpty"
         25 DICT_MAP: "map"
         26 DICT_JOINMAP: "joinMap"
+        27 DICT_KEEP_IF: "keepIf"
+        28 DICT_DROP_IF: "dropIf"
     }
     9 SET: "Set" => {
         0 SET_SET: "Set" exposed_type=true // the Set.Set type alias
@@ -1506,6 +1508,8 @@ define_builtins! {
         17 SET_IS_EMPTY: "isEmpty"
         18 SET_MAP: "map"
         19 SET_JOIN_MAP: "joinMap"
+        20 SET_KEEP_IF: "keepIf"
+        21 SET_DROP_IF: "dropIf"
     }
     10 BOX: "Box" => {
         0 BOX_BOX_TYPE: "Box" exposed_apply_type=true // the Box.Box opaque type