add tests

Author: datokrat

src/Std/Data/Iterators/Lemmas/Consumers/Loop.lean

@@ -138,12 +138,13 @@ theorem Iter.fold_eq_foldM {α β γ : Type w} [Iterator α Id β]
     it.fold (init := init) f = (it.foldM (m := Id) (init := init) (pure <| f · ·)).run := by
   simp [foldM_eq_forIn, fold_eq_forIn]
 
-theorem Iter.forIn_yield_eq_fold {α β γ : Type w} [Iterator α Id β]
+@[simp]
+theorem Iter.forIn_pure_yield_eq_fold {α β γ : Type w} [Iterator α Id β]
     [Finite α Id] [IteratorLoop α Id Id]
     [LawfulIteratorLoop α Id Id] {f : β → γ → γ} {init : γ}
     {it : Iter (α := α) β} :
-    ForIn.forIn (m := Id) it init (fun c b => .yield (f c b)) =
-      it.fold (fun b c => f c b) init := by
+    ForIn.forIn (m := Id) it init (fun c b => pure (.yield (f c b))) =
+      pure (it.fold (fun b c => f c b) init) := by
   simp only [fold_eq_forIn]
   rfl
 

src/Std/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean

@@ -105,6 +105,7 @@ theorem IterM.fold_eq_foldM {α β γ : Type w} {m : Type w → Type w'} [Iterat
     it.fold (init := init) f = it.foldM (init := init) (pure <| f · ·) := by
   simp [foldM_eq_forIn, fold_eq_forIn]
 
+@[simp]
 theorem IterM.forIn_pure_yield_eq_fold {α β γ : Type w} {m : Type w → Type w'} [Iterator α m β]
     [Finite α m] [Monad m] [LawfulMonad m] [IteratorLoop α m m]
     [LawfulIteratorLoop α m m] {f : β → γ → γ} {init : γ}

src/Std/Data/Iterators/Lemmas/Producers/List.lean

@@ -30,17 +30,17 @@ theorem _root_.List.step_iter_cons {x : β} {xs : List β} :
   simp only [List.iter, List.iterM, IterM.step, Iterator.step]; rfl
 
 @[simp]
-theorem _root_.List.toArray_iter {m : Type w → Type w'} [Monad m] [LawfulMonad m] {l : List β} :
+theorem _root_.List.toArray_iter {l : List β} :
     l.iter.toArray = l.toArray := by
   simp [List.iter, List.toArray_iterM, Iter.toArray_eq_toArray_toIterM]
 
 @[simp]
-theorem _root_.List.toList_iter {m : Type w → Type w'} [Monad m] [LawfulMonad m] {l : List β} :
+theorem _root_.List.toList_iter {l : List β} :
     l.iter.toList = l := by
   simp [List.iter, List.toList_iterM]
 
 @[simp]
-theorem _root_.List.toListRev_iter {m : Type w → Type w'} [Monad m] [LawfulMonad m] {l : List β} :
+theorem _root_.List.toListRev_iter {l : List β} :
     l.iter.toListRev = l.reverse := by
   simp [List.iter, Iter.toListRev_eq_toListRev_toIterM, List.toListRev_iterM]
 

tests/lean/run/iterators.lean

@@ -34,14 +34,32 @@ section ListIteratorBasic
 #guard_msgs in
 #eval [1, 2, 3].iterM IO |>.toArray
 
+example : [1, 2, 3].iter.toList = [1, 2, 3] := by
+  simp
+
+example : [1, 2, 3].iter.toArray = #[1, 2, 3] := by
+  simp
+
+example : [1, 2, 3].iter.toListRev = [3, 2, 1] := by
+  simp
+
+-- This does not work for `IO` because `LawfulMonad IO` is in Batteries
+example : ([1, 2, 3].iterM (StateM Nat)).toList = pure [1, 2, 3] := by
+  simp
+
+example : ([1, 2, 3].iterM (StateM Nat)).toArray = pure #[1, 2, 3] := by
+  simp
+
+example : ([1, 2, 3].iterM (StateM Nat)).toListRev = pure [3, 2, 1] := by
+  simp
+
 end ListIteratorBasic
 
 section WellFoundedRecursion
 
 def sum (l : List Nat) : Nat :=
   go l.iter 0
 where
-  @[specialize] -- The old code generator seems to need this.
   go it acc :=
     match it.step with
     | .yield it' out _ => go it' (acc + out)
@@ -64,4 +82,27 @@ def sumFold (l : List Nat) : Nat :=
 #guard_msgs in
 #eval [1, 2, 3].iter.fold (init := 0) (· + · )
 
+example (l : List Nat) : l.iter.fold (init := 0) (· + ·) = l.sum := by
+  simp [← Std.Iterators.Iter.foldl_toList, List.sum]
+  -- It remains to show that List.foldl (· + ·) = List.foldr (· + ·)
+  generalize 0 = init
+  induction l generalizing init
+  · rfl
+  · rename_i x l ih
+    simp only [List.foldl_cons, ih, List.foldr_cons]
+    clear ih
+    induction l generalizing x
+    · simp only [List.foldr_nil]
+      omega
+    · simp only [List.foldr_cons, *]
+      omega
+
+example (l : List Nat) :
+  (Id.run do
+    let mut s := 0
+    for x in l.iter do
+      s := s + x
+    return s) = l.iter.fold (init := 0) (· + ·) := by
+  simp
+
 end Loop