feat: equivalence of iterators (#8545)

This PR provides the means to reason about "equivalent" iterators.
Simply speaking, two iterators are equivalent if they behave the same as
long as consumers do not introspect their states.

Author: datokrat

src/Std/Data/Internal/LawfulMonadLiftFunction.lean

@@ -10,7 +10,7 @@ import Init.NotationExtra
 import Init.Control.Lawful.MonadLift
 
 /-!
-# Typeclass for lawfule monad lifting functions
+# Typeclass for lawful monad lifting functions
 
 This module provides a typeclass `LawfulMonadLiftFunction f` that asserts that a function `f`
 mapping values from one monad to another monad commutes with `pure` and `bind`. This equivalent to
@@ -60,9 +60,13 @@ theorem LawfulMonadLiftFunction.lift_seqRight [LawfulMonad m] [LawfulMonad n]
     lift (x *> y) = (lift x : n α) *> (lift y : n β) := by
   simp only [seqRight_eq, lift_map, lift_seq]
 
-def instMonadLiftOfFunction {lift : ⦃α : Type u⦄ -> m α → n α} :
-    MonadLift m n where
-  monadLift := lift (α := _)
+abbrev idToMonad [Monad m] ⦃α : Type u⦄ (x : Id α) : m α :=
+    pure x.run
+
+def LawfulMonadLiftFunction.idToMonad [Monad m] [LawfulMonad m] :
+    LawfulMonadLiftFunction (m := Id) (n := m) idToMonad where
+  lift_pure := by simp [Internal.idToMonad]
+  lift_bind := by simp [Internal.idToMonad]
 
 instance [LawfulMonadLiftFunction lift] :
     letI : MonadLift m n := ⟨lift (α := _)⟩

src/Std/Data/Iterators.lean

@@ -97,6 +97,8 @@ All of the following module names are prefixed with `Std.Data.Iterators`.
 
 `Lemmas` provides the means to verify programs that use iterators.
 
+In particular, `Lemmas.Equivalence` develops the theory of equivalences of iterators.
+
 ### Implementation details
 
 `Internal` contains code that should not be relied upon because it may change in the future.

src/Std/Data/Iterators/Basic.lean

@@ -6,6 +6,7 @@ Authors: Paul Reichert
 prelude
 import Init.Core
 import Init.Classical
+import Init.Ext
 import Init.NotationExtra
 import Init.TacticsExtra
 
@@ -58,6 +59,7 @@ def x := [1, 2, 3].iterM IO
 def x := ([1, 2, 3].iterM IO : IterM IO Nat)
 ```
 -/
+@[ext]
 structure IterM {α : Type w} (m : Type w → Type w') (β : Type w) where
   /-- Internal implementation detail of the iterator. -/
   internalState : α
@@ -204,6 +206,12 @@ theorem IterStep.mapIterator_mapIterator {α' : Type u'} {α'' : Type u''}
     (step.mapIterator f).mapIterator g = step.mapIterator (g ∘ f) := by
   cases step <;> rfl
 
+theorem IterStep.mapIterator_comp {α' : Type u'} {α'' : Type u''}
+    {f : α → α'} {g : α' → α''} :
+    IterStep.mapIterator (β := β) (g ∘ f) = mapIterator g ∘ mapIterator f := by
+  apply funext
+  exact fun _ => mapIterator_mapIterator.symm
+
 @[simp]
 theorem IterStep.mapIterator_id {step : IterStep α β} :
     step.mapIterator id = step := by

src/Std/Data/Iterators/Combinators/FilterMap.lean

@@ -51,12 +51,12 @@ it                                ---a --b--c --d-e--⊥
 it.filterMapWithPostcondition     ---a'-----c'-------⊥
 ```
 
-(given that `f a = pure (some a)'`, `f c = pure (some c')` and `f b = f d = d e = pure none`)
+(given that `f a = pure (some a')`, `f c = pure (some c')` and `f b = f d = d e = pure none`)
 
 **Termination properties:**
 
 * `Finite` instance: only if `it` is finite
-* `Productive` instance: only if `it` is finite`
+* `Productive` instance: only if `it` is finite
 
 For certain mapping functions `f`, the resulting iterator will be finite (or productive) even though
 no `Finite` (or `Productive`) instance is provided. For example, if `f` never returns `none`, then

src/Std/Data/Iterators/Lemmas.lean

@@ -9,3 +9,4 @@ import Std.Data.Iterators.Lemmas.Monadic
 import Std.Data.Iterators.Lemmas.Consumers
 import Std.Data.Iterators.Lemmas.Combinators
 import Std.Data.Iterators.Lemmas.Producers
+import Std.Data.Iterators.Lemmas.Equivalence

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

@@ -181,7 +181,7 @@ theorem Iter.step_mapM {f : β → n γ}
     match it.step with
     | .yield it' out h => do
       let out' ← f out
-      pure <| .yield (it'.mapM f) out' (.yieldSome h ⟨_, rfl⟩)
+      pure <| .yield (it'.mapM f) out' (.yieldSome h ⟨⟨out', True.intro⟩, rfl⟩)
     | .skip it' h =>
       pure <| .skip (it'.mapM f) (.skip h)
     | .done h =>
@@ -220,7 +220,7 @@ theorem Iter.step_filterMap {f : β → Option γ} :
   · simp
   · simp
 
-def Iter.step_map {f : β → γ} :
+theorem Iter.step_map {f : β → γ} :
     (it.map f).step = match it.step with
       | .yield it' out h =>
         .yield (it'.map f) (f out) (.yieldSome (out := out) h ⟨⟨f out, rfl⟩, rfl⟩)

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

@@ -4,9 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Paul Reichert
 -/
 prelude
+import Std.Data.Internal.LawfulMonadLiftFunction
 import Std.Data.Iterators.Combinators.Monadic.FilterMap
 import Std.Data.Iterators.Lemmas.Consumers.Monadic
-import Std.Data.Internal.LawfulMonadLiftFunction
+import Std.Data.Iterators.Lemmas.Equivalence.StepCongr
 
 namespace Std.Iterators
 open Std.Internal
@@ -145,7 +146,7 @@ theorem IterM.step_mapM {γ : Type w} {f : β → n γ}
     match ← it.step with
     | .yield it' out h => do
       let out' ← f out
-      pure <| .yield (it'.mapM f) out' (.yieldSome h ⟨_, rfl⟩)
+      pure <| .yield (it'.mapM f) out' (.yieldSome h ⟨⟨out', True.intro⟩, rfl⟩)
     | .skip it' h =>
       pure <| .skip (it'.mapM f) (.skip h)
     | .done h =>
@@ -157,6 +158,8 @@ theorem IterM.step_mapM {γ : Type w} {f : β → n γ}
     simp only [PostconditionT.operation_map, bind_map_left, bind_pure_comp]
     simp only [PostconditionT.lift, Functor.map, Functor.map_map,
       bind_map_left, bind_pure_comp]
+    simp only [PostconditionT.operation_map, Functor.map_map, PlausibleIterStep.skip,
+      PlausibleIterStep.yield, bind_map_left, bind_pure_comp]
     rfl
   | .skip it' h => rfl
   | .done h => rfl
@@ -178,7 +181,7 @@ theorem IterM.step_filterMap [Monad m] [LawfulMonad m] {f : β → Option β'} :
   apply bind_congr
   intro step
   split
-  · simp only [pure_bind]
+  · simp only [PostconditionT.pure, PlausibleIterStep.skip, PlausibleIterStep.yield, pure_bind]
     split <;> split <;> simp_all
   · simp
   · simp
@@ -229,7 +232,6 @@ end Step
 
 section Lawful
 
-
 instance {α β γ : Type w} {m : Type w → Type w'} {n : Type w → Type w''} {o : Type w → Type x}
     [Monad m] [Monad n] [Monad o] [LawfulMonad n] [LawfulMonad o] [Iterator α m β] [Finite α m]
     [IteratorCollect α m o] [LawfulIteratorCollect α m o]
@@ -408,4 +410,195 @@ theorem IterM.toArray_filter {α : Type w} {m : Type w → Type w'} [Monad m] [L
 
 end ToArray
 
+section Equivalence
+
+theorem stepAsHetT_filterMapWithPostcondition [Monad m] [LawfulMonad m] [Monad n]
+    [LawfulMonad n] [Iterator α m β] [MonadLiftT m n] [LawfulMonadLiftT m n]
+    {f : β → PostconditionT n (Option γ)} {it : IterM (α := α) m β} :
+    (IterM.stepAsHetT (it.filterMapWithPostcondition f)) =
+      (((IterM.stepAsHetT it).liftInner n : HetT n _)).bind (match · with
+        | .yield it' out => do match ← HetT.ofPostconditionT (f out) with
+          | some out' => return .yield (it'.filterMapWithPostcondition f) out'
+          | none => return .skip (it'.filterMapWithPostcondition f)
+        | .skip it' => return .skip (it'.filterMapWithPostcondition f)
+        | .done => return .done) := by
+  simp only [HetT.ext_iff, Equivalence.property_step, Equivalence.prun_step, HetT.prun_bind,
+    HetT.property_liftInner, Equivalence.prun_liftInner_step, HetT.property_bind]
+  refine ⟨?_, ?_⟩
+  · ext step
+    constructor
+    · intro h
+      cases h
+      case yieldNone it' out h h' =>
+        refine ⟨_, h, ?_⟩
+        simp only [bind, HetT.property_bind, HetT.property_ofPostconditionT]
+        exact ⟨none, by simp [Pure.pure]; exact ⟨h', rfl⟩⟩
+      case yieldSome it' out out' h h' =>
+        refine ⟨_, h, ?_⟩
+        simp only [bind, HetT.property_bind, HetT.property_ofPostconditionT]
+        exact ⟨some out', by simp [Pure.pure]; exact ⟨h', rfl⟩⟩
+      case skip it' h =>
+        exact ⟨_, h, by simp [Pure.pure]; rfl⟩
+      case done h =>
+        exact ⟨_, h, by simp [Pure.pure]⟩
+    · rintro ⟨step', h, h'⟩
+      cases step'
+      case yield it' out =>
+        simp only [bind, HetT.property_bind, HetT.property_ofPostconditionT] at h'
+        rcases h' with ⟨out', h'⟩
+        cases out'
+        · simp only [pure, HetT.property_pure] at h'
+          cases h'.2
+          exact .yieldNone h h'.1
+        · simp only [pure, HetT.property_pure] at h'
+          cases h'.2
+          exact .yieldSome h h'.1
+      case skip it' =>
+        simp only [pure, HetT.property_pure] at h'
+        cases h'
+        exact .skip h
+      case done =>
+        simp only [pure, HetT.property_pure] at h'
+        cases h'
+        exact .done h
+  · intro β f
+    simp only [IterM.step_filterMapWithPostcondition, PlausibleIterStep.skip,
+      PlausibleIterStep.yield, PlausibleIterStep.done, bind_assoc]
+    apply bind_congr
+    intro step
+    cases step using PlausibleIterStep.casesOn
+    · simp only [bind_assoc, bind, HetT.prun_bind, HetT.property_ofPostconditionT,
+        HetT.prun_ofPostconditionT]
+      apply bind_congr
+      rintro ⟨out, _⟩
+      cases out <;> simp [pure]
+    · simp [pure]
+    · simp [pure]
+
+theorem IterM.Equiv.filterMapWithPostcondition {α₁ α₂ β γ : Type w}
+    {m : Type w → Type w'} {n : Type w → Type w''}
+    [Monad m] [LawfulMonad m] [Monad n] [LawfulMonad n] [Iterator α₁ m β] [Iterator α₂ m β]
+    [MonadLiftT m n] [LawfulMonadLiftT m n]
+    {f : β → PostconditionT n (Option γ)} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.filterMapWithPostcondition f) (itb.filterMapWithPostcondition f) := by
+  rw [IterM.Equiv]
+  refine BundledIterM.Equiv.fixpoint_induct n γ ?R ?implies (.ofIterM _) (.ofIterM _) ?hR
+  case R =>
+    intro ita' itb'
+    exact ∃ (ita : IterM (α := α₁) m β) (itb : IterM (α := α₂) m β),
+      ita' = .ofIterM (ita.filterMapWithPostcondition f) ∧
+      itb' = .ofIterM (itb.filterMapWithPostcondition f) ∧
+      IterM.Equiv ita itb
+  case hR =>
+    exact ⟨_, _, rfl, rfl, h⟩
+  case implies =>
+    rintro _ _ ⟨ita, itb, rfl, rfl, h'⟩
+    replace h := h'
+    simp only [BundledIterM.step, BundledIterM.iterator_ofIterM, HetT.map_eq_pure_bind,
+      HetT.bind_assoc, Function.comp_apply, HetT.pure_bind, IterStep.mapIterator_mapIterator]
+    rw [stepAsHetT_filterMapWithPostcondition, stepAsHetT_filterMapWithPostcondition]
+    simp only [HetT.bind_assoc]
+    simp only [Equiv, BundledIterM.Equiv, BundledIterM.step, BundledIterM.iterator_ofIterM,
+      HetT.map_eq_pure_bind, HetT.bind_assoc, Function.comp_apply, HetT.pure_bind,
+      IterStep.mapIterator_mapIterator] at h'
+    apply liftInner_stepAsHetT_bind_congr h
+    intro sa hsa sb hsb hs
+    simp only [IterStep.bundledQuotient] at hs
+    cases sa <;> cases sb <;> (try exfalso; simp_all; done)
+    case yield =>
+      simp only [IterStep.mapIterator_yield, Function.comp_apply, IterStep.yield.injEq,
+        BundledIterM.Equiv.quotMk_eq_iff] at hs
+      cases hs.2
+      simp only [bind, HetT.bind_assoc]
+      congr
+      ext out
+      cases out
+      all_goals
+        simp only [pure, HetT.pure_bind, IterStep.mapIterator_skip, IterStep.mapIterator_yield,
+          Function.comp_apply]
+        congr 2
+        apply Quot.sound
+        exact ⟨_, _, rfl, rfl, hs.1⟩
+    case skip =>
+      simp only [IterStep.mapIterator_skip, Function.comp_apply, IterStep.skip.injEq,
+        BundledIterM.Equiv.quotMk_eq_iff] at hs
+      simp only [pure, HetT.pure_bind, IterStep.mapIterator_skip, Function.comp_apply]
+      congr 2
+      apply Quot.sound
+      exact ⟨_, _, rfl, rfl, hs⟩
+    case done =>
+      simp [Pure.pure]
+
+theorem IterM.Equiv.filterWithPostcondition {α₁ α₂ β : Type w}
+    {m : Type w → Type w'} {n : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [Monad n] [LawfulMonad n] [Iterator α₁ m β] [Iterator α₂ m β]
+    [MonadLiftT m n] [LawfulMonadLiftT m n]
+    {f : β → PostconditionT n (ULift Bool)} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.filterWithPostcondition f) (itb.filterWithPostcondition f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+theorem IterM.Equiv.mapWithPostcondition {α₁ α₂ β γ : Type w}
+    {m : Type w → Type w'} {n : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [Monad n] [LawfulMonad n] [Iterator α₁ m β] [Iterator α₂ m β]
+    [MonadLiftT m n] [LawfulMonadLiftT m n]
+    {f : β → PostconditionT n γ} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.mapWithPostcondition f) (itb.mapWithPostcondition f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+theorem IterM.Equiv.filterMapM {α₁ α₂ β γ : Type w}
+    {m : Type w → Type w'} {n : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [Monad n] [LawfulMonad n] [Iterator α₁ m β] [Iterator α₂ m β]
+    [MonadLiftT m n] [LawfulMonadLiftT m n]
+    {f : β → n (Option γ)} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.filterMapM f) (itb.filterMapM f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+theorem IterM.Equiv.filterM {α₁ α₂ β : Type w}
+    {m : Type w → Type w'} {n : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [Monad n] [LawfulMonad n] [Iterator α₁ m β] [Iterator α₂ m β]
+    [MonadLiftT m n] [LawfulMonadLiftT m n]
+    {f : β → n (ULift Bool)} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.filterM f) (itb.filterM f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+theorem IterM.Equiv.mapM {α₁ α₂ β γ : Type w}
+    {m : Type w → Type w'} {n : Type w → Type w''} [Monad m] [LawfulMonad m]
+    [Monad n] [LawfulMonad n] [Iterator α₁ m β] [Iterator α₂ m β]
+    [MonadLiftT m n] [LawfulMonadLiftT m n]
+    {f : β → n γ} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.mapM f) (itb.mapM f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+theorem IterM.Equiv.filterMap {α₁ α₂ β γ : Type w}
+    {m : Type w → Type w'} [Monad m] [LawfulMonad m]
+    [Iterator α₁ m β] [Iterator α₂ m β]
+    {f : β → Option γ} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.filterMap f) (itb.filterMap f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+theorem IterM.Equiv.filter {α₁ α₂ β : Type w}
+    {m : Type w → Type w'} [Monad m] [LawfulMonad m]
+    [Iterator α₁ m β] [Iterator α₂ m β]
+    {f : β → Bool} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.filter f) (itb.filter f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+theorem IterM.Equiv.map {α₁ α₂ β γ : Type w}
+    {m : Type w → Type w'} [Monad m] [LawfulMonad m]
+    [Iterator α₁ m β] [Iterator α₂ m β]
+    {f : β → γ} {ita : IterM (α := α₁) m β} {itb : IterM (α := α₂) m β}
+    (h : IterM.Equiv ita itb) :
+    IterM.Equiv (ita.map f) (itb.map f) :=
+  IterM.Equiv.filterMapWithPostcondition h
+
+end Equivalence
+
 end Std.Iterators

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

@@ -48,6 +48,13 @@ theorem Iter.toArray_toList {α β} [Iterator α Id β] [Finite α Id] [Iterator
     it.toList.toArray = it.toArray := by
   simp [toArray_eq_toArray_toIterM, toList_eq_toList_toIterM, ← IterM.toArray_toList]
 
+@[simp]
+theorem Iter.reverse_toListRev [Iterator α Id β] [Finite α Id]
+    [IteratorCollect α Id Id] [LawfulIteratorCollect α Id Id]
+    {it : Iter (α := α) β} :
+    it.toListRev.reverse = it.toList := by
+  simp [toListRev_eq_toListRev_toIterM, toList_eq_toList_toIterM, ← IterM.reverse_toListRev]
+
 theorem Iter.toListRev_eq {α β} [Iterator α Id β] [Finite α Id] [IteratorCollect α Id Id]
     [LawfulIteratorCollect α Id Id] {it : Iter (α := α) β} :
     it.toListRev = it.toList.reverse := by
@@ -102,4 +109,30 @@ theorem Iter.toList_eq_of_atIdxSlow?_eq {α₁ α₂ β}
     it₁.toList = it₂.toList := by
   ext; simp [getElem?_toList_eq_atIdxSlow?, h]
 
+section Equivalence
+
+theorem Iter.Equiv.toListRev_eq
+    [Iterator α₁ Id β] [Iterator α₂ Id β] [Finite α₁ Id] [Finite α₂ Id]
+    {ita : Iter (α := α₁) β} {itb : Iter (α := α₂) β} (h : Iter.Equiv ita itb) :
+    ita.toListRev = itb.toListRev := by
+  simp [Iter.toListRev_eq_toListRev_toIterM, h.toIterM.toListRev_eq]
+
+theorem Iter.Equiv.toList_eq
+    [Iterator α₁ Id β] [Iterator α₂ Id β] [Finite α₁ Id] [Finite α₂ Id]
+    [IteratorCollect α₁ Id Id] [LawfulIteratorCollect α₁ Id Id]
+    [IteratorCollect α₂ Id Id] [LawfulIteratorCollect α₂ Id Id]
+    {ita : Iter (α := α₁) β} {itb : Iter (α := α₂) β} (h : Iter.Equiv ita itb) :
+    ita.toList = itb.toList := by
+  simp only [← Iter.reverse_toListRev, toListRev_eq h]
+
+theorem Iter.Equiv.toArray_eq
+    [Iterator α₁ Id β] [Iterator α₂ Id β] [Finite α₁ Id] [Finite α₂ Id]
+    [IteratorCollect α₁ Id Id] [LawfulIteratorCollect α₁ Id Id]
+    [IteratorCollect α₂ Id Id] [LawfulIteratorCollect α₂ Id Id]
+    {ita : Iter (α := α₁) β} {itb : Iter (α := α₂) β} (h : Iter.Equiv ita itb) :
+    ita.toArray = itb.toArray := by
+  simp only [← Iter.toArray_toList, toList_eq h]
+
+end Equivalence
+
 end Std.Iterators