While loop (no measure)

Author: gaetanserre

RandomDo/Monad/Instances.lean

@@ -29,6 +29,11 @@ instance {m : Type u → Type v} [Monad m] :
   mPure := pure
   mBind := bind
 
+/-- The unbounded loop behind `while … rdo`, at a core monad, is core's loop over `Lean.Loop`. -/
+instance {m : Type u → Type v} [Monad m] :
+    MeasurableSpaceForIn (Monad.toMeasurableSpaceMonad m) Lean.Loop Unit where
+  forIn xs b f := ForIn.forIn (m := m) xs b f
+
 /-- A measurable space monad for pseudo random number generation. -/
 abbrev PseudoRandomM := Monad.toMeasurableSpaceMonad Rand
 

RandomDo/Monad/Notation.lean

@@ -260,6 +260,17 @@ def rdoForDecl := leading_parser
           dec.continueWithUnit
   mkBindApp σ γ forIn rest
 
+/-- parser for `rdo` while loops -/
+@[doElem_parser] def rdoWhile := leading_parser
+  "while " >> withForbidden "rdo" doIfCond >> " rdo " >> doSeq
+
+/-- Define expander for `while` loops in `rdo` notation. As in core, `while c rdo body` is a loop
+over `Loop.mk` that runs `body` while `c` holds and breaks otherwise. -/
+@[macro rdoWhile] def expandRDoWhile : Macro
+  | `(rdoWhile| while%$tk $cond:doIfCond rdo $body) =>
+    `(doElem| for%$tk _ in Lean.Loop.mk rdo if $cond:doIfCond then $body else break)
+  | _ => Macro.throwUnsupported
+
 /-- Infer the `ControlInfo` of an `rdo` loop as that of the core `for` loop with the same body. -/
 @[doElem_control_info rdoFor] def controlInfoRDoFor : ControlInfoHandler := fun stx => do
   let `(rdoFor| for $_:rdoForDecl,* rdo $body) := stx | throwUnsupportedSyntax

Test/Gaps.lean

@@ -20,22 +20,42 @@ namespace Test.Gaps
 
 /-! ## Unbounded and conditional iteration
 
-TODO: `while`, `repeat` and `repeat … until` all expand to `for _ in Loop.mk do …`, which reaches
-core's `doFor` and so asks for a `ForIn` instance. Supporting them needs the macros re-pointed at
-`rdoFor` and, at `Measure`, a denotation for an iteration that need not terminate.
+`while … rdo` is a loop over `Lean.Loop`, which has an instance at the core monads only.
+TODO: at `Measure`, a denotation for an iteration that need not terminate: the least fixed point of
+its unfolding, where the runs that never stop carry no mass. And `repeat` and `repeat … until`,
+which still expand to core's `for _ in Loop.mk do …`, need `rdo` counterparts.
 -/
 
+/--
+error: failed to synthesize instance of type class
+  MeasurableSpaceForIn Measure Lean.Loop ?α
+
+Hint: Type class instance resolution failures can be inspected with the `set_option trace.Meta.synthInstance true` command.
+-/
+#guard_msgs in
+noncomputable def whileAtMeasure : Measure ℕ := rdo
+  let mut n := 0
+  let mut go := true
+  while go rdo
+    let b ← fairCoin
+    n := n + 1
+    if b then
+      go := false
+  return n
+
 /--
 error: failed to synthesize instance of type class
   ForIn IdM Lean.Loop ?α
 
 Hint: Type class instance resolution failures can be inspected with the `set_option trace.Meta.synthInstance true` command.
 -/
 #guard_msgs in
-def whileLoop : IdM ℕ := rdo
+def repeatLoop : IdM ℕ := rdo
   let mut i := 0
-  while i < 3 do
+  repeat
     i := i + 1
+    if 3 ≤ i then
+      break
   return i
 
 /-! ## Exceptions

Test/Loops.lean

@@ -1,14 +1,18 @@
 module
 
 public import Test.Common
+meta import Test.Common
+public import Std.Tactic.Do
 
 set_option linter.style.header false
+set_option linter.hashCommand false
 
 /-!
-# `rdo`: `for` loops over a single collection
+# `rdo`: `for` and `while` loops
 
 `rdo` has its own `for … rdo …` parser, expander and elaborator, mirroring core's but emitting
-`MeasurableSpaceForIn.forIn`. Instances exist for `List`, `Array` and `Vector`.
+`MeasurableSpaceForIn.forIn`. Instances exist for `List`, `Array` and `Vector`, and for `Lean.Loop`,
+which `while … rdo` loops over, at the core monads.
 
 A loop can sit under another construct, including another loop: the enclosing one learns what the
 loop does to the control flow from the `ControlInfo` handler of `rdoFor`, which is that of core's
@@ -235,6 +239,100 @@ noncomputable def countPairsOfHeads (n : ℕ) : Measure ℕ := rdo
         c := c + 1
   return c
 
+/-! ## `while` loops
+
+`while c rdo body` is a loop over `Lean.Loop`, as in core. At a core monad it is core's loop, which
+the kernel cannot unfold, so these programs are checked with `#guard` rather than `rfl`, and proved
+through `mvcgen`. There is no instance at `Measure` yet.
+-/
+
+/-- A `while` loop, counting down from `n`. -/
+def countdown (n : ℕ) : IdM ℕ := rdo
+  let mut i := n
+  let mut steps := 0
+  while 0 < i rdo
+    i := i - 1
+    steps := steps + 1
+  return steps
+
+#guard IdM.run (countdown 5) = 5
+
+#guard IdM.run (countdown 0) = 0
+
+open Std.Do in
+set_option mvcgen.warning false in
+theorem countdown_eq (n : ℕ) : IdM.run (countdown n) = n := by
+  generalize h : IdM.run (countdown n) = r
+  apply Id.of_wp_run_eq h
+  simp only [countdown, MeasurableSpaceForIn.forIn, MeasurableSpaceBind.mBind,
+    MeasurableSpacePure.mPure]
+  dsimp only [IdM, Monad.toMeasurableSpaceMonad]
+  mvcgen invariants
+  · fun st => ⟨st.1⟩
+  · ⇓ c => match c with
+      | .inl st => ⌜st.1 + st.2 = n⌝
+      | .inr st => ⌜st.2 = n⌝
+  all_goals simp_all <;> omega
+
+/-- `break` out of a `while` loop. -/
+def halveUntilOdd (n : ℕ) : IdM ℕ := rdo
+  let mut k := n
+  while 0 < k rdo
+    if k % 2 = 1 then
+      break
+    k := k / 2
+  return k
+
+#guard IdM.run (halveUntilOdd 24) = 3
+
+#guard IdM.run (halveUntilOdd 0) = 0
+
+/-- An early `return` out of a `while` loop. -/
+def firstSquareAbove (n : ℕ) : IdM ℕ := rdo
+  let mut k := 0
+  while true rdo
+    if k * k > n then
+      return k
+    k := k + 1
+  return 0
+
+#guard IdM.run (firstSquareAbove 10) = 4
+
+/-- `while let`, consuming a list one element at a time. -/
+def sumByPopping (xs : List ℕ) : IdM ℕ := rdo
+  let mut rest := xs
+  let mut s := 0
+  while let x :: xs' := rest rdo
+    s := s + x
+    rest := xs'
+  return s
+
+#guard IdM.run (sumByPopping [1, 2, 3]) = 6
+
+/-- `while h : c`, which hands the body a proof of the condition. -/
+def countdownWithProof (n : ℕ) : IdM ℕ := rdo
+  let mut i := n
+  let mut steps := 0
+  while h : 0 < i rdo
+    have : i - 1 < i := Nat.sub_lt h Nat.one_pos
+    i := i - 1
+    steps := steps + 1
+  return steps
+
+#guard IdM.run (countdownWithProof 4) = 4
+
+/-- A `while` loop nested inside a `for` loop. -/
+def sumOfLogs (xs : List ℕ) : IdM ℕ := rdo
+  let mut s := 0
+  for x in xs rdo
+    let mut k := x
+    while 1 < k rdo
+      k := k / 2
+      s := s + 1
+  return s
+
+#guard IdM.run (sumOfLogs [1, 2, 8]) = 4
+
 end Test.Loops
 
 end