feat: add ForIn' instance that is similar to ForIn but provides a proof that the iterated elements are in the collection

src/Init/Core.lean

@@ -89,6 +89,12 @@ class ForIn (m : Type u₁ → Type u₂) (ρ : Type u) (α : outParam (Type v))
 
 export ForIn (forIn)
 
+class ForIn' (m : Type u₁ → Type u₂) (ρ : Type u) (α : outParam (Type v)) [d : outParam $ Membership α ρ] where
+  forIn' {β} [Monad m] (x : ρ) (b : β) (f : (a : α) → a ∈ x → β → m (ForInStep β)) : m β
+
+export ForIn' (forIn')
+
+
 /- Auxiliary type used to compile `do` notation. -/
 inductive DoResultPRBC (α β σ : Type u) where
   | «pure»     : α → σ → DoResultPRBC α β σ

src/Init/Data/List/Basic.lean

@@ -81,6 +81,11 @@ theorem append_assoc (as bs cs : List α) : (as ++ bs) ++ cs = as ++ (bs ++ cs)
   | nil => rfl
   | cons a as ih => simp [ih]
 
+theorem append_cons (as : List α) (b : α) (bs : List α) : as ++ b :: bs = as ++ [b] ++ bs := by
+  induction as with
+  | nil => simp
+  | cons a as ih => simp [ih]
+
 instance : EmptyCollection (List α) := ⟨List.nil⟩
 
 protected def erase {α} [BEq α] : List α → α → List α
@@ -236,6 +241,18 @@ theorem elem_eq_true_of_mem [DecidableEq α] {a : α} {as : List α} (h : a ∈
 instance [DecidableEq α] (a : α) (as : List α) : Decidable (a ∈ as) :=
   decidable_of_decidable_of_iff (Iff.intro mem_of_elem_eq_true elem_eq_true_of_mem)
 
+theorem mem_append_of_mem_left {a : α} {as : List α} (bs : List α) : a ∈ as → a ∈ as ++ bs := by
+  intro h
+  induction h with
+  | head => apply Mem.head
+  | tail => apply Mem.tail; assumption
+
+theorem mem_append_of_mem_right {b : α} {bs : List α} (as : List α) : b ∈ bs → b ∈ as ++ bs := by
+  intro h
+  induction as with
+  | nil  => simp [h]
+  | cons => apply Mem.tail; assumption
+
 def eraseDupsAux {α} [BEq α] : List α → List α → List α
   | [],    bs => bs.reverse
   | a::as, bs => match bs.elem a with

src/Init/Data/List/Control.lean

@@ -158,6 +158,24 @@ instance : ForIn m (List α) α where
     : forIn (a::as) b f = f a b >>= fun | ForInStep.done b => pure b | ForInStep.yield b => forIn as b f :=
   rfl
 
+@[inline] protected def forIn' {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : List α) (init : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β :=
+  let rec @[specialize] loop : (as' : List α) → (b : β) → Exists (fun bs => bs ++ as' = as) → m β
+    | [], b, _    => pure b
+    | a::as', b, h => do
+      have : a ∈ as := by
+        have ⟨bs, h⟩ := h
+        subst h
+        exact mem_append_of_mem_right _ (Mem.head ..)
+      match (← f a this b) with
+      | ForInStep.done b  => pure b
+      | ForInStep.yield b =>
+        have : Exists (fun bs => bs ++ as' = as) := have ⟨bs, h⟩ := h; ⟨bs ++ [a], by rw [← h, append_cons bs a as']⟩
+        loop as' b this
+  loop as init ⟨[], rfl⟩
+
+instance : ForIn' m (List α) α where
+  forIn' := List.forIn'
+
 instance : ForM m (List α) α where
   forM := List.forM