From 3ecf8ac8ecfd9c7de8a514b144e4d0711dc98660 Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Mon, 5 Aug 2019 07:41:41 -0700 Subject: [PATCH] feat(library/init/data/persistentarray/basic): add `mfoldlFrom` and `foldlFrom` --- library/init/data/persistentarray/basic.lean | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/library/init/data/persistentarray/basic.lean b/library/init/data/persistentarray/basic.lean index 2bb6b04675..ff205a04e3 100644 --- a/library/init/data/persistentarray/basic.lean +++ b/library/init/data/persistentarray/basic.lean @@ -205,6 +205,20 @@ do b ← t.tail.mfindRev f; | none => mfindRevAux f t.root | some b => pure (some b) +partial def mfoldlFromAux (f : β → α → m β) : PersistentArrayNode α → USize → USize → β → m β +| (node cs) i shift b := do + let j := (div2Shift i shift).toNat; + b ← mfoldlFromAux (cs.get j) (mod2Shift i shift) (shift - initShift) b; + cs.mfoldlFrom (fun b c => mfoldlAux f c b) b (j+1) +| (leaf vs) i _ b := vs.mfoldlFrom f b i.toNat + +def mfoldlFrom (t : PersistentArray α) (f : β → α → m β) (b : β) (ini : Nat) : m β := +if ini >= t.tailOff then + t.tail.mfoldlFrom f b (ini - t.tailOff) +else do + b ← mfoldlFromAux f t.root (USize.ofNat ini) t.shift b; + t.tail.mfoldl f b + end @[inline] def foldl {β} (t : PersistentArray α) (f : β → α → β) (b : β) : β := @@ -216,6 +230,9 @@ Id.run (t.mfind f) @[inline] def findRev {β} (t : PersistentArray α) (f : α → (Option β)) : Option β := Id.run (t.mfindRev f) +@[inline] def foldlFrom {β} (t : PersistentArray α) (f : β → α → β) (b : β) (ini : Nat) : β := +Id.run (t.mfoldlFrom f b ini) + def toList (t : PersistentArray α) : List α := (t.foldl (fun xs x => x :: xs) []).reverse