add tests

Author: datokrat

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

@@ -17,9 +17,14 @@ namespace Std.Iterators
 
 variable {α : Type w} {m : Type w → Type w'} {β : Type w}
 
+/--
+The internal state of the `IterM.drop` combinator.
+-/
 @[unbox]
 structure Drop (α : Type w) (m : Type w → Type w') (β : Type w) where
+  /-- Internal implementation detail of the iterator library -/
   remaining : Nat
+  /-- Internal implementation detail of the iterator library -/
   inner : IterM (α := α) m β
 
 /--

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

@@ -41,6 +41,7 @@ theorem Iter.atIdxSlow?_drop {α β}
     rw [atIdxSlow?.eq_def, atIdxSlow?.eq_def, step_drop]
     cases it.step using PlausibleIterStep.casesOn <;> simp [*]
 
+@[simp]
 theorem Iter.toList_drop {α β} [Iterator α Id β] {n : Nat}
     [Finite α Id] [IteratorCollect α Id] [LawfulIteratorCollect α Id]
     {it : Iter (α := α) β} :
@@ -49,6 +50,7 @@ theorem Iter.toList_drop {α β} [Iterator α Id β] {n : Nat}
   simp only [getElem?_toList_eq_atIdxSlow?, List.getElem?_drop, atIdxSlow?_drop]
   rw [Nat.add_comm]
 
+@[simp]
 theorem Iter.toListRev_drop {α β} [Iterator α Id β] {n : Nat}
     [Finite α Id] [IteratorCollect α Id] [LawfulIteratorCollect α Id]
     {it : Iter (α := α) β} :
@@ -57,18 +59,19 @@ theorem Iter.toListRev_drop {α β} [Iterator α Id β] {n : Nat}
 
 -- TODO: Init.Data.Array.Lemmas
 theorem _root_.List.drop_toArray {l : List α} {k : Nat} :
-    l.toArray.drop k = (l.drop k).toArray := by
+    (l.drop k).toArray = l.toArray.extract k := by
   induction l generalizing k
   case nil => simp
   case cons l' ih =>
     match k with
     | 0 => simp
-    | k' + 1 => simp [List.drop_succ_cons, ← ih]
+    | k' + 1 => simp [List.drop_succ_cons, ih]
 
+@[simp]
 theorem Iter.toArray_drop {α β} [Iterator α Id β] {n : Nat}
     [Finite α Id] [IteratorCollect α Id] [LawfulIteratorCollect α Id]
     {it : Iter (α := α) β} :
-    (it.drop n).toArray = it.toArray.drop n := by
-  rw [← toArray_toList, ← toArray_toList, List.drop_toArray, toList_drop]
+    (it.drop n).toArray = it.toArray.extract n := by
+  rw [← toArray_toList, ← toArray_toList, ← List.drop_toArray, toList_drop]
 
 end Std.Iterators

tests/lean/run/iterators.lean

@@ -144,3 +144,41 @@ example : ([1, 2, 3].iter.take 2).toListRev = [2, 1] := by
   simp
 
 end Take
+
+section Drop
+
+def sumDropRec (l : List Nat) : Nat :=
+  go (l.iter.drop 1) 0
+where
+  go it acc :=
+    match it.step with
+    | .yield it' out _ => go it' (acc + out)
+    | .skip it' _ => go it' acc
+    | .done _ => acc
+  termination_by it.finitelyManySteps
+
+def sumDropFold (l : List Nat) : Nat :=
+  l.iter.drop 1 |>.fold (init := 0) (· + ·)
+
+/-- info: [2, 3] -/
+#guard_msgs in
+#eval [1, 2, 3].iter.drop 1 |>.toList
+
+/-- info: 5 -/
+#guard_msgs in
+#eval sumDropRec [1, 2, 3]
+
+/-- info: 5 -/
+#guard_msgs in
+#eval sumDropFold [1, 2, 3]
+
+example : ([1, 2, 3].iter.drop 1).toList = [2, 3] := by
+  simp
+
+example : ([1, 2, 3].iter.drop 1).toArray = #[2, 3] := by
+  simp
+
+example : ([1, 2, 3].iter.drop 1).toListRev = [3, 2] := by
+  simp
+
+end Drop