- Added filter/partition functionality to refinement types.

1 files changed, 35 insertions, 1 deletions

diff --git a/stdlib/source/lux/type/refinement.lux b/stdlib/source/lux/type/refinement.lux
index 051166eb9..7431c93fe 100644
--- a/stdlib/source/lux/type/refinement.lux
+++ b/stdlib/source/lux/type/refinement.lux
@@ -16,10 +16,13 @@
   {#value t
    #predicate (Pred t)}
 
+  (type: #export (Refiner t r)
+    (-> t (Maybe (Ref t r))))
+
   (def: #export (refinement predicate)
     (All [t]
       (Ex [r]
-        (-> (Pred t) (-> t (Maybe (Ref t r))))))
+        (-> (Pred t) (Refiner t r))))
     (function (_ un-refined)
       (if (predicate un-refined)
         (#.Some (@abstraction {#value un-refined
@@ -48,6 +51,37 @@
           #.None))))
   )
 
+(def: #export (filter refiner values)
+  (All [t r] (-> (Refiner t r) (List t) (List (Ref t r))))
+  (case values
+    #.Nil
+    #.Nil
+
+    (#.Cons head tail)
+    (case (refiner head)
+      (#.Some refined)
+      (#.Cons refined (filter refiner tail))
+      
+      #.None
+      (filter refiner tail))))
+
+(def: #export (partition refiner values)
+  (All [t r] (-> (Refiner t r) (List t) [(List (Ref t r)) (List t)]))
+  (case values
+    #.Nil
+    [#.Nil #.Nil]
+
+    (#.Cons head tail)
+    (let [[yes no] (partition refiner tail)]
+      (case (refiner head)
+        (#.Some refined)
+        [(#.Cons refined yes)
+         no]
+        
+        #.None
+        [yes
+         (#.Cons head no)]))))
+
 (def: (refinement-type constructor-type)
   (-> Type (Error Type))
   (<| (poly.run constructor-type)