wip

Author: datokrat

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

@@ -255,12 +255,33 @@ def Iter.Partial.find? {α β : Type w} [Iterator α Id β] [IteratorLoopPartial
     (it : Iter.Partial (α := α) β) (f : β → Bool) : Option β :=
   Id.run (it.findM? (pure <| .up <| f ·))
 
-@[always_inline, inline, expose, inherit_doc IterM.size]
+/--
+Computes how many elements the iterator returns. In contrast to `count`, this function might take
+shortcuts instead of actually stepping through the whole iterator.
+
+**Performance**:
+
+Depends on the concrete type of the iterator. The default implementation, which is identical to
+`count`, is linear, but it is sometimes possible to provide an O(1) implementation.
+-/
+@[always_inline, inline, expose]
 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]
+/--
+Computes how many elements the iterator returns. In contrast to `count`, this function might take
+shortcuts instead of actually stepping through the whole iterator.
+
+This is the partial version of `size`. It does not require a proof of finiteness and might loop
+forever. It is not possible to verify the behavior in Lean because it uses `partial`.
+
+**Performance**:
+
+Depends on the concrete type of the iterator. The default implementation, which is identical to
+`count`, is linear, but it is sometimes possible to provide an O(1) implementation.
+-/
+@[always_inline, inline]
 def Iter.Partial.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorSizePartial α Id]
     (it : Iter (α := α) β) : Nat :=
   (IteratorSizePartial.size it.toIterM).run.down

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

@@ -639,7 +639,7 @@ section Size
 This is the implementation of the default instance `IteratorSize.defaultImplementation`.
 -/
 @[always_inline, inline]
-def IterM.DefaultConsumers.size {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
+def IterM.DefaultConsumers.count {α : 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)
@@ -648,70 +648,89 @@ def IterM.DefaultConsumers.size {α : Type w} {m : Type w → Type w'} [Monad m]
 This is the implementation of the default instance `IteratorSizePartial.defaultImplementation`.
 -/
 @[always_inline, inline]
-def IterM.DefaultConsumers.sizePartial {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
+def IterM.DefaultConsumers.countPartial {α : 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)
 
+/--
+Steps through the whole iterator, counting the number of outputs emitted.
+
+**Performance**:
+
+linear in the number of steps taken by the iterator
+-/
+@[always_inline, inline]
+def IterM.count {α : Type} {m : Type → Type w'} {β : Type} [Iterator α m β] [IteratorLoop α m m]
+    [Monad m] [Finite α m] (it : IterM (α := α) m β) : m Nat :=
+  ULift.down <$> IterM.DefaultConsumers.count it
+
+/--
+Steps through the whole iterator, counting the number of outputs emitted.
+
+**Performance**:
+
+linear in the number of steps taken by the iterator
+-/
+@[always_inline, inline]
+def IterM.Partial.count {α : Type} {m : Type → Type w'} {β : Type} [Iterator α m β]
+    [IteratorLoopPartial α m m] [Monad m] (it : IterM.Partial (α := α) m β)
+    [IteratorSizePartial α m] : m Nat :=
+  ULift.down <$> IterM.DefaultConsumers.countPartial it.it
+
+set_option pp.universes true
 /--
 This is the default implementation of the `IteratorSize` class.
 It simply iterates using `IteratorLoop` and counts the elements.
 For certain iterators, more efficient implementations are possible and should be used instead.
 -/
 @[always_inline, inline]
-def IteratorSize.defaultImplementation {α β : Type w} {m : Type w → Type w'} [Monad m]
+def IteratorSize.defaultImplementation {α β : Type w} [Monad m]
     [Iterator α m β] [Finite α m] [IteratorLoop α m m] :
     IteratorSize α m where
-  size := IterM.DefaultConsumers.size
-
+  size := IterM.DefaultConsumers.count
 
 /--
 This is the default implementation of the `IteratorSizePartial` class.
 It simply iterates using `IteratorLoopPartial` and counts the elements.
 For certain iterators, more efficient implementations are possible and should be used instead.
 -/
 @[always_inline, inline]
-instance IteratorSizePartial.defaultImplementation {α β : Type w} {m : Type w → Type w'} [Monad m]
+instance IteratorSizePartial.defaultImplementation {α β : Type w} [Monad m]
     [Iterator α m β] [IteratorLoopPartial α m m] :
     IteratorSizePartial α m where
-  size := IterM.DefaultConsumers.sizePartial
+  size := IterM.DefaultConsumers.countPartial
 
 /--
-Computes how many elements the iterator returns. In monadic situations, it is unclear which effects
-are caused by calling `size`, and if the monad is nondeterministic, it is also unclear what the
-returned value should be. The reference implementation, `IteratorSize.defaultImplementation`,
-simply iterates over the whole iterator monadically, counting the number of emitted values.
-An `IteratorSize` instance is considered lawful if it is equal to the reference implementation.
+Computes how many elements the iterator returns. In contrast to `count`, this function might take
+shortcuts instead of actually stepping through the whole iterator.
 
 **Performance**:
 
-Default performance is linear in the number of steps taken by the iterator.
+Depends on the concrete type of the iterator. The default implementation, which is identical to
+`count`, is linear, but it is sometimes possible to provide an O(1) implementation.
 -/
 @[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
+def IterM.size {α : Type} {β : Type} [Iterator α Id β]
+    (it : IterM (α := α) Id β) [IteratorSize α Id] : Nat :=
+  (IteratorSize.size it).run.down
 
 /--
-Computes how many elements the iterator emits.
-
-With monadic iterators (`IterM`), it is unclear which effects
-are caused by calling `size`, and if the monad is nondeterministic, it is also unclear what the
-returned value should be. The reference implementation, `IteratorSize.defaultImplementation`,
-simply iterates over the whole iterator monadically, counting the number of emitted values.
-An `IteratorSize` instance is considered lawful if it is equal to the reference implementation.
+Computes how many elements the iterator returns. In contrast to `count`, this function might take
+shortcuts instead of actually stepping through the whole iterator.
 
 This is the partial version of `size`. It does not require a proof of finiteness and might loop
 forever. It is not possible to verify the behavior in Lean because it uses `partial`.
 
 **Performance**:
 
-Default performance is linear in the number of steps taken by the iterator.
+Depends on the concrete type of the iterator. The default implementation, which is identical to
+`count`, is linear, but it is sometimes possible to provide an O(1) implementation.
 -/
 @[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
+def IterM.Partial.size {α : Type} {β : Type} [Iterator α Id β]
+    (it : IterM.Partial (α := α) Id β) [IteratorSizePartial α Id] : Nat :=
+  (IteratorSizePartial.size it.it).run.down
 
 end Size
 

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

@@ -154,12 +154,12 @@ instance Types.StepSizeIterator.instIteratorLoopPartial {m n} [Iterator α m β]
     IteratorLoopPartial (Types.StepSizeIterator α m β) m n :=
   .defaultImplementation
 
-instance Types.StepSizeIterator.instIteratorSize {m} [Iterator α m β]
-    [IteratorAccess α m] [Monad m] [Finite (Types.StepSizeIterator α m β) m] :
-    IteratorSize (Types.StepSizeIterator α m β) m :=
+instance Types.StepSizeIterator.instIteratorSize [Iterator α Id β]
+    [IteratorAccess α Id] [Monad Id] [Finite (Types.StepSizeIterator α Id β) Id] :
+    IteratorSize (Types.StepSizeIterator α Id β) Id :=
   .defaultImplementation
 
-instance Types.StepSizeIterator.instIteratorSizePartial {m} [Iterator α m β]
+instance Types.StepSizeIterator.instIteratorSizePartial [Iterator α m β]
     [IteratorAccess α m] [Monad m] :
     IteratorSizePartial (Types.StepSizeIterator α m β) m :=
   .defaultImplementation