diff --git a/library/init/data/array/basic.lean b/library/init/data/array/basic.lean index d328525642..336c27f1e0 100644 --- a/library/init/data/array/basic.lean +++ b/library/init/data/array/basic.lean @@ -88,6 +88,9 @@ def update (a : Array α) (i : @& Fin a.size) (v : α) : Array α := { sz := a.sz, data := λ j, if h : i = j then v else a.data j } +theorem szUpdateEq (a : Array α) (i : Fin a.size) (v : α) : (update a i v).size = a.size := +rfl + /- Low-level version of `update` which is as fast as a C array update. `Fin` values are represented as tag pointers in the Lean runtime. Thus, `update` may be slightly slower than `updt`. -/ @@ -114,8 +117,13 @@ if h₂ : j < a.size then fswap a ⟨i, h₁⟩ ⟨j, h₂⟩ else a else a -theorem szUpdateEq (a : Array α) (i : Fin a.size) (v : α) : (update a i v).size = a.size := -rfl +@[inline] def fswapAt {α : Type} (a : Array α) (i : Fin a.size) (v : α) : α × Array α := +let e := a.index i in +let a := a.update i v in +(e, a) + +@[inline] def swapAt {α : Type} (a : Array α) (i : Nat) (v : α) : α × Array α := +if h : i < a.size then fswapAt a ⟨i, h⟩ v else (v, a) @[extern cpp inline "lean::array_pop(#2)"] def pop (a : Array α) : Array α := diff --git a/library/init/lean/compiler/ir.lean b/library/init/lean/compiler/ir.lean index 1ea19ac48a..6f4220994f 100644 --- a/library/init/lean/compiler/ir.lean +++ b/library/init/lean/compiler/ir.lean @@ -273,10 +273,9 @@ partial def reshapeAux : Array FnBody → Nat → FnBody → FnBody | a i b := if i == 0 then b else - let i := i - 1 in - let curr := a.get i in - let a := a.set i (default _) in - let b := curr.setBody b in + let i := i - 1 in + let (curr, a) := a.swapAt i (default _) in + let b := curr.setBody b in reshapeAux a i b def reshape (bs : Array FnBody) (term : FnBody) : FnBody :=