docstrings, move option WF lemmas

datokrat · datokrat · commit 0a98a2c67fa6 · 2025-05-26T14:08:26.000+02:00
diff --git a/src/Init/Data/Option/Basic.lean b/src/Init/Data/Option/Basic.lean
@@ -256,6 +256,31 @@ instance (r : α → β → Prop) [s : DecidableRel r] : DecidableRel (Option.lt
   | some _, none   => isFalse not_false
   | none,   none   => isFalse not_false
 
+namespace SomeLtNone
+
+/--
+Lifts an ordering relation to `Option`, such that `none` is the greatest element.
+
+It can be understood as adding a distinguished greatest element, represented by `none`, to both `α`
+and `β`.
+
+Caution: Given `LT α`, `Option.SomeLtNone.lt LT.lt` differs from the `LT (Option α)` instance,
+which is implemented by `Option.lt Lt.lt`.
+
+Examples:
+ * `Option.lt (fun n k : Nat => n < k) none none = False`
+ * `Option.lt (fun n k : Nat => n < k) none (some 3) = False`
+ * `Option.lt (fun n k : Nat => n < k) (some 3) none = True`
+ * `Option.lt (fun n k : Nat => n < k) (some 4) (some 5) = True`
+ * `Option.lt (fun n k : Nat => n < k) (some 4) (some 4) = False`
+-/
+def lt {α} (r : α → α → Prop) : Option α → Option α → Prop
+  | none, _ => false
+  | some _, none => true
+  | some a', some a => r a' a
+
+end SomeLtNone
+
 /--
 Applies a function to a two optional values if both are present. Otherwise, if one value is present,
 it is returned and the function is not used.
diff --git a/src/Init/Data/Option/Lemmas.lean b/src/Init/Data/Option/Lemmas.lean
@@ -1902,4 +1902,40 @@ theorem map_min [Min α] [Min β] {o o' : Option α} {f : α → β} (hf : ∀ x
     (min o o').map f = min (o.map f) (o'.map f) := by
   cases o <;> cases o' <;> simp [*]
 
+theorem wellFounded_lt {α} {rel : α → α → Prop} (h : WellFounded rel) :
+    WellFounded (Option.lt rel) := by
+  refine ⟨fun x => ?_⟩
+  have hn : Acc (Option.lt rel) none := by
+    refine Acc.intro none ?_
+    intro y hyx
+    cases y <;> cases hyx
+  cases x
+  · exact hn
+  · rename_i x
+    induction h.apply x
+    rename_i x' h ih
+    refine Acc.intro _ ?_
+    intro y hyx'
+    cases y
+    · exact hn
+    · exact ih _ hyx'
+
+theorem SomeLtNone.wellFounded_lt {α} {r : α → α → Prop} (h : WellFounded r) :
+    WellFounded (SomeLtNone.lt r) := by
+  refine ⟨?_⟩
+  intro x
+  constructor
+  intro x' hlt
+  match x' with
+  | none => contradiction
+  | some x' =>
+    clear hlt
+    induction h.apply x'
+    rename_i ih
+    constructor
+    intro x'' hlt'
+    match x'' with
+    | none => contradiction
+    | some x'' => exact ih x'' hlt'
+
 end Option
diff --git a/src/Std/Data/Iterators/Combinators/Monadic/Zip.lean b/src/Std/Data/Iterators/Combinators/Monadic/Zip.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Paul Reichert
 -/
 prelude
+import Init.Data.Option.Lemmas
 import Std.Data.Iterators.Basic
 import Std.Data.Iterators.Consumers.Collect
 import Std.Data.Iterators.Consumers.Loop
@@ -58,32 +59,42 @@ instance Zip.instIterator [Monad m] :
       | .done hp =>
           pure <| .done (.doneRight hm hp)
 
-@[inline]
+/--
+Given two iterators `left` and `right`, `left.zip right` is an iterator that yields pairs of
+outputs of `left` and `right`. When one of them terminates,
+the `zip` iterator will also terminate.
+
+**Marble diagram:**
+
+```text
+left               --a        ---b        --c
+right                 --x         --y        --⊥
+left.zip right     -----(a, x)------(b, y)-----⊥
+```
+
+**Termination properties:**
+
+* `Finite` instance: only if either `left` or `right` is finite and the other is productive
+* `Productive` instance: only if `left` and `right` are productive
+
+There are situations where `left.zip right` is finite (or productive) but none of the instances
+above applies. For example, if the computation happens in an `Except` monad and `left` immediately
+fails when calling `step`, then `left.zip right` will also do so. In such a case, the `Finite`
+(or `Productive`) instance needs to be proved manually.
+
+**Performance:**
+
+This combinator incurs an additional O(1) cost with each step taken by `left` or `right`.
+
+Right now, the compiler does not unbox the internal state, leading to worse performance than
+possible.
+-/
+@[always_inline, inline]
 def IterM.zip
     (left : IterM (α := α₁) m β₁) (right : IterM (α := α₂) m β₂) :
     IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) :=
   toIterM ⟨left, none, right⟩ m _
 
--- TODO: put this into core. This is also duplicated in FlatMap
-theorem Zip.wellFounded_optionLt {α} {rel : α → α → Prop} (h : WellFounded rel) :
-    WellFounded (Option.lt rel) := by
-  refine ⟨?_⟩
-  intro x
-  have hn : Acc (Option.lt rel) none := by
-    refine Acc.intro none ?_
-    intro y hyx
-    cases y <;> cases hyx
-  cases x
-  · exact hn
-  · rename_i x
-    induction h.apply x
-    rename_i x' h ih
-    refine Acc.intro _ ?_
-    intro y hyx'
-    cases y
-    · exact hn
-    · exact ih _ hyx'
-
 variable (m) in
 def Zip.Rel₁ [Finite α₁ m] [Productive α₂ m] :
     IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) → IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) → Prop :=
@@ -119,7 +130,7 @@ def Zip.instFinitenessRelation₁ [Monad m] [Finite α₁ m] [Productive α₂ m
     refine ⟨fun (a, b) => Prod.lexAccessible (WellFounded.apply ?_ a) (WellFounded.apply ?_) b⟩
     · exact WellFoundedRelation.wf
     · refine ⟨fun (a, b) => Prod.lexAccessible (WellFounded.apply ?_ a) (WellFounded.apply ?_) b⟩
-      · apply Zip.wellFounded_optionLt
+      · apply Option.wellFounded_lt
         exact emptyWf.wf
       · exact WellFoundedRelation.wf
   subrelation {it it'} h := by
@@ -151,35 +162,12 @@ instance Zip.instFinite₁ [Monad m] [Finite α₁ m] [Productive α₂ m] :
     Finite (Zip α₁ m α₂ β₂) m :=
   Finite.of_finitenessRelation Zip.instFinitenessRelation₁
 
-def Zip.lt_with_top {α} (r : α → α → Prop) : Option α → Option α → Prop
-  | none, _ => false
-  | some _, none => true
-  | some a', some a => r a' a
-
-theorem Zip.wellFounded_lt_with_top {α} {r : α → α → Prop} (h : WellFounded r) :
-    WellFounded (lt_with_top r) := by
-  refine ⟨?_⟩
-  intro x
-  constructor
-  intro x' hlt
-  match x' with
-  | none => contradiction
-  | some x' =>
-    clear hlt
-    induction h.apply x'
-    rename_i ih
-    constructor
-    intro x'' hlt'
-    match x'' with
-    | none => contradiction
-    | some x'' => exact ih x'' hlt'
-
 variable (m) in
 def Zip.Rel₂ [Productive α₁ m] [Finite α₂ m] :
     IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) → IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) → Prop :=
   InvImage (Prod.Lex
       IterM.TerminationMeasures.Finite.Rel
-      (Prod.Lex (Zip.lt_with_top emptyRelation) IterM.TerminationMeasures.Productive.Rel))
+      (Prod.Lex (Option.SomeLtNone.lt emptyRelation) IterM.TerminationMeasures.Productive.Rel))
     (fun it => (it.internalState.right.finitelyManySteps, (it.internalState.memoizedLeft, it.internalState.left.finitelyManySkips)))
 
 theorem Zip.rel₂_of_right [Productive α₁ m] [Finite α₂ m] {it' it : IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂)}
@@ -188,7 +176,7 @@ theorem Zip.rel₂_of_right [Productive α₁ m] [Finite α₂ m] {it' it : Iter
 
 theorem Zip.rel₂_of_memoizedLeft [Productive α₁ m] [Finite α₂ m]
     {right : IterM (α := α₂) m β₂} {b' b} {left' left : IterM (α := α₁) m β₁}
-    (h : lt_with_top emptyRelation b' b) :
+    (h : Option.SomeLtNone.lt emptyRelation b' b) :
     Zip.Rel₂ m ⟨left, b', right⟩ ⟨left', b, right⟩ :=
   Prod.Lex.right _ <| Prod.Lex.left _ _ h
 
@@ -208,7 +196,7 @@ def Zip.instFinitenessRelation₂ [Monad m] [Productive α₁ m] [Finite α₂ m
     refine ⟨fun (a, b) => Prod.lexAccessible (WellFounded.apply ?_ a) (WellFounded.apply ?_) b⟩
     · exact WellFoundedRelation.wf
     · refine ⟨fun (a, b) => Prod.lexAccessible (WellFounded.apply ?_ a) (WellFounded.apply ?_) b⟩
-      · apply Zip.wellFounded_lt_with_top
+      · apply Option.SomeLtNone.wellFounded_lt
         exact emptyWf.wf
       · exact WellFoundedRelation.wf
   subrelation {it it'} h := by
@@ -217,7 +205,7 @@ def Zip.instFinitenessRelation₂ [Monad m] [Productive α₁ m] [Finite α₂ m
     case yieldLeft hm it' out hp =>
       cases h
       apply Zip.rel₂_of_memoizedLeft
-      simp_all [Zip.lt_with_top]
+      simp_all [Option.SomeLtNone.lt]
     case skipLeft hm it' hp =>
       cases h
       apply Zip.rel₂_of_left
@@ -244,12 +232,12 @@ variable (m) in
 def Zip.Rel₃ [Productive α₁ m] [Productive α₂ m] :
     IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) → IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) → Prop :=
   InvImage (Prod.Lex
-      (lt_with_top emptyRelation)
+      (Option.SomeLtNone.lt emptyRelation)
       (Prod.Lex (IterM.TerminationMeasures.Productive.Rel) IterM.TerminationMeasures.Productive.Rel))
     (fun it => (it.internalState.memoizedLeft, (it.internalState.left.finitelyManySkips, it.internalState.right.finitelyManySkips)))
 
 theorem Zip.rel₃_of_memoizedLeft [Productive α₁ m] [Productive α₂ m] {it' it : IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂)}
-    (h : lt_with_top emptyRelation it'.internalState.memoizedLeft it.internalState.memoizedLeft) :
+    (h : Option.SomeLtNone.lt emptyRelation it'.internalState.memoizedLeft it.internalState.memoizedLeft) :
     Zip.Rel₃ m it' it :=
   Prod.Lex.left _ _ h
 
@@ -273,7 +261,7 @@ def Zip.instProductivenessRelation [Monad m] [Productive α₁ m] [Productive α
   wf := by
     apply InvImage.wf
     refine ⟨fun (a, b) => Prod.lexAccessible (WellFounded.apply ?_ a) (WellFounded.apply ?_) b⟩
-    · apply Zip.wellFounded_lt_with_top
+    · apply Option.SomeLtNone.wellFounded_lt
       exact emptyWf.wf
     · refine ⟨fun (a, b) => Prod.lexAccessible (WellFounded.apply ?_ a) (WellFounded.apply ?_) b⟩
       · exact WellFoundedRelation.wf
@@ -282,7 +270,7 @@ def Zip.instProductivenessRelation [Monad m] [Productive α₁ m] [Productive α
     cases h
     case yieldLeft hm it' out hp =>
       apply Zip.rel₃_of_memoizedLeft
-      simp [lt_with_top, hm]
+      simp [Option.SomeLtNone.lt, hm]
     case skipLeft hm it' hp =>
       obtain ⟨⟨left, memoizedLeft, right⟩⟩ := it
       simp only at hm
diff --git a/src/Std/Data/Iterators/Combinators/Zip.lean b/src/Std/Data/Iterators/Combinators/Zip.lean
@@ -8,6 +8,36 @@ import Std.Data.Iterators.Combinators.Monadic.Zip
 
 namespace Std.Iterators
 
+/--
+Given two iterators `left` and `right`, `left.zip right` is an iterator that yields pairs of
+outputs of `left` and `right`. When one of them terminates,
+the `zip` iterator will also terminate.
+
+**Marble diagram:**
+
+```text
+left               --a        ---b        --c
+right                 --x         --y        --⊥
+left.zip right     -----(a, x)------(b, y)-----⊥
+```
+
+**Termination properties:**
+
+* `Finite` instance: only if either `left` or `right` is finite and the other is productive
+* `Productive` instance: only if `left` and `right` are productive
+
+There are situations where `left.zip right` is finite (or productive) but none of the instances
+above applies. For example, if `left` immediately terminates but `right` always skips, then
+`left.zip.right` is finite even though no `Finite` (or even `Productive`) instance is available.
+Such instances need to be proved manually.
+
+**Performance:**
+
+This combinator incurs an additional O(1) cost with each step taken by `left` or `right`.
+
+Right now, the compiler does not unbox the internal state, leading to worse performance than
+theoretically possible.
+-/
 @[always_inline, inline]
 def Iter.zip {α₁ : Type w} {β₁: Type w} {α₂ : Type w} {β₂ : Type w}
     [Iterator α₁ Id β₁] [Iterator α₂ Id β₂]
