ForIn' and size

Author: datokrat

src/Init/Data/Iterators/Basic.lean

@@ -372,6 +372,13 @@ def Iter.IsPlausibleStep {α : Type w} {β : Type w} [Iterator α Id β]
     (it : Iter (α := α) β) (step : IterStep (Iter (α := α) β) β) : Prop :=
   it.toIterM.IsPlausibleStep (step.mapIterator Iter.toIterM)
 
+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
+
 /--
 The type of the step object returned by `Iter.step`, containing an `IterStep`
 and a proof that this is a plausible step for the given iterator.
@@ -429,6 +436,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/Init/Data/Iterators/Consumers/Loop.lean

@@ -26,9 +26,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] :
@@ -113,4 +115,14 @@ def Iter.Partial.fold {α : Type w} {β : Type w} {γ : Type w} [Iterator α Id
     (init : γ) (it : Iter.Partial (α := α) β) : γ :=
   ForIn.forIn (m := Id) it init (fun x acc => ForInStep.yield (f acc x))
 
+@[always_inline, inline, inherit_doc IterM.size]
+def Iter.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorSize α Id]
+    (it : Iter (α := α) β) : Nat :=
+  (IteratorSize.size it.toIterM).run.down
+
+@[always_inline, inline, inherit_doc IterM.Partial.size]
+def Iter.Partial.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorSizePartial α Id]
+    (it : Iter (α := α) β) : Nat :=
+  (IteratorSizePartial.size it.toIterM).run.down
+
 end Std.Iterators

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

@@ -65,7 +65,7 @@ class IteratorLoop (α : Type w) (m : Type w → Type w') {β : Type w} [Iterato
       (plausible_forInStep : β → γ → ForInStep γ → Prop) →
       IteratorLoop.WellFounded α m plausible_forInStep →
       (it : IterM (α := α) m β) → γ →
-      ((b : β) → (c : γ) → n (Subtype (plausible_forInStep b c))) →
+      ((b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (Subtype (plausible_forInStep b c))) →
       n γ
 
 /--
@@ -80,7 +80,13 @@ class IteratorLoopPartial (α : Type w) (m : Type w → Type w') {β : Type w} [
     (n : Type w → Type w'') where
   forInPartial : ∀ (_lift : (γ : Type w) → m γ → n γ) {γ : Type w},
       (it : IterM (α := α) m β) → γ →
-      ((b : β) → (c : γ) → n (ForInStep γ)) → n γ
+      ((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)
+
+class IteratorSizePartial (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] where
+  size : IterM (α := α) m β → m (ULift Nat)
 
 end Typeclasses
 
@@ -113,20 +119,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 it' out h =>
+      match ← f out (.direct ⟨_, h⟩) init with
       | ⟨.yield c, _⟩ =>
         IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' c
-          (fun out acc => f out acc)
+          (fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc)
       | ⟨.done c, _⟩ => return c
-    | .skip it' _ =>
+    | .skip it' h =>
       IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' init
-        (fun out acc => f out acc)
+        (fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc)
     | .done _ => return init
 termination_by IteratorLoop.WFRel.mk wf it init
 decreasing_by
@@ -161,19 +167,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 it' out h =>
+      match ← f out (.direct ⟨_, h⟩) init with
       | .yield c =>
         IterM.DefaultConsumers.forInPartial lift _ it' c
-          fun out acc => f out acc
+          fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc
       | .done c => return c
-    | .skip it' _ =>
+    | .skip it' h =>
       IterM.DefaultConsumers.forInPartial lift _ it' init
-        fun out acc => f out acc
+        fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc
     | .done _ => return init
 
 /--
@@ -208,28 +214,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 (fun out acc => (⟨·, .intro⟩) <$> f out acc)
+      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''}
@@ -340,4 +346,42 @@ def IterM.Partial.drain {α : Type w} {m : Type w → Type w'} [Monad m] {β : T
     m PUnit :=
   it.fold (γ := PUnit) (fun _ _ => .unit) .unit
 
+section Size
+
+@[always_inline, inline]
+def IterM.DefaultConsumers.size {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
+    [Iterator α m β] [IteratorLoop α m m] [Finite α m] (it : IterM (α := α) m β) :
+    m (ULift Nat) :=
+  it.fold (init := .up 0) fun acc _ => .up (acc.down + 1)
+
+@[always_inline, inline]
+def IterM.DefaultConsumers.sizePartial {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
+    [Iterator α m β] [IteratorLoopPartial α m m] (it : IterM (α := α) m β) :
+    m (ULift Nat) :=
+  it.allowNontermination.fold (init := .up 0) fun acc _ => .up (acc.down + 1)
+
+@[always_inline, inline]
+instance IteratorSize.defaultImplementation {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [Finite α m] [IteratorLoop α m m] :
+    IteratorSize α m where
+  size := IterM.DefaultConsumers.size
+
+@[always_inline, inline]
+instance IteratorSize.Partial.defaultImplementation {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [Finite α m] [IteratorLoopPartial α m m] :
+    IteratorSizePartial α m where
+  size := IterM.DefaultConsumers.sizePartial
+
+@[always_inline, inline]
+def IterM.size {α : Type} {m : Type → Type w'} {β : Type} [Iterator α m β] [Monad m]
+    (it : IterM (α := α) m β) [IteratorSize α m] : m Nat :=
+  ULift.down <$> IteratorSize.size it
+
+@[always_inline, inline]
+def IterM.Partial.size {α : Type} {m : Type → Type w'} {β : Type} [Iterator α m β] [Monad m]
+    (it : IterM.Partial (α := α) m β) [IteratorSizePartial α m] : m Nat :=
+  ULift.down <$> IteratorSizePartial.size it.it
+
+end Size
+
 end Std.Iterators

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

@@ -10,14 +10,25 @@ import Init.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 : (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 acc => (⟨·, .intro⟩) <$>
+          (fun out _ acc => (⟨·, .intro⟩) <$>
             f out acc) := by
   cases hl.lawful; rfl
 

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

@@ -18,18 +18,18 @@ theorem IterM.DefaultConsumers.forIn'_eq_match_step {α β : Type w} {m : Type w
     {plausible_forInStep : β → γ → ForInStep γ → Prop}
     {wf : IteratorLoop.WellFounded α m plausible_forInStep}
     {it : IterM (α := α) m β} {init : γ}
-    {f : (b : β) → (c : γ) → n (Subtype (plausible_forInStep b c))} :
+    {f : (b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (Subtype (plausible_forInStep b c))} :
     IterM.DefaultConsumers.forIn lift γ plausible_forInStep wf it init f = (do
       match ← lift _ it.step with
-      | .yield it' out _ =>
-        match ← f out init with
+      | .yield it' out h =>
+        match ← f out (.direct ⟨_, h⟩) init with
         | ⟨.yield c, _⟩ =>
           IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' c
-            fun out acc => f out acc
+            fun out h'' acc => f out (.indirect ⟨_, rfl, h⟩ h'') acc
         | ⟨.done c, _⟩ => return c
-      | .skip it' _ =>
+      | .skip it' h =>
         IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' init
-          fun out acc => f out acc
+          fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc
       | .done _ => return init) := by
   rw [forIn]
   apply bind_congr
@@ -39,17 +39,17 @@ theorem IterM.DefaultConsumers.forIn'_eq_match_step {α β : Type w} {m : Type w
 theorem IterM.forIn'_eq {α β : Type w} {m : Type w → Type w'} [Iterator α m β] [Finite α m]
     {n : Type w → Type w''} [Monad n] [IteratorLoop α m n] [hl : LawfulIteratorLoop α m n]
     [MonadLiftT m n] {γ : Type w} {it : IterM (α := α) m β} {init : γ}
-    {f : (b : β) → γ → n (ForInStep γ)} :
-    ForIn.forIn it init f = IterM.DefaultConsumers.forIn (fun _ => monadLift) γ (fun _ _ _ => True)
-        IteratorLoop.wellFounded_of_finite it init ((⟨·, .intro⟩) <$> f · ·) := by
+    {f : (b : β) → it.IsPlausibleIndirectOutput b → γ → n (ForInStep γ)} :
+    ForIn'.forIn' it init f = IterM.DefaultConsumers.forIn (fun _ => monadLift) γ (fun _ _ _ => True)
+        IteratorLoop.wellFounded_of_finite it init ((⟨·, .intro⟩) <$> f · · ·) := by
   cases hl.lawful; rfl
 
 theorem IterM.forIn_eq {α β : Type w} {m : Type w → Type w'} [Iterator α m β] [Finite α m]
     {n : Type w → Type w''} [Monad n] [IteratorLoop α m n] [hl : LawfulIteratorLoop α m n]
     [MonadLiftT m n] {γ : Type w} {it : IterM (α := α) m β} {init : γ}
     {f : β → γ → n (ForInStep γ)} :
     ForIn.forIn it init f = IterM.DefaultConsumers.forIn (fun _ => monadLift) γ (fun _ _ _ => True)
-        IteratorLoop.wellFounded_of_finite it init (fun out acc => (⟨·, .intro⟩) <$> f out acc) := by
+        IteratorLoop.wellFounded_of_finite it init (fun out _ acc => (⟨·, .intro⟩) <$> f out acc) := by
   cases hl.lawful; rfl
 
 theorem IterM.forIn_eq_match_step {α β : Type w} {m : Type w → Type w'} [Iterator α m β]

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

@@ -239,14 +239,7 @@ instance TakeWhile.instIteratorLoop [Monad m] [Monad n] [Iterator α m β]
 
 instance TakeWhile.instIteratorForPartial [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