- Added "while" and "do-while" loops for stateful computations.

- Improved tests for lux/control/state.

2 files changed, 103 insertions, 28 deletions

diff --git a/stdlib/source/lux/control/state.lux b/stdlib/source/lux/control/state.lux
index 9ee12e93d..37135ac06 100644
--- a/stdlib/source/lux/control/state.lux
+++ b/stdlib/source/lux/control/state.lux
@@ -122,3 +122,19 @@
     (do Monad<M>
       [a ma]
       (wrap [state a]))))
+
+(def: #export (while condition body)
+  (All [s] (-> (State s Bool) (State s Unit) (State s Unit)))
+  (do Monad<State>
+    [execute? condition]
+    (if execute?
+      (do @
+        [_ body]
+        (while condition body))
+      (wrap []))))
+
+(def: #export (do-while condition body)
+  (All [s] (-> (State s Bool) (State s Unit) (State s Unit)))
+  (do Monad<State>
+    [_ body]
+    (while condition body)))
diff --git a/stdlib/test/test/lux/control/state.lux b/stdlib/test/test/lux/control/state.lux
index e02dfdf55..de1560f48 100644
--- a/stdlib/test/test/lux/control/state.lux
+++ b/stdlib/test/test/lux/control/state.lux
@@ -7,43 +7,102 @@
        (data [text "Text/" Monoid<Text>]
              text/format
              [number]
-             [product]))
+             [product])
+       ["R" math/random])
   lux/test)
 
-(test: "State"
+(def: (with-conditions [state output] computation)
+  (-> [Nat Nat] (&;State Nat Nat) Bool)
+  (|> computation
+      (&;run state)
+      product;right
+      (n.= output)))
+
+(test: "Basics"
+  [state R;nat
+   value R;nat]
+  ($_ seq
+      (assert "Can get the state as a value."
+              (with-conditions [state state]
+                &;get))
+      (assert "Can replace the state."
+              (with-conditions [state value]
+                (do &;Monad<State>
+                  [_ (&;put value)]
+                  &;get)))
+      (assert "Can update the state."
+              (with-conditions [state (n.* value state)]
+                (do &;Monad<State>
+                  [_ (&;update (n.* value))]
+                  &;get)))
+      (assert "Can use the state."
+              (with-conditions [state (n.inc state)]
+                (&;use n.inc)))
+      (assert "Can use a temporary (local) state."
+              (with-conditions [state (n.* value state)]
+                (&;local (n.* value)
+                         &;get)))
+      ))
+
+(test: "Structures"
+  [state R;nat
+   value R;nat]
   ($_ seq
-      (assert "" (i.= 123 (product;right (&;run 123 &;get))))
-      (assert "" (i.= 321 (product;right (&;run 123 (do &;Monad<State>
-                                                      [_ (&;put 321)]
-                                                      &;get)))))
-      (assert "" (i.= 369 (product;right (&;run 123 (do &;Monad<State>
-                                                      [_ (&;update (i.* 3))]
-                                                      &;get)))))
-      (assert "" (i.= 124 (product;right (&;run 123 (&;use i.inc)))))
-      (assert "" (i.= 246 (product;right (&;run 123 (&;local (i.* 2) &;get)))))
-      (assert "" (i.= 124 (product;right (&;run 123 (:: &;Functor<State> map i.inc &;get)))))
-      (assert "" (i.= 10 (product;right (&;run 123 (:: &;Applicative<State> wrap 10)))))
-      (assert "" (i.= 30 (product;right (&;run 123 (let [(^open "&/") &;Applicative<State>]
-                                                     (&/apply (&/wrap (i.+ 10)) (&/wrap 20)))))))
-      (assert "" (i.= 30 (product;right (&;run 123 (: (&;State Int Int)
-                                                      (do &;Monad<State>
-                                                        [f (wrap i.+)
-                                                         x (wrap 10)
-                                                         y (wrap 20)]
-                                                        (wrap (f x y))))))))
+      (assert "Can use functor."
+              (with-conditions [state (n.inc state)]
+                (:: &;Functor<State> map n.inc &;get)))
+      (assert "Can use applicative."
+              (let [(^open "&/") &;Applicative<State>]
+                (and (with-conditions [state value]
+                       (&/wrap value))
+                     (with-conditions [state (n.+ value value)]
+                       (&/apply (&/wrap (n.+ value))
+                                (&/wrap value))))))
+      (assert "Can use monad."
+              (with-conditions [state (n.+ value value)]
+                (: (&;State Nat Nat)
+                   (do &;Monad<State>
+                     [f (wrap n.+)
+                      x (wrap value)
+                      y (wrap value)]
+                     (wrap (f x y))))))
       ))
 
 (test: "Monad transformer"
+  [state R;nat
+   left R;nat
+   right R;nat]
   (let [lift (&;lift-state io;Monad<IO>)
         (^open "io/") io;Monad<IO>]
     (assert "Can add state functionality to any monad."
-            (|> (: (&;State' io;IO Text Int)
+            (|> (: (&;State' io;IO Nat Nat)
                    (do (&;StateT io;Monad<IO>)
-                     [a (lift (io/wrap 123))
-                      b (wrap 456)]
-                     (wrap (i.+ a b))))
-                (&;run' "")
+                     [a (lift (io/wrap left))
+                      b (wrap right)]
+                     (wrap (n.+ a b))))
+                (&;run' state)
                 io;run
-                (case> ["" 579] true
-                       _        false)))
+                (case> [state' output']
+                       (and (n.= state state')
+                            (n.= (n.+ left right) output')))))
     ))
+
+(test: "Loops"
+  [limit (|> R;nat (:: @ map (n.% +10)))
+   #let [condition (do &;Monad<State>
+                     [state &;get]
+                     (wrap (n.< limit state)))]]
+  ($_ seq
+      (assert "'while' will only execute if the condition is true."
+              (|> (&;while condition (&;update n.inc))
+                  (&;run +0)
+                  (case> [state' output']
+                         (n.= limit state'))))
+      (assert "'do-while' will execute at least once."
+              (|> (&;do-while condition (&;update n.inc))
+                  (&;run +0)
+                  (case> [state' output']
+                         (or (n.= limit state')
+                             (and (n.= +0 limit)
+                                  (n.= +1 state'))))))
+      ))