From 264e451d3c57c80dd6801fc6d3bdc3357ec3d8d1 Mon Sep 17 00:00:00 2001
From: Sebastian Graf <sgraf1337@gmail.com>
Date: Mon, 11 Aug 2025 16:34:34 +0200
Subject: [PATCH] feat: Add `@[spec]` lemmas for `forIn` at `Std.PRange`
 (#9848)

This PR adds `@[spec]` lemmas for `forIn` and `forIn'` at `Std.PRange`.
---
 src/Std/Do/Triple/SpecLemmas.lean | 41 +++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/src/Std/Do/Triple/SpecLemmas.lean b/src/Std/Do/Triple/SpecLemmas.lean
index 71e8afaf30..8af5ec97ac 100644
--- a/src/Std/Do/Triple/SpecLemmas.lean
+++ b/src/Std/Do/Triple/SpecLemmas.lean
@@ -8,6 +8,7 @@ module
 prelude
 public import Std.Do.Triple.Basic
 public import Std.Do.WP
+import Init.Data.Range.Polymorphic
 
 @[expose] public section
 
@@ -481,6 +482,46 @@ theorem Spec.forIn_range {β : Type} {m : Type → Type v} {ps : PostShape}
   simp only [Std.Range.forIn_eq_forIn_range', Std.Range.size]
   apply Spec.forIn_list inv step
 
+open Std.PRange in
+@[spec]
+theorem Spec.forIn'_prange {α β : Type u}
+    [Monad m] [WPMonad m ps]
+    [UpwardEnumerable α]
+    [SupportsUpperBound su α] [SupportsLowerBound sl α] [HasFiniteRanges su α]
+    [BoundedUpwardEnumerable sl α] [LawfulUpwardEnumerable α]
+    [LawfulUpwardEnumerableLowerBound sl α] [LawfulUpwardEnumerableUpperBound su α]
+    {xs : PRange ⟨sl, su⟩ α} {init : β} {f : (a : α) → a ∈ xs → β → m (ForInStep β)}
+    (inv : Invariant xs.toList β ps)
+    (step : ∀ b rpref x (hx : x ∈ xs) suff (h : xs.toList = rpref.reverse ++ x :: suff),
+        ⦃inv.1 (⟨rpref, x::suff, by simp [h]⟩, b)}
+        f x hx b
+        ⦃(fun r => match r with
+                   | .yield b' => inv.1 (⟨x::rpref, suff, by simp [h]⟩, b')
+                   | .done b' => inv.1 (⟨xs.toList.reverse, [], by simp⟩, b'), inv.2)}) :
+    ⦃inv.1 (⟨[], xs.toList, by simp⟩, init)} forIn' xs init f ⦃(fun b => inv.1 (⟨xs.toList.reverse, [], by simp⟩, b), inv.2)} := by
+  simp only [forIn'_eq_forIn'_toList]
+  apply Spec.forIn'_list inv (fun b rpref x hx suff h => step b rpref x (mem_toList_iff_mem.mp hx) suff h)
+
+open Std.PRange in
+@[spec]
+theorem Spec.forIn_prange {α β : Type u}
+    [Monad m] [WPMonad m ps]
+    [UpwardEnumerable α]
+    [SupportsUpperBound su α] [SupportsLowerBound sl α] [HasFiniteRanges su α]
+    [BoundedUpwardEnumerable sl α] [LawfulUpwardEnumerable α]
+    [LawfulUpwardEnumerableLowerBound sl α] [LawfulUpwardEnumerableUpperBound su α]
+    {xs : PRange ⟨sl, su⟩ α} {init : β} {f : α → β → m (ForInStep β)}
+    (inv : Invariant xs.toList β ps)
+    (step : ∀ b rpref x suff (h : xs.toList = rpref.reverse ++ x :: suff),
+        ⦃inv.1 (⟨rpref, x::suff, by simp [h]⟩, b)}
+        f x b
+        ⦃(fun r => match r with
+                   | .yield b' => inv.1 (⟨x::rpref, suff, by simp [h]⟩, b')
+                   | .done b' => inv.1 (⟨xs.toList.reverse, [], by simp⟩, b'), inv.2)}) :
+    ⦃inv.1 (⟨[], xs.toList, by simp⟩, init)} forIn xs init f ⦃(fun b => inv.1 (⟨xs.toList.reverse, [], by simp⟩, b), inv.2)} := by
+  simp only [forIn]
+  apply Spec.forIn'_prange inv (fun b rpref x _hx suff h => step b rpref x suff h)
+
 @[spec]
 theorem Spec.forIn'_array {α β : Type u}
     [Monad m] [WPMonad m ps]