From 83e4c63d277ba52486c60915aa65aff4ea5d686e Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 4 Jun 2019 10:40:46 -0700
Subject: [PATCH] feat(library/init/data/persistentarray/basic): add
 `PersistentArray.modify`

---
 library/init/data/persistentarray/basic.lean | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/library/init/data/persistentarray/basic.lean b/library/init/data/persistentarray/basic.lean
index 4080d6cff6..fa77aa2aba 100644
--- a/library/init/data/persistentarray/basic.lean
+++ b/library/init/data/persistentarray/basic.lean
@@ -68,6 +68,20 @@ if i >= t.tailOff then
 else
   { root := setAux t.root (USize.ofNat i) t.shift a, .. t }
 
+@[specialize] partial def modifyAux [Inhabited α] (f : α → α) : PersistentArrayNode α → USize → USize → PersistentArrayNode α
+| (node cs) i shift :=
+  let j     := div2Shift i shift in
+  let i     := mod2Shift i shift in
+  let shift := shift - initShift in
+  node $ cs.modify j.toNat $ λ c, modifyAux c i shift
+| (leaf cs) i _     := leaf (cs.modify i.toNat f)
+
+@[specialize] def modify [Inhabited α] (t : PersistentArray α) (i : Nat) (f : α → α) : PersistentArray α :=
+if i >= t.tailOff then
+  { tail := t.tail.modify (i - t.tailOff) f, .. t }
+else
+  { root := modifyAux f t.root (USize.ofNat i) t.shift, .. t }
+
 partial def mkNewPath : USize → Array α → PersistentArrayNode α
 | shift a :=
   if shift == 0 then