basic forIn support

Author: datokrat

src/Std/Data/Iterators/Basic.lean

@@ -360,6 +360,13 @@ def IterM.IsPlausibleSuccessorOf {α : Type w} {m : Type w → Type w'} {β : Ty
     (it' it : IterM (α := α) m β) : Prop :=
   ∃ step, step.successor = some it' ∧ it.IsPlausibleStep step
 
+inductive IterM.IsPlausibleIndirectOutput {α β : Type w} {m : Type w → Type w'} [Iterator α m β]
+    : IterM (α := α) m β → β → Prop where
+  | direct {it : IterM (α := α) m β} {out : β} : it.IsPlausibleOutput out →
+      it.IsPlausibleIndirectOutput out
+  | indirect {it it' : IterM (α := α) m β} {out : β} : it'.IsPlausibleSuccessorOf it →
+      it'.IsPlausibleIndirectOutput out → it.IsPlausibleIndirectOutput out
+
 def IterM.IsInLineageOf {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
     (it' it : IterM (α := α) m β) : Prop :=
   it' = it ∨ Relation.TransGen IterM.IsPlausibleSuccessorOf it' it
@@ -475,6 +482,34 @@ def Iter.IsPlausibleSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β]
     (it' it : Iter (α := α) β) : Prop :=
   it'.toIterM.IsPlausibleSuccessorOf it.toIterM
 
+inductive Iter.IsPlausibleIndirectOutput {α β : Type w} [Iterator α Id β] :
+    Iter (α := α) β → β → Prop where
+  | direct {it : Iter (α := α) β} {out : β} : it.IsPlausibleOutput out →
+      it.IsPlausibleIndirectOutput out
+  | indirect {it it' : Iter (α := α) β} {out : β} : it'.IsPlausibleSuccessorOf it →
+      it'.IsPlausibleIndirectOutput out → it.IsPlausibleIndirectOutput out
+
+theorem Iter.isPlausibleIndirectOutput_iff_isPlausibleIndirectOutput_toIterM {α β : Type w}
+    [Iterator α Id β] {it : Iter (α := α) β} {out : β} :
+    it.IsPlausibleIndirectOutput out ↔ it.toIterM.IsPlausibleIndirectOutput out := by
+  constructor
+  · intro h
+    induction h with
+    | direct h =>
+      exact .direct h
+    | indirect h _ ih =>
+      exact .indirect h ih
+  · intro h
+    rw [← Iter.toIter_toIterM (it := it)]
+    generalize it.toIterM = it at ⊢ h
+    induction h with
+    | direct h =>
+      exact .direct h
+    | indirect h h' ih =>
+      rename_i it it' out
+      replace h : it'.toIter.IsPlausibleSuccessorOf it.toIter := h
+      exact .indirect (α := α) h ih
+
 /--
 Asserts that a certain iterator `it'` could plausibly be the directly succeeding iterator of another
 given iterator `it` while no value is emitted (see `IterStep.skip`).

src/Std/Data/Iterators/Combinators/Monadic/Take.lean

@@ -145,13 +145,7 @@ instance Take.instIteratorLoop [Monad m] [Monad n] [Iterator α m β]
 
 instance Take.instIteratorForPartial [Monad m] [Monad n] [Iterator α m β]
     [IteratorLoopPartial α m n] [MonadLiftT m n] :
-    IteratorLoopPartial (Take α m β) m n where
-  forInPartial lift {γ} it init f := do
-    Prod.fst <$> IteratorLoopPartial.forInPartial lift it.internalState.inner (γ := γ × Nat)
-        (init, it.internalState.remaining)
-        fun out acc =>
-          match acc.snd with
-          | 0 => pure <| .done acc
-          | n + 1 => (fun | .yield x => .yield ⟨x, n⟩ | .done x => .done ⟨x, n⟩) <$> f out acc.fst
+    IteratorLoopPartial (Take α m β) m n :=
+  .defaultImplementation
 
 end Std.Iterators

src/Std/Data/Iterators/Combinators/Monadic/TakeWhile.lean

@@ -155,7 +155,7 @@ it terminates.
 -/
 @[always_inline, inline]
 def IterM.takeWhile [Monad m] (P : β → Bool) (it : IterM (α := α) m β) :=
-  (it.takeWhileM (pure ∘ ULift.up ∘ P) : IterM m β)
+  (it.takeWhileWithPostcondition (fun x => PostconditionT.pure (.up <| P x)) : IterM m β)
 
 /--
 `it.PlausibleStep step` is the proposition that `step` is a possible next step from the
@@ -239,14 +239,7 @@ instance TakeWhile.instIteratorLoop [Monad m] [Monad n] [Iterator α m β]
 
 instance TakeWhile.instIteratorLoopPartial [Monad m] [Monad n] [Iterator α m β]
     [IteratorLoopPartial α m n] [MonadLiftT m n] {P} :
-    IteratorLoopPartial (TakeWhile α m β P) m n where
-  forInPartial lift {γ} it init f := do
-    IteratorLoopPartial.forInPartial lift it.internalState.inner (γ := γ)
-        init
-        fun out acc => do match ← (P out).operation with
-          | ⟨.up true, _⟩ => match ← f out acc with
-            | .yield c => pure (.yield c)
-            | .done c => pure (.done c)
-          | ⟨.up false, _⟩ => pure (.done acc)
+    IteratorLoopPartial (TakeWhile α m β P) m n :=
+  .defaultImplementation
 
 end Std.Iterators

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

@@ -24,9 +24,11 @@ namespace Std.Iterators
 
 instance (α : Type w) (β : Type w) (n : Type w → Type w') [Monad n]
     [Iterator α Id β] [Finite α Id] [IteratorLoop α Id n] :
-    ForIn n (Iter (α := α) β) β where
-  forIn it init f :=
-    IteratorLoop.finiteForIn (fun δ (c : Id δ) => pure c.run) |>.forIn it.toIterM init f
+    ForIn' n (Iter (α := α) β) β ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ where
+  forIn' it init f :=
+    IteratorLoop.finiteForIn' (fun δ (c : Id δ) => pure c.run) |>.forIn' it.toIterM init
+        fun out h acc =>
+          f out (Iter.isPlausibleIndirectOutput_iff_isPlausibleIndirectOutput_toIterM.mpr h) acc
 
 instance (α : Type w) (β : Type w) (n : Type w → Type w') [Monad n]
     [Iterator α Id β] [IteratorLoopPartial α Id n] :

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

@@ -62,8 +62,9 @@ class IteratorLoop (α : Type w) (m : Type w → Type w') {β : Type w} [Iterato
   forIn : ∀ (_lift : (γ : Type w) → m γ → n γ) (γ : Type w),
       (plausible_forInStep : β → γ → ForInStep γ → Prop) →
       IteratorLoop.WellFounded α m plausible_forInStep →
-      IterM (α := α) m β → γ →
-      ((b : β) → (c : γ) → n (Subtype (plausible_forInStep b c))) → n γ
+      (it : IterM (α := α) m β) → γ →
+      ((b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (Subtype (plausible_forInStep b c))) →
+      n γ
 
 /--
 `IteratorLoopPartial α m` provides efficient implementations of loop-based consumers for `α`-based
@@ -76,7 +77,8 @@ provided by the standard library.
 class IteratorLoopPartial (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β]
     (n : Type w → Type w'') where
   forInPartial : ∀ (_lift : (γ : Type w) → m γ → n γ) {γ : Type w},
-      IterM (α := α) m β → γ → ((b : β) → (c : γ) → n (ForInStep γ)) → n γ
+      (it : IterM (α := α) m β) → γ →
+      ((b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (ForInStep γ)) → n γ
 
 class IteratorSize (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] where
   size : IterM (α := α) m β → m (ULift Nat)
@@ -115,16 +117,20 @@ def IterM.DefaultConsumers.forIn {m : Type w → Type w'} {α : Type w} {β : Ty
     (plausible_forInStep : β → γ → ForInStep γ → Prop)
     (wf : IteratorLoop.WellFounded α m plausible_forInStep)
     (it : IterM (α := α) m β) (init : γ)
-    (f : (b : β) → (c : γ) → n (Subtype (plausible_forInStep b c))) : n γ :=
+    (f : (b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (Subtype (plausible_forInStep b c))) : n γ :=
   haveI : WellFounded _ := wf
   letI : MonadLift m n := ⟨fun {γ} => lift γ⟩
   do
     match ← it.step with
-    | .yield it' out _ =>
-      match ← f out init with
-      | ⟨.yield c, _⟩ => IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' c f
+    | .yield it' out h =>
+      match ← f out (.direct ⟨_, h⟩) init with
+      | ⟨.yield c, _⟩ =>
+        IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' c
+          (fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc)
       | ⟨.done c, _⟩ => return c
-    | .skip it' _ => IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' init f
+    | .skip it' h =>
+      IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' init
+        (fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc)
     | .done _ => return init
 termination_by IteratorLoop.WFRel.mk wf it init
 decreasing_by
@@ -159,15 +165,19 @@ partial def IterM.DefaultConsumers.forInPartial {m : Type w → Type w'} {α : T
     {n : Type w → Type w''} [Monad n]
     (lift : ∀ γ, m γ → n γ) (γ : Type w)
     (it : IterM (α := α) m β) (init : γ)
-    (f : (b : β) → (c : γ) → n (ForInStep γ)) : n γ :=
+    (f : (b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (ForInStep γ)) : n γ :=
   letI : MonadLift m n := ⟨fun {γ} => lift γ⟩
   do
     match ← it.step with
-    | .yield it' out _ =>
-      match ← f out init with
-      | .yield c => IterM.DefaultConsumers.forInPartial lift _ it' c f
+    | .yield it' out h =>
+      match ← f out (.direct ⟨_, h⟩) init with
+      | .yield c =>
+        IterM.DefaultConsumers.forInPartial lift _ it' c
+          fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc
       | .done c => return c
-    | .skip it' _ => IterM.DefaultConsumers.forInPartial lift _ it' init f
+    | .skip it' h =>
+      IterM.DefaultConsumers.forInPartial lift _ it' init
+        fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc
     | .done _ => return init
 
 /--
@@ -202,27 +212,28 @@ theorem IteratorLoop.wellFounded_of_finite {m : Type w → Type w'}
     exact WellFoundedRelation.wf
 
 /--
-This `ForIn`-style loop construct traverses a finite iterator using an `IteratorLoop` instance.
+This `ForIn'`-style loop construct traverses a finite iterator using an `IteratorLoop` instance.
 -/
 @[always_inline, inline]
-def IteratorLoop.finiteForIn {m : Type w → Type w'} {n : Type w → Type w''}
+def IteratorLoop.finiteForIn' {m : Type w → Type w'} {n : Type w → Type w''}
     {α : Type w} {β : Type w} [Iterator α m β] [Finite α m] [IteratorLoop α m n]
     (lift : ∀ γ, m γ → n γ) :
-    ForIn n (IterM (α := α) m β) β where
-  forIn {γ} [Monad n] it init f :=
+    ForIn' n (IterM (α := α) m β) β ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ where
+  forIn' {γ} [Monad n] it init f :=
     IteratorLoop.forIn (α := α) (m := m) lift γ (fun _ _ _ => True)
       wellFounded_of_finite
-      it init ((⟨·, .intro⟩) <$> f · ·)
+      it init (fun out h acc => (⟨·, .intro⟩) <$> f out h acc)
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
     {α : Type w} {β : Type w} [Iterator α m β] [Finite α m] [IteratorLoop α m n]
     [MonadLiftT m n] :
-    ForIn n (IterM (α := α) m β) β := IteratorLoop.finiteForIn (fun _ => monadLift)
+    ForIn' n (IterM (α := α) m β) β ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ :=
+  IteratorLoop.finiteForIn' (fun _ => monadLift)
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
     {α : Type w} {β : Type w} [Iterator α m β] [IteratorLoopPartial α m n] [MonadLiftT m n] :
-    ForIn n (IterM.Partial (α := α) m β) β where
-  forIn it init f :=
+    ForIn' n (IterM.Partial (α := α) m β) β ⟨fun it out => it.it.IsPlausibleIndirectOutput out⟩ where
+  forIn' it init f :=
     IteratorLoopPartial.forInPartial (α := α) (m := m) (fun _ => monadLift) it.it init f
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}

src/Std/Data/Iterators/Lemmas/Combinators/Monadic/TakeWhile.lean

@@ -48,18 +48,35 @@ theorem IterM.step_takeWhile {α m β} [Monad m] [LawfulMonad m] [Iterator α m
     {it : IterM (α := α) m β} {P} :
     (it.takeWhile P).step = (do
       match ← it.step with
-      | .yield it' out h => match P out with
-        | true => pure <| .yield (it'.takeWhile P) out (.yield h True.intro)
-        | false => pure <| .done (.rejected h True.intro)
+      | .yield it' out h => match h' : P out with
+        | true => pure <| .yield (it'.takeWhile P) out (.yield h (congrArg ULift.up h'))
+        | false => pure <| .done (.rejected h (congrArg ULift.up h'))
       | .skip it' h => pure <| .skip (it'.takeWhile P) (.skip h)
       | .done h => pure <| .done (.done h)) := by
-  simp only [takeWhile, step_takeWhileM]
+  simp only [takeWhile, step_takeWhileWithPostcondition]
   apply bind_congr
   intro step
   cases step using PlausibleIterStep.casesOn
   · simp only [Function.comp_apply, PlausibleIterStep.yield, PlausibleIterStep.done, pure_bind]
-    cases P _ <;> rfl
+    simp only [PostconditionT.pure, pure_bind]
+    split <;> split <;> simp_all
   · simp
   · simp
 
+theorem TakeWhile.Monadic.isPlausibleSuccessorOf_inner_of_isPlausibleSuccessorOf {α m β P} [Monad m] [Iterator α m β]
+    {it' it : IterM (α := TakeWhile α m β P) m β} (h : it'.IsPlausibleSuccessorOf it) :
+    it'.internalState.inner.IsPlausibleSuccessorOf it.internalState.inner := by
+  rcases h with ⟨step, h, h'⟩
+  cases step
+  · rename_i out
+    cases h
+    cases h'
+    rename_i it' h' h''
+    exact ⟨.yield it' out, rfl, h'⟩
+  · cases h
+    cases h'
+    rename_i it' h'
+    exact ⟨.skip it', rfl, h'⟩
+  · cases h
+
 end Std.Iterators

src/Std/Data/Iterators/Lemmas/Combinators/TakeWhile.lean

@@ -18,16 +18,16 @@ theorem Iter.takeWhile_eq {α β} [Iterator α Id β] {P}
 theorem Iter.step_takeWhile {α β} [Iterator α Id β] {P}
     {it : Iter (α := α) β} :
     (it.takeWhile P).step = (match it.step with
-      | .yield it' out h => match P out with
-        | true => .yield (it'.takeWhile P) out (.yield h True.intro)
-        | false => .done (.rejected h True.intro)
+      | .yield it' out h => match h' : P out with
+        | true => .yield (it'.takeWhile P) out (.yield h (congrArg _ h'))
+        | false => .done (.rejected h (congrArg _ h'))
       | .skip it' h => .skip (it'.takeWhile P) (.skip h)
       | .done h => .done (.done h)) := by
   simp [Iter.takeWhile_eq, Iter.step, toIterM_toIter, IterM.step_takeWhile]
   generalize it.toIterM.step.run = step
   cases step using PlausibleIterStep.casesOn
   · simp only [IterM.Step.toPure_yield, PlausibleIterStep.yield, toIter_toIterM, toIterM_toIter]
-    cases P _ <;> rfl
+    split <;> split <;> (try exfalso; simp_all; done) <;> simp_all
   · simp
   · simp
 
@@ -43,36 +43,41 @@ theorem Iter.atIdxSlow?_takeWhile {α β}
     simp only [Option.any_some]
     apply Eq.symm
     split
-    · cases h' : P out
-      · exfalso; simp_all
-      · simp
-    · cases h' : P out
+    · -- Ugh, this is so ugly
+      have : (match h' : P out with | true => Unit.unit | false => Unit.unit) = .unit := sorry
+      split at this
       · simp
       · exfalso; simp_all
+    · have : (match h' : P out with | true => Unit.unit | false => Unit.unit) = .unit := sorry
+      split at this
+      · exfalso; simp_all
+      · simp
   case case2 it it' out h h' l ih =>
     simp only [Nat.succ_eq_add_one, atIdxSlow?.eq_def (it := it.takeWhile P), step_takeWhile, h']
     simp only [atIdxSlow?.eq_def (it := it), h']
-    cases hP : P out
-    · simp
-      intro h
-      specialize h 0 (Nat.zero_le _)
-      simp at h
-      exfalso; simp_all
-    · simp [ih]
+    have : (match h' : P out with | true => Unit.unit | false => Unit.unit) = .unit := sorry
+    split at this
+    · rename_i hP
+      simp only [ih]
       split
       · rename_i h
         rw [if_pos]
         intro k hk
         split
         · exact hP
-        · simp at hk
+        · simp only [Nat.succ_eq_add_one, Nat.add_le_add_iff_right] at hk
           exact h _ hk
       · rename_i hl
         rw [if_neg]
         intro hl'
         apply hl
         intro k hk
         exact hl' (k + 1) (Nat.succ_le_succ hk)
+    · simp only [right_eq_ite_iff]
+      intro h
+      specialize h 0 (Nat.zero_le _)
+      simp only at h
+      exfalso; simp_all
   case case3 l it it' h h' ih =>
     simp only [atIdxSlow?.eq_def (it := it.takeWhile P), step_takeWhile, h', ih]
     simp only [atIdxSlow?.eq_def (it := it), h']
@@ -131,4 +136,10 @@ theorem Iter.toArray_takeWhile_of_finite {α β} [Iterator α Id β] {P}
     (it.takeWhile P).toArray = it.toArray.takeWhile P := by
   rw [← toArray_toList, ← toArray_toList, List.takeWhile_toArray, toList_takeWhile_of_finite]
 
+theorem TakeWhile.isPlausibleSuccessorOf_inner_of_isPlausibleSuccessorOf {α β P}
+    [Iterator α Id β]
+    {it' it : Iter (α := TakeWhile α Id β P) β} (h : it'.IsPlausibleSuccessorOf it) :
+    it'.internalState.inner.IsPlausibleSuccessorOf it.internalState.inner :=
+  TakeWhile.Monadic.isPlausibleSuccessorOf_inner_of_isPlausibleSuccessorOf h
+
 end Std.Iterators

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

@@ -13,13 +13,24 @@ import Std.Data.Iterators.Consumers.Loop
 
 namespace Std.Iterators
 
+theorem Iter.forIn'_eq {α β : Type w} [Iterator α Id β] [Finite α Id]
+    {m : Type w → Type w''} [Monad m] [IteratorLoop α Id m] [hl : LawfulIteratorLoop α Id m]
+    {γ : Type w} {it : Iter (α := α) β} {init : γ}
+    {f : (b : β) → it.IsPlausibleIndirectOutput b → γ → m (ForInStep γ)} :
+    ForIn'.forIn' it init f =
+      IterM.DefaultConsumers.forIn (fun _ c => pure c.run) γ (fun _ _ _ => True)
+        IteratorLoop.wellFounded_of_finite it.toIterM init
+          (fun out h acc => (⟨·, .intro⟩) <$>
+            f out (Iter.isPlausibleIndirectOutput_iff_isPlausibleIndirectOutput_toIterM.mpr h) acc) := by
+  cases hl.lawful; rfl
+
 theorem Iter.forIn_eq {α β : Type w} [Iterator α Id β] [Finite α Id]
     {m : Type w → Type w''} [Monad m] [IteratorLoop α Id m] [hl : LawfulIteratorLoop α Id m]
     {γ : Type w} {it : Iter (α := α) β} {init : γ}
     {f : β → γ → m (ForInStep γ)} :
     ForIn.forIn it init f =
       IterM.DefaultConsumers.forIn (fun _ c => pure c.run) γ (fun _ _ _ => True)
-        IteratorLoop.wellFounded_of_finite it.toIterM init ((⟨·, .intro⟩) <$> f · ·) := by
+        IteratorLoop.wellFounded_of_finite it.toIterM init (fun out _ acc => (⟨·, .intro⟩) <$> f out acc) := by
   cases hl.lawful; rfl
 
 theorem Iter.forIn_eq_forIn_toIterM {α β : Type w} [Iterator α Id β]