/-
Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Paul Reichert
-/
module

prelude
public import Init.NotationExtra
public import Init.WFTactics
import Init.Ext
import Init.PropLemmas

set_option doc.verso true
set_option linter.missingDocs true

public section

/-!
# Definition of iterators

This module defines iterators and what it means for an iterator to be finite and productive.
-/

namespace Std

private opaque Internal.idOpaque {α} : { f : α → α // f = id } := ⟨id, rfl⟩

/--
Currently, {lean}`Shrink α` is just a wrapper around {lean}`α`.

In the future, {name}`Shrink` should allow shrinking {lean}`α` into a potentially smaller universe,
given a proof that {name}`α` is actually small, just like Mathlib's {lit}`Shrink`, except that
the latter's conversion functions are noncomputable. Until then, {lean}`Shrink α` is always in the
same universe as {name}`α`.

This no-op type exists so that fewer breaking changes will be needed when the
real {lit}`Shrink` type is available and the iterators will be made more flexible with regard to
universes.

The conversion functions {name (scope := "Init.Data.Iterators.Basic")}`Shrink.deflate` and
{name (scope := "Init.Data.Iterators.Basic")}`Shrink.inflate` form an equivalence between
{name}`α` and {lean}`Shrink α`, but this equivalence is intentionally not definitional.
-/
public def Shrink (α : Type u) : Type u := Internal.idOpaque.1 α

/-- Converts elements of {name}`α` into elements of {lean}`Shrink α`. -/
@[always_inline]
public def Shrink.deflate {α} (x : α) : Shrink α :=
  cast (by simp [Shrink, Internal.idOpaque.property]) x

/-- Converts elements of {lean}`Shrink α` into elements of {name}`α`. -/
@[always_inline]
public def Shrink.inflate {α} (x : Shrink α) : α :=
  cast (by simp [Shrink, Internal.idOpaque.property]) x

@[simp, grind =]
public theorem Shrink.deflate_inflate {α} {x : Shrink α} :
    Shrink.deflate x.inflate = x := by
  simp [deflate, inflate]

@[simp, grind =]
public theorem Shrink.inflate_deflate {α} {x : α} :
    (Shrink.deflate x).inflate = x := by
  simp [deflate, inflate]

public theorem Shrink.inflate_inj {α} {x y : Shrink α} :
    x.inflate = y.inflate ↔ x = y := by
  apply Iff.intro
  · intro h
    simpa using congrArg Shrink.deflate h
  · rintro rfl
    rfl

public theorem Shrink.deflate_inj {α} {x y : α} :
    Shrink.deflate x = Shrink.deflate y ↔ x = y := by
  apply Iff.intro
  · intro h
    simpa using congrArg Shrink.inflate h
  · rintro rfl
    rfl

-- It is not fruitful to move the following docstrings to verso right now because there are lots of
-- forward references that cannot be realized nicely.
set_option doc.verso false

/--
An iterator that sequentially emits values of type `β` in the monad `m`. It may be finite
or infinite.

See the root module `Std.Data.Iterators` for a more comprehensive overview over the iterator
framework.

See `Std.Data.Iterators.Producers` for ways to iterate over common data structures.
By convention, the monadic iterator associated with an object can be obtained via dot notation.
For example, `List.iterM IO` creates an iterator over a list in the monad `IO`.

See `Init.Data.Iterators.Consumers` for ways to use an iterator. For example, `it.toList` will
convert an iterator `it` into a list and `it.ensureTermination.toList` guarantees that this
operation will terminate, given a proof that the iterator is finite.
It is also always possible to manually iterate using
`it.step`, relying on the termination measures `it.finitelyManySteps` and `it.finitelyManySkips`.

See `Iter` for a more convenient interface in case that no monadic effects are needed (`m = Id`).

Internally, `IterM m β` wraps an element of type `α` containing state information.
The type `α` determines the implementation of the iterator using a typeclass mechanism.
The concrete typeclass implementing the iterator is `Iterator α m β`.

When using combinators, `α` can become very complicated. It is an implicit parameter
of `α` so that the pretty printer will not print this large type by default. If a declaration
returns an iterator, the following will not work:

```lean
def x : IterM IO Nat := [1, 2, 3].iterM IO
```

Instead the declaration type needs to be completely omitted:

```lean
def x := [1, 2, 3].iterM IO

-- if you want to ensure that `x` is an iterator in `IO` emitting `Nat`
def x := ([1, 2, 3].iterM IO : IterM IO Nat)
```
-/
@[ext]
structure IterM {α : Type w} (m : Type w → Type w') (β : Type w) where
  mk' ::
  /-- Internal implementation detail of the iterator. -/
  internalState : α

/--
An iterator that sequentially emits values of type `β`. It may be finite
or infinite.

See the root module `Std.Data.Iterators` for a more comprehensive overview over the iterator
framework.

See `Std.Data.Iterators.Producers` for ways to iterate over common data structures.
By convention, the monadic iterator associated with an object can be obtained via dot notation.
For example, `List.iterM IO` creates an iterator over a list in the monad `IO`.

See `Init.Data.Iterators.Consumers` for ways to use an iterator. For example, `it.toList` will
convert an iterator `it` into a list and `it.ensureTermination.toList` guarantees that this
operation will terminate, given a proof that the iterator is finite.
It is also always possible to manually iterate using
`it.step`, relying on the termination measures `it.finitelyManySteps` and `it.finitelyManySkips`.

See `IterM` for iterators that operate in a monad.

Internally, `Iter β` wraps an element of type `α` containing state information.
The type `α` determines the implementation of the iterator using a typeclass mechanism.
The concrete typeclass implementing the iterator is `Iterator α m β`.

When using combinators, `α` can become very complicated. It is an implicit parameter
of `α` so that the pretty printer will not print this large type by default. If a declaration
returns an iterator, the following will not work:

```lean
def x : Iter Nat := [1, 2, 3].iter
```

Instead the declaration type needs to be completely omitted:

```lean
def x := [1, 2, 3].iter

-- if you want to ensure that `x` is an iterator emitting `Nat`
def x := ([1, 2, 3].iter : Iter Nat)
```
-/
structure Iter {α : Type w} (β : Type w) where
  /-- Internal implementation detail of the iterator. -/
  internalState : α

/--
Converts a pure iterator (`Iter β`) into a monadic iterator (`IterM Id β`) in the
identity monad `Id`.
-/
@[expose]
def Iter.toIterM {α : Type w} {β : Type w} (it : Iter (α := α) β) : IterM (α := α) Id β :=
  ⟨it.internalState⟩

/--
Converts a monadic iterator (`IterM Id β`) over `Id` into a pure iterator (`Iter β`).
-/
@[expose]
def IterM.toIter {α : Type w} {β : Type w} (it : IterM (α := α) Id β) : Iter (α := α) β :=
  ⟨it.internalState⟩

@[simp]
theorem Iter.toIter_toIterM {α : Type w} {β : Type w} (it : Iter (α := α) β) :
    it.toIterM.toIter = it :=
  rfl

@[simp]
theorem Iter.toIter_comp_toIterM {α : Type w} {β : Type w} :
    IterM.toIter ∘ Iter.toIterM (α := α) (β := β) = id :=
  rfl

@[simp]
theorem Iter.toIterM_toIter {α : Type w} {β : Type w} (it : IterM (α := α) Id β) :
    it.toIter.toIterM = it :=
  rfl

@[simp]
theorem Iter.toIterM_comp_toIter {α : Type w} {β : Type w} :
    Iter.toIterM ∘ IterM.toIter (α := α) (β := β) = id :=
  rfl

section IterStep

variable {α : Type u} {β : Type w}

/--
`IterStep α β` represents a step taken by an iterator (`Iter β` or `IterM m β`).
-/
inductive IterStep (α β) where
  /--
  `IterStep.yield it out` describes the situation that an iterator emits `out` and provides `it`
  as the succeeding iterator.
  -/
  | yield : (it : α) → (out : β) → IterStep α β
  /--
  `IterStep.skip it` describes the situation that an iterator does not emit anything in this
  iteration and provides `it'` as the succeeding iterator.

  Allowing `skip` steps is necessary to generate efficient code from a loop over an iterator.
  -/
  | skip : (it : α) → IterStep α β
  /--
  `IterStep.done` describes the situation that an iterator has finished and will neither emit
  more values nor cause any monadic effects. In this case, no succeeding iterator is provided.
  -/
  | done : IterStep α β

/--
Returns the succeeding iterator stored in an iterator step or `none` if the step is `.done`
and the iterator has finished.
-/
@[expose]
def IterStep.successor : IterStep α β → Option α
  | .yield it _ => some it
  | .skip it => some it
  | .done => none

@[simp]
theorem IterStep.successor_yield {it : α} {out : β} :
  (IterStep.yield it out).successor = some it := rfl

@[simp]
theorem IterStep.successor_skip {it : α} : (IterStep.skip (β := β) it).successor = some it := rfl

@[simp]
theorem IterStep.successor_done : (IterStep.done (α := α) (β := β)).successor = none := rfl

/--
If present, applies `f` to the iterator of an `IterStep` and replaces the iterator
with the result of the application of `f`.
-/
@[always_inline, inline, expose]
def IterStep.mapIterator {α' : Type u'} (f : α → α') : IterStep α β → IterStep α' β
  | .yield it out => .yield (f it) out
  | .skip it => .skip (f it)
  | .done => .done

@[simp]
theorem IterStep.mapIterator_yield {α' : Type u'} {f : α → α'} {it : α} {out : β} :
    (IterStep.yield it out).mapIterator f = IterStep.yield (f it) out :=
  rfl

@[simp]
theorem IterStep.mapIterator_skip {α' : Type u'} {f : α → α'} {it : α} :
    (IterStep.skip it (β := β)).mapIterator f = IterStep.skip (f it) :=
  rfl

@[simp]
theorem IterStep.mapIterator_done {α' : Type u'} {f : α → α'} :
    (IterStep.done (α := α) (β := β)).mapIterator f = IterStep.done :=
  rfl

@[simp]
theorem IterStep.mapIterator_mapIterator {α' : Type u'} {α'' : Type u''}
    {f : α → α'} {g : α' → α''} {step : IterStep α β} :
    (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
  cases step <;> rfl

@[simp]
theorem IterStep.mapIterator_id' {step : IterStep α β} :
    step.mapIterator (fun x => x) = step := by
  cases step <;> rfl

/--
A variant of `IterStep` that bundles the step together with a proof that it is "plausible".
The plausibility predicate will later be chosen to assert that a state is a plausible successor
of another state. Having this proof bundled up with the step is important for termination proofs.

See `IterM.Step` and `Iter.Step` for the concrete choice of the plausibility predicate.
-/
@[expose]
def PlausibleIterStep (IsPlausibleStep : IterStep α β → Prop) := Subtype IsPlausibleStep

/--
Match pattern for the `yield` case. See also `IterStep.yield`.
-/
@[match_pattern, simp, spec, expose]
def PlausibleIterStep.yield {IsPlausibleStep : IterStep α β → Prop}
    (it' : α) (out : β) (h : IsPlausibleStep (.yield it' out)) :
    PlausibleIterStep IsPlausibleStep :=
  ⟨.yield it' out, h⟩

/--
Match pattern for the `skip` case. See also `IterStep.skip`.
-/
@[match_pattern, simp, grind =, expose]
def PlausibleIterStep.skip {IsPlausibleStep : IterStep α β → Prop}
    (it' : α) (h : IsPlausibleStep (.skip it')) : PlausibleIterStep IsPlausibleStep :=
  ⟨.skip it', h⟩

/--
Match pattern for the `done` case. See also `IterStep.done`.
-/
@[match_pattern, simp, grind =, expose]
def PlausibleIterStep.done {IsPlausibleStep : IterStep α β → Prop}
    (h : IsPlausibleStep .done) : PlausibleIterStep IsPlausibleStep :=
  ⟨.done, h⟩

/--
A more convenient `cases` eliminator for `PlausibleIterStep`.
-/
@[elab_as_elim, cases_eliminator]
abbrev PlausibleIterStep.casesOn {IsPlausibleStep : IterStep α β → Prop}
    {motive : PlausibleIterStep IsPlausibleStep → Sort x} (s : PlausibleIterStep IsPlausibleStep)
    (yield : ∀ it' out h, motive ⟨.yield it' out, h⟩)
    (skip : ∀ it' h, motive ⟨.skip it', h⟩)
    (done : ∀ h, motive ⟨.done, h⟩) : motive s :=
  match s with
  | .yield it' out h => yield it' out h
  | .skip it' h => skip it' h
  | .done h => done h

end IterStep

/--
The step function of an iterator in `Iter (α := α) β` or `IterM (α := α) m β`.

In order to allow intrinsic termination proofs when iterating with the `step` function, the
step object is bundled with a proof that it is a "plausible" step for the given current iterator.
-/
class Iterator (α : Type w) (m : Type w → Type w') (β : outParam (Type w)) where
  /--
  A relation that governs the allowed steps from a given iterator.

  The "plausible" steps are those which make sense for a given state; plausibility can ensure
  properties such as the successor iterator being drawn from the same collection, that an iterator
  resulting from a skip will return the same next value, or that the next item yielded is next one
  in the original collection.
  -/
  IsPlausibleStep : IterM (α := α) m β → IterStep (IterM (α := α) m β) β → Prop
  /--
  Carries out a step of iteration.
  -/
  step : (it : IterM (α := α) m β) → m (Shrink <| PlausibleIterStep <| IsPlausibleStep it)

section Monadic

/--
Wraps the state of an iterator into an `IterM` object.
-/
@[always_inline, inline, expose]
def IterM.mk {α : Type w} (it : α) (m : Type w → Type w') (β : Type w) :
    IterM (α := α) m β :=
  ⟨it⟩

@[deprecated IterM.mk (since := "2025-12-01"), inline, expose, inherit_doc IterM.mk]
def Iterators.toIterM := @IterM.mk

@[simp]
theorem IterM.mk_internalState {α m β} (it : IterM (α := α) m β) :
    .mk it.internalState m β = it :=
  rfl

set_option linter.missingDocs false in
@[deprecated IterM.mk_internalState (since := "2025-12-01")]
def Iterators.toIterM_internalState := @IterM.mk_internalState

@[simp]
theorem internalState_toIterM {α m β} (it : α) :
    (IterM.mk it m β).internalState = it :=
  rfl

/--
Asserts that certain step is plausibly the successor of a given iterator. What "plausible" means
is up to the `Iterator` instance but it should be strong enough to allow termination proofs.
-/
@[expose]
abbrev IterM.IsPlausibleStep {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β] :
    IterM (α := α) m β → IterStep (IterM (α := α) m β) β → Prop :=
  Iterator.IsPlausibleStep (α := α) (m := m)

/--
The type of the step object returned by `IterM.step`, containing an `IterStep`
and a proof that this is a plausible step for the given iterator.
-/
@[expose]
abbrev IterM.Step {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    (it : IterM (α := α) m β) :=
  PlausibleIterStep it.IsPlausibleStep

/--
Makes a single step with the given iterator `it`, potentially emitting a value and providing a
succeeding iterator. If this function is used recursively, termination can sometimes be proved with
the termination measures `it.finitelyManySteps` and `it.finitelyManySkips`.
-/
@[always_inline, inline, expose]
def IterM.step {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    (it : IterM (α := α) m β) : m (Shrink it.Step) :=
  Iterator.step it

/--
Asserts that a certain output value could plausibly be emitted by the given iterator in its next
step.
-/
@[expose]
def IterM.IsPlausibleOutput {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    (it : IterM (α := α) m β) (out : β) : Prop :=
  ∃ it', it.IsPlausibleStep (.yield it' out)

/--
Asserts that a certain iterator `it'` could plausibly be the directly succeeding iterator of another
given iterator `it`.
-/
@[expose]
def IterM.IsPlausibleSuccessorOf {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    (it' it : IterM (α := α) m β) : Prop :=
  ∃ step, step.successor = some it' ∧ it.IsPlausibleStep step

/--
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`).
-/
@[expose]
def IterM.IsPlausibleSkipSuccessorOf {α : Type w} {m : Type w → Type w'} {β : Type w}
    [Iterator α m β] (it' it : IterM (α := α) m β) : Prop :=
  it.IsPlausibleStep (.skip it')

end Monadic

section Pure

/--
Asserts that certain step is plausibly the successor of a given iterator. What "plausible" means
is up to the `Iterator` instance but it should be strong enough to allow termination proofs.
-/
@[expose]
def Iter.IsPlausibleStep {α : Type w} {β : Type w} [Iterator α Id β]
    (it : Iter (α := α) β) (step : IterStep (Iter (α := α) β) β) : Prop :=
  it.toIterM.IsPlausibleStep (step.mapIterator Iter.toIterM)

/--
Asserts that a certain iterator `it` could plausibly yield the value `out` after an arbitrary
number of steps.
-/
inductive IterM.IsPlausibleIndirectOutput {α β : Type w} {m : Type w → Type w'} [Iterator α m β]
    : IterM (α := α) m β → β → Prop where
  /-- The output value could plausibly be emitted in the next step. -/
  | direct {it : IterM (α := α) m β} {out : β} : it.IsPlausibleOutput out →
      it.IsPlausibleIndirectOutput out
  /-- The output value could plausibly be emitted in a step after the next step. -/
  | indirect {it it' : IterM (α := α) m β} {out : β} : it'.IsPlausibleSuccessorOf it →
      it'.IsPlausibleIndirectOutput out → it.IsPlausibleIndirectOutput out

/--
Asserts that an iterator `it'` could plausibly produce `it'` as a successor iterator after
finitely many steps. This relation is reflexive.
-/
inductive IterM.IsPlausibleIndirectSuccessorOf {α β : Type w} {m : Type w → Type w'}
    [Iterator α m β] : IterM (α := α) m β → IterM (α := α) m β → Prop where
  /-- Every iterator is a plausible indirect successor of itself. -/
  | refl (it : IterM (α := α) m β) : it.IsPlausibleIndirectSuccessorOf it
  /-- The iterator is a plausible successor of one of the current iterator's successors. -/
  | cons_right {it'' it' it : IterM (α := α) m β} (h' : it''.IsPlausibleIndirectSuccessorOf it')
      (h : it'.IsPlausibleSuccessorOf it) : it''.IsPlausibleIndirectSuccessorOf it

theorem IterM.IsPlausibleIndirectSuccessorOf.trans {α β : Type w} {m : Type w → Type w'}
    [Iterator α m β] {it'' it' it : IterM (α := α) m β}
    (h' : it''.IsPlausibleIndirectSuccessorOf it') (h : it'.IsPlausibleIndirectSuccessorOf it) :
    it''.IsPlausibleIndirectSuccessorOf it := by
  induction h
  case refl => exact h'
  case cons_right ih => exact IsPlausibleIndirectSuccessorOf.cons_right ih ‹_›

theorem IterM.IsPlausibleIndirectSuccessorOf.single {α β : Type w} {m : Type w → Type w'}
    [Iterator α m β] {it' it : IterM (α := α) m β}
    (h : it'.IsPlausibleSuccessorOf it) :
    it'.IsPlausibleIndirectSuccessorOf it :=
  .cons_right (.refl _) h

theorem IterM.IsPlausibleIndirectOutput.trans {α β : Type w} {m : Type w → Type w'}
    [Iterator α m β]
    {it' it : IterM (α := α) m β} {out : β} (h : it'.IsPlausibleIndirectSuccessorOf it)
    (h' : it'.IsPlausibleIndirectOutput out) : it.IsPlausibleIndirectOutput out := by
  induction h
  case refl => exact h'
  case cons_right ih => exact IsPlausibleIndirectOutput.indirect ‹_› ih

/--
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.
-/
@[expose]
def Iter.Step {α : Type w} {β : Type w} [Iterator α Id β] (it : Iter (α := α) β) :=
  PlausibleIterStep (Iter.IsPlausibleStep it)

/--
Converts an `Iter.Step` into an `IterM.Step`.
-/
@[always_inline, inline, expose]
def Iter.Step.toMonadic {α : Type w} {β : Type w} [Iterator α Id β] {it : Iter (α := α) β}
    (step : it.Step) : it.toIterM.Step :=
  ⟨step.val.mapIterator Iter.toIterM, step.property⟩

/--
Converts an `IterM.Step` into an `Iter.Step`.
-/
@[always_inline, inline, expose]
def IterM.Step.toPure {α : Type w} {β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
    (step : it.Step) : it.toIter.Step :=
  ⟨step.val.mapIterator IterM.toIter, (by simp [Iter.IsPlausibleStep, step.property])⟩

@[simp]
theorem IterM.Step.toPure_yield {α β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
    {it' out h} : IterM.Step.toPure (⟨.yield it' out, h⟩ : it.Step) = .yield it'.toIter out h :=
  rfl

@[simp]
theorem IterM.Step.toPure_skip {α β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
    {it' h} : IterM.Step.toPure (⟨.skip it', h⟩ : it.Step) = .skip it'.toIter h :=
  rfl

@[simp]
theorem IterM.Step.toPure_done {α β : Type w} [Iterator α Id β] {it : IterM (α := α) Id β}
    {h} : IterM.Step.toPure (⟨.done, h⟩ : it.Step) = .done h :=
  rfl

/--
Asserts that a certain output value could plausibly be emitted by the given iterator in its next
step.
-/
@[expose]
def Iter.IsPlausibleOutput {α : Type w} {β : Type w} [Iterator α Id β]
    (it : Iter (α := α) β) (out : β) : Prop :=
  it.toIterM.IsPlausibleOutput out

theorem Iter.isPlausibleOutput_iff_exists {α : Type w} {β : Type w} [Iterator α Id β]
    {it : Iter (α := α) β} {out : β} :
    it.IsPlausibleOutput out ↔ ∃ it', it.IsPlausibleStep (.yield it' out) := by
  simp only [IsPlausibleOutput, IterM.IsPlausibleOutput]
  constructor
  · rintro ⟨it', h⟩
    exact ⟨it'.toIter, h⟩
  · rintro ⟨it', h⟩
    exact ⟨it'.toIterM, h⟩

/--
Asserts that a certain iterator `it'` could plausibly be the directly succeeding iterator of another
given iterator `it`.
-/
@[expose]
def Iter.IsPlausibleSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β]
    (it' it : Iter (α := α) β) : Prop :=
  it'.toIterM.IsPlausibleSuccessorOf it.toIterM

theorem Iter.isPlausibleSuccessorOf_eq_invImage {α : Type w} {β : Type w} [Iterator α Id β] :
    IsPlausibleSuccessorOf (α := α) (β := β) =
      InvImage (IterM.IsPlausibleSuccessorOf (α := α) (β := β) (m := Id)) Iter.toIterM := rfl

theorem Iter.isPlausibleSuccessorOf_iff_exists {α : Type w} {β : Type w} [Iterator α Id β]
    {it' it : Iter (α := α) β} :
    it'.IsPlausibleSuccessorOf it ↔ ∃ step, step.successor = some it' ∧ it.IsPlausibleStep step := by
  simp only [IsPlausibleSuccessorOf, IterM.IsPlausibleSuccessorOf]
  constructor
  · rintro ⟨step, h₁, h₂⟩
    exact ⟨step.mapIterator IterM.toIter,
      by cases step <;> simp_all [IterStep.successor, Iter.IsPlausibleStep]⟩
  · rintro ⟨step, h₁, h₂⟩
    exact ⟨step.mapIterator Iter.toIterM,
      by cases step <;> simp_all [IterStep.successor, Iter.IsPlausibleStep]⟩

theorem Iter.IsPlausibleStep.isPlausibleSuccessor_of_yield {α : Type w} {β : Type w}
    [Iterator α Id β] {it' it : Iter (α := α) β} {out : β}
    (h : it.IsPlausibleStep (.yield it' out)) : it'.IsPlausibleSuccessorOf it := by
  simpa [isPlausibleSuccessorOf_iff_exists] using ⟨.yield it' out, by simp [h]⟩

theorem Iter.IsPlausibleStep.isPlausibleSuccessor_of_skip {α : Type w} {β : Type w}
    [Iterator α Id β] {it' it : Iter (α := α) β} (h : it.IsPlausibleStep (.skip it')) :
    it'.IsPlausibleSuccessorOf it := by
  simpa [isPlausibleSuccessorOf_iff_exists] using ⟨.skip it', by simp [h]⟩

/--
Asserts that a certain iterator `it` could plausibly yield the value `out` after an arbitrary
number of steps.
-/
inductive Iter.IsPlausibleIndirectOutput {α β : Type w} [Iterator α Id β] :
    Iter (α := α) β → β → Prop where
  /-- The output value could plausibly be emitted in the next step. -/
  | direct {it : Iter (α := α) β} {out : β} : it.IsPlausibleOutput out →
      it.IsPlausibleIndirectOutput out
  /-- The output value could plausibly be emitted in a step after the next step. -/
  | 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 an iterator `it'` could plausibly produce `it'` as a successor iterator after
finitely many steps. This relation is reflexive.
-/
inductive Iter.IsPlausibleIndirectSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β] :
    Iter (α := α) β → Iter (α := α) β → Prop where
  /-- Every iterator is a plausible indirect successor of itself. -/
  | refl (it : Iter (α := α) β) : IsPlausibleIndirectSuccessorOf it it
  /-- The iterator is a plausible indirect successor of one of the current iterator's successors. -/
  | cons_right {it'' it' it : Iter (α := α) β} (h' : it''.IsPlausibleIndirectSuccessorOf it')
      (h : it'.IsPlausibleSuccessorOf it) : it''.IsPlausibleIndirectSuccessorOf it

theorem Iter.isPlausibleIndirectSuccessor_iff_isPlausibleIndirectSuccessor_toIterM {α β : Type w}
    [Iterator α Id β] {it' it : Iter (α := α) β} :
    it'.IsPlausibleIndirectSuccessorOf it ↔ it'.toIterM.IsPlausibleIndirectSuccessorOf it.toIterM := by
  constructor
  · intro h
    induction h with
    | refl => exact .refl _
    | cons_right _ h ih => exact .cons_right ih h
  · intro h
    rw [← Iter.toIter_toIterM (it := it), ← Iter.toIter_toIterM (it := it')]
    generalize it.toIterM = it at ⊢ h
    induction h with
    | refl => exact .refl _
    | cons_right _ h ih => exact .cons_right ih h

theorem Iter.IsPlausibleIndirectSuccessorOf.trans {α : Type w} {β : Type w} [Iterator α Id β]
    {it'' it' it : Iter (α := α) β} (h' : it''.IsPlausibleIndirectSuccessorOf it')
    (h : it'.IsPlausibleIndirectSuccessorOf it) : it''.IsPlausibleIndirectSuccessorOf it := by
  induction h
  case refl => exact h'
  case cons_right ih => exact IsPlausibleIndirectSuccessorOf.cons_right ih ‹_›

theorem Iter.IsPlausibleIndirectOutput.trans {α : Type w} {β : Type w} [Iterator α Id β]
    {it' it : Iter (α := α) β} {out : β} (h : it'.IsPlausibleIndirectSuccessorOf it)
    (h' : it'.IsPlausibleIndirectOutput out) : it.IsPlausibleIndirectOutput out := by
  induction h
  case refl => exact h'
  case cons_right ih => exact IsPlausibleIndirectOutput.indirect ‹_› 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`).
-/
def Iter.IsPlausibleSkipSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β]
    (it' it : Iter (α := α) β) : Prop :=
  it'.toIterM.IsPlausibleSkipSuccessorOf it.toIterM

/--
Makes a single step with the given iterator `it`, potentially emitting a value and providing a
succeeding iterator. If this function is used recursively, termination can sometimes be proved with
the termination measures `it.finitelyManySteps` and `it.finitelyManySkips`.
-/
@[always_inline, inline, expose]
def Iter.step {α β : Type w} [Iterator α Id β] (it : Iter (α := α) β) : it.Step :=
  it.toIterM.step.run.inflate.toPure

end Pure

section Finite

/--
`Finite α m` asserts that `IterM (α := α) m` terminates after finitely many steps. Technically,
this means that the relation of plausible successors is well-founded.
Given this typeclass, termination proofs for well-founded recursion over an iterator `it` can use
`it.finitelyManySteps` as a termination measure.
-/
class Iterators.Finite (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] : Prop where
  /-- The relation of plausible successors is well-founded. -/
  wf : WellFounded (IterM.IsPlausibleSuccessorOf (α := α) (m := m))

theorem Iterators.Finite.wf_of_id {α : Type w} {β : Type w} [Iterator α Id β] [Finite α Id] :
    WellFounded (Iter.IsPlausibleSuccessorOf (α := α)) := by
  simpa [Iter.isPlausibleSuccessorOf_eq_invImage] using InvImage.wf _ Finite.wf

/--
This type is a wrapper around `IterM` so that it becomes a useful termination measure for
recursion over finite iterators. See also `IterM.finitelyManySteps` and `Iter.finitelyManySteps`.
-/
structure IterM.TerminationMeasures.Finite
    (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] where
  /--
  The wrapped iterator.

  In the wrapper, its finiteness is used as a termination measure.
  -/
  it : IterM (α := α) m β

/--
The relation of plausible successors on `IterM.TerminationMeasures.Finite`. It is well-founded
if there is a `Finite` instance.
-/
@[expose]
def IterM.TerminationMeasures.Finite.Rel
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β] :
    TerminationMeasures.Finite α m → TerminationMeasures.Finite α m → Prop :=
  Relation.TransGen <| InvImage IterM.IsPlausibleSuccessorOf IterM.TerminationMeasures.Finite.it

instance IterM.TerminationMeasures.instWellFoundedRelationFinite {α : Type w} {m : Type w → Type w'}
    {β : Type w} [Iterator α m β] [Iterators.Finite α m] :
    WellFoundedRelation (IterM.TerminationMeasures.Finite α m) where
  rel := IterM.TerminationMeasures.Finite.Rel
  wf := by exact (InvImage.wf _ Iterators.Finite.wf).transGen

/--
Termination measure to be used in well-founded recursive functions recursing over a finite iterator
(see also `Finite`).
-/
@[expose]
def IterM.finitelyManySteps {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    [Iterators.Finite α m] (it : IterM (α := α) m β) : IterM.TerminationMeasures.Finite α m :=
  ⟨it⟩

/--
Termination measure to be used in recursive functions built with `WellFounded.extrinsicFix`
recursing over a finite iterator without requiring a proof of finiteness (see also `Finite`).
-/
@[expose]
def IterM.finitelyManySteps! {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    (it : IterM (α := α) m β) : IterM.TerminationMeasures.Finite α m :=
  ⟨it⟩

/--
This theorem is used by a `decreasing_trivial` extension. It powers automatic termination proofs
with `IterM.finitelyManySteps`.
-/
theorem IterM.TerminationMeasures.Finite.rel_of_yield
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    {it it' : IterM (α := α) m β} {out : β} (h : it.IsPlausibleStep (.yield it' out)) :
    Rel ⟨it'⟩ ⟨it⟩ := by
  exact .single ⟨_, rfl, h⟩

@[inherit_doc IterM.TerminationMeasures.Finite.rel_of_yield]
theorem IterM.TerminationMeasures.Finite.rel_of_skip
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    {it it' : IterM (α := α) m β} (h : it.IsPlausibleStep (.skip it')) :
    Rel ⟨it'⟩ ⟨it⟩ := by
  exact .single ⟨_, rfl, h⟩

theorem IterM.TerminationMeasures.Finite.rel_trans
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    {it it' it'' : TerminationMeasures.Finite α m} :
    it.Rel it' → it'.Rel it'' → it.Rel it'' :=
  .trans

macro_rules | `(tactic| decreasing_trivial) => `(tactic|
  first
  | exact IterM.TerminationMeasures.Finite.rel_of_yield ‹_›
  | exact IterM.TerminationMeasures.Finite.rel_of_skip ‹_›
  | fail)

@[inherit_doc IterM.finitelyManySteps, expose]
def Iter.finitelyManySteps {α : Type w} {β : Type w} [Iterator α Id β] [Iterators.Finite α Id]
    (it : Iter (α := α) β) : IterM.TerminationMeasures.Finite α Id :=
  it.toIterM.finitelyManySteps

@[inherit_doc IterM.finitelyManySteps!, expose]
def Iter.finitelyManySteps! {α : Type w} {β : Type w} [Iterator α Id β]
    (it : Iter (α := α) β) : IterM.TerminationMeasures.Finite α Id :=
  it.toIterM.finitelyManySteps!

/--
This theorem is used by a `decreasing_trivial` extension. It powers automatic termination proofs
with `IterM.finitelyManySteps`.
-/
theorem Iter.TerminationMeasures.Finite.rel_of_yield
    {α : Type w} {β : Type w} [Iterator α Id β]
    {it it' : Iter (α := α) β} {out : β} (h : it.IsPlausibleStep (.yield it' out)) :
    IterM.TerminationMeasures.Finite.Rel ⟨it'.toIterM⟩ ⟨it.toIterM⟩ :=
  IterM.TerminationMeasures.Finite.rel_of_yield h

@[inherit_doc Iter.TerminationMeasures.Finite.rel_of_yield]
theorem Iter.TerminationMeasures.Finite.rel_of_skip
    {α : Type w} {β : Type w} [Iterator α Id β]
    {it it' : Iter (α := α) β} (h : it.IsPlausibleStep (.skip it')) :
    IterM.TerminationMeasures.Finite.Rel ⟨it'.toIterM⟩ ⟨it.toIterM⟩ :=
  IterM.TerminationMeasures.Finite.rel_of_skip h

macro_rules | `(tactic| decreasing_trivial) => `(tactic|
  first
  | exact Iter.TerminationMeasures.Finite.rel_of_yield ‹_›
  | exact Iter.TerminationMeasures.Finite.rel_of_skip ‹_›
  | fail)

theorem IterM.isPlausibleSuccessorOf_of_yield
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    {it' it : IterM (α := α) m β} {out : β} (h : it.IsPlausibleStep (.yield it' out)) :
    it'.IsPlausibleSuccessorOf it :=
  ⟨_, rfl, h⟩

theorem IterM.isPlausibleSuccessorOf_of_skip
    {α m β} [Iterator α m β] {it it' : IterM (α := α) m β}
    (h : it.IsPlausibleStep (.skip it')) :
    it'.IsPlausibleSuccessorOf it :=
  ⟨_, rfl, h⟩

end Finite

section Productive

/--
`Productive α m` asserts that `IterM (α := α) m` terminates or emits a value after finitely many
skips. Technically, this means that the relation of plausible successors during skips is
well-founded.
Given this typeclass, termination proofs for well-founded recursion over an iterator `it` can use
`it.finitelyManySkips` as a termination measure.
-/
class Iterators.Productive (α m) {β} [Iterator α m β] : Prop where
  /-- The relation of plausible successors during skips is well-founded. -/
  wf : WellFounded (IterM.IsPlausibleSkipSuccessorOf (α := α) (m := m))

/--
This type is a wrapper around `IterM` so that it becomes a useful termination measure for
recursion over productive iterators. See also `IterM.finitelyManySkips` and `Iter.finitelyManySkips`.
-/
structure IterM.TerminationMeasures.Productive
    (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] where
  /--
  The wrapped iterator.

  In the wrapper, its productivity is used as a termination measure.
  -/
  it : IterM (α := α) m β

/--
The relation of plausible successors while skipping on `IterM.TerminationMeasures.Productive`.
It is well-founded if there is a `Productive` instance.
-/
@[expose]
def IterM.TerminationMeasures.Productive.Rel
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β] :
    TerminationMeasures.Productive α m → TerminationMeasures.Productive α m → Prop :=
  Relation.TransGen <| InvImage IterM.IsPlausibleSkipSuccessorOf IterM.TerminationMeasures.Productive.it

theorem IterM.TerminationMeasures.Finite.Rel.of_productive
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β] {a b : Finite α m} :
    Productive.Rel ⟨a.it⟩ ⟨b.it⟩ → Finite.Rel a b := by
  generalize ha' : Productive.mk a.it = a'
  generalize hb' : Productive.mk b.it = b'
  have ha : a = ⟨a'.it⟩ := by simp [← ha']
  have hb : b = ⟨b'.it⟩ := by simp [← hb']
  rw [ha, hb]
  clear ha hb ha' hb' a b
  rw [Productive.Rel, Finite.Rel]
  intro h
  induction h
  · rename_i ih
    exact .single ⟨_, rfl, ih⟩
  · rename_i hab ih
    refine .trans ih ?_
    exact .single ⟨_, rfl, hab⟩

instance IterM.TerminationMeasures.instWellFoundedRelationProductive {α : Type w}
    {m : Type w → Type w'} {β : Type w} [Iterator α m β] [Iterators.Productive α m] :
    WellFoundedRelation (IterM.TerminationMeasures.Productive α m) where
  rel := IterM.TerminationMeasures.Productive.Rel
  wf := by exact (InvImage.wf _ Iterators.Productive.wf).transGen

/--
Termination measure to be used in well-founded recursive functions recursing over a productive
iterator (see also `Productive`).
-/
@[expose]
def IterM.finitelyManySkips {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    [Iterators.Productive α m] (it : IterM (α := α) m β) : IterM.TerminationMeasures.Productive α m :=
  ⟨it⟩

/--
Termination measure to be used in recursive functions built with `WellFounded.extrinsicFix`
recursing over a productive iterator without requiring a proof of productiveness
(see also `Productive`).
-/
@[expose]
def IterM.finitelyManySkips! {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    (it : IterM (α := α) m β) : IterM.TerminationMeasures.Productive α m :=
  ⟨it⟩

/--
This theorem is used by a `decreasing_trivial` extension. It powers automatic termination proofs
with `IterM.finitelyManySkips`.
-/
theorem IterM.TerminationMeasures.Productive.rel_of_skip
    {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
    {it it' : IterM (α := α) m β} (h : it.IsPlausibleStep (.skip it')) :
    Rel ⟨it'⟩ ⟨it⟩ :=
  .single h

macro_rules | `(tactic| decreasing_trivial) => `(tactic|
  first
    | exact IterM.TerminationMeasures.Productive.rel_of_skip ‹_›
    | fail)

@[inherit_doc IterM.finitelyManySkips, expose]
def Iter.finitelyManySkips {α : Type w} {β : Type w} [Iterator α Id β] [Iterators.Productive α Id]
    (it : Iter (α := α) β) : IterM.TerminationMeasures.Productive α Id :=
  it.toIterM.finitelyManySkips

@[inherit_doc IterM.finitelyManySkips!, expose]
def Iter.finitelyManySkips! {α : Type w} {β : Type w} [Iterator α Id β]
    (it : Iter (α := α) β) : IterM.TerminationMeasures.Productive α Id :=
  it.toIterM.finitelyManySkips!

/--
This theorem is used by a `decreasing_trivial` extension. It powers automatic termination proofs
with `Iter.finitelyManySkips`.
-/
theorem Iter.TerminationMeasures.Productive.rel_of_skip
    {α : Type w} {β : Type w} [Iterator α Id β]
    {it it' : Iter (α := α) β} (h : it.IsPlausibleStep (.skip it')) :
    IterM.TerminationMeasures.Productive.Rel ⟨it'.toIterM⟩ ⟨it.toIterM⟩ :=
  IterM.TerminationMeasures.Productive.rel_of_skip h

macro_rules | `(tactic| decreasing_trivial) => `(tactic|
  first
    | exact Iter.TerminationMeasures.Productive.rel_of_skip ‹_›
    | fail)

instance Iterators.instProductiveOfFinte [Iterator α m β] [Iterators.Finite α m] :
    Iterators.Productive α m where
  wf := by
    apply Subrelation.wf (r := IterM.IsPlausibleSuccessorOf)
    · intro it' it h
      exact IterM.isPlausibleSuccessorOf_of_skip h
    · exact Iterators.Finite.wf

end Productive

/--
This typeclass characterizes iterators that have deterministic return values. This typeclass does
*not* guarantee that there are no monadic side effects such as exceptions.

General monadic iterators can be nondeterministic, so that `it.IsPlausibleStep step` will be true
for no or more than one choice of `step`. This typeclass ensures that there is exactly one such
choice.

This is an experimental instance and it should not be explicitly used downstream of the standard
library.
-/
class LawfulDeterministicIterator (α : Type w) (m : Type w → Type w') [Iterator α m β]
    where
  /--
  Every iterator with state `α` in monad `m` has exactly one plausible step.
  -/
  isPlausibleStep_eq_eq : ∀ it : IterM (α := α) m β, ∃ step, it.IsPlausibleStep = (· = step)

namespace Iterators

/--
This structure provides a more convenient way to define `Finite α m` instances using
`Finite.of_finitenessRelation : FinitenessRelation α m → Finite α m`.
-/
structure FinitenessRelation (α : Type w) (m : Type w → Type w') {β : Type w}
    [Iterator α m β] where
  /--
  A well-founded relation such that if `it'` is a successor iterator of `it`, then `Rel it' it`.
  -/
  Rel (it' it : IterM (α := α) m β) : Prop
  /-- `Rel` is well-founded. -/
  wf : WellFounded Rel
  /-- If `it'` is a successor iterator of `it`, then `Rel it' it`. -/
  subrelation : ∀ {it it'}, it'.IsPlausibleSuccessorOf it → Rel it' it

theorem Finite.of_finitenessRelation
    {α : Type w} {m : Type w → Type w'} {β : Type w}
    [Iterator α m β] (r : FinitenessRelation α m) : Finite α m where
  wf := by
    refine Subrelation.wf (r := r.Rel) ?_ ?_
    · intro x y h
      apply FinitenessRelation.subrelation
      exact h
    · apply InvImage.wf
      exact r.wf

/--
This structure provides a more convenient way to define `Productive α m` instances using
`Productive.of_productivenessRelation : ProductivenessRelation α m → Productive α m`.
-/
structure ProductivenessRelation (α : Type w) (m : Type w → Type w') {β : Type w}
    [Iterator α m β] where
  /--
  A well-founded relation such that if `it'` is obtained from `it` by skipping, then `Rel it' it`.
  -/
  Rel : (IterM (α := α) m β) → (IterM (α := α) m β) → Prop
  /-- `Rel` is well-founded. -/
  wf : WellFounded Rel
  /-- If `it'` is obtained from `it` by skipping, then `Rel it' it`. -/
  subrelation : ∀ {it it'}, it'.IsPlausibleSkipSuccessorOf it → Rel it' it

theorem Productive.of_productivenessRelation
    {α : Type w} {m : Type w → Type w'} {β : Type w}
    [Iterator α m β] (r : ProductivenessRelation α m) : Productive α m where
  wf := by
    refine Subrelation.wf (r := r.Rel) ?_ ?_
    · intro x y h
      apply ProductivenessRelation.subrelation
      exact h
    · apply InvImage.wf
      exact r.wf

end Std.Iterators