From 16cb6fe75901a20d3f78a1f1cecdbd3cc699bdc1 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 30 Dec 2019 08:11:23 -0800
Subject: [PATCH] feat: add `foldrRangeM`

---
 src/Init/Data/Array/Basic.lean | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean
index 055d8d34ab..54ca8b4bb2 100644
--- a/src/Init/Data/Array/Basic.lean
+++ b/src/Init/Data/Array/Basic.lean
@@ -343,6 +343,19 @@ iterateRevMAux a f a.size (Nat.leRefl _) b
 @[inline] def foldrM (f : α → β → m β) (init : β) (a : Array α) : m β :=
 iterateRevM a init (fun _ => f)
 
+@[specialize] private def foldrRangeMAux (a : Array α) (f : α → β → m β) (beginIdx : Nat) : ∀ (i : Nat), i ≤ a.size → β → m β
+| 0,   h, b => pure b
+| j+1, h, b => do
+  let i : Fin a.size := ⟨j, h⟩;
+  b ← f (a.get i) b;
+  if j == beginIdx then pure b else foldrRangeMAux j (Nat.leOfLt h) b
+
+@[inline] def foldrRangeM (beginIdx endIdx : Nat) (f : α → β → m β) (ini : β) (a : Array α) : m β :=
+if h : endIdx ≤ a.size then
+  foldrRangeMAux a f beginIdx endIdx h ini
+else
+  foldrRangeMAux a f beginIdx a.size (Nat.leRefl _) ini
+
 end
 
 @[inline] def iterateRev {β} (a : Array α) (b : β) (f : ∀ (i : Fin a.size), α → β → β) : β :=
@@ -351,6 +364,9 @@ Id.run $ iterateRevM a b f
 @[inline] def foldr {β} (f : α → β → β) (init : β) (a : Array α) : β :=
 Id.run $ foldrM f init a
 
+@[inline] def foldrRange {β} (beginIdx endIdx : Nat) (f : α → β → β) (init : β) (a : Array α) : β :=
+Id.run $ foldrRangeM beginIdx endIdx f init a
+
 def toList (a : Array α) : List α :=
 a.foldr List.cons []