diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean index 7a1556efa513..8b35d1a91704 100644 --- a/src/Init/Data/Array/Basic.lean +++ b/src/Init/Data/Array/Basic.lean @@ -727,33 +727,38 @@ def takeWhile (p : α → Bool) (as : Array α) : Array α := termination_by as.size - i go 0 #[] -def eraseIdxAux (i : Nat) (a : Array α) : Array α := - if h : i < a.size then - let idx : Fin a.size := ⟨i, h⟩; - let idx1 : Fin a.size := ⟨i - 1, by exact Nat.lt_of_le_of_lt (Nat.pred_le i) h⟩; - let a' := a.swap idx idx1 - eraseIdxAux (i+1) a' +/-- Remove the element at a given index from an array without bounds checks, using a `Fin` index. + + This function takes worst case O(n) time because + it has to backshift all elements at positions greater than `i`.-/ +def feraseIdx (a : Array α) (i : Fin a.size) : Array α := + if h : i.val + 1 < a.size then + let a' := a.swap ⟨i.val + 1, h⟩ i + let i' : Fin a'.size := ⟨i.val + 1, by simp [a', h]⟩ + have : a'.size - i' < a.size - i := by + simp [a', Nat.sub_succ_lt_self _ _ i.isLt] + a'.feraseIdx i' else a.pop -termination_by a.size - i +termination_by a.size - i.val -def feraseIdx (a : Array α) (i : Fin a.size) : Array α := - eraseIdxAux (i.val + 1) a +derive_functional_induction feraseIdx -def eraseIdx (a : Array α) (i : Nat) : Array α := - if i < a.size then eraseIdxAux (i+1) a else a +theorem size_feraseIdx (a : Array α) (i : Fin a.size) : (a.feraseIdx i).size = a.size - 1 := by + induction a, i using feraseIdx.induct with + | @case1 a i h a' _ _ ih => + unfold feraseIdx + simp [h, a', ih] + | case2 a i h => + unfold feraseIdx + simp [h] -def eraseIdxSzAux (a : Array α) (i : Nat) (r : Array α) (heq : r.size = a.size) : { r : Array α // r.size = a.size - 1 } := - if h : i < r.size then - let idx : Fin r.size := ⟨i, h⟩; - let idx1 : Fin r.size := ⟨i - 1, by exact Nat.lt_of_le_of_lt (Nat.pred_le i) h⟩; - eraseIdxSzAux a (i+1) (r.swap idx idx1) ((size_swap r idx idx1).trans heq) - else - ⟨r.pop, (size_pop r).trans (heq ▸ rfl)⟩ -termination_by r.size - i +/-- Remove the element at a given index from an array, or do nothing if the index is out of bounds. -def eraseIdx' (a : Array α) (i : Fin a.size) : { r : Array α // r.size = a.size - 1 } := - eraseIdxSzAux a (i.val + 1) a rfl + This function takes worst case O(n) time because + it has to backshift all elements at positions greater than `i`.-/ +def eraseIdx (a : Array α) (i : Nat) : Array α := + if h : i < a.size then a.feraseIdx ⟨i, h⟩ else a def erase [BEq α] (as : Array α) (a : α) : Array α := match as.indexOf? a with diff --git a/src/Lean/Data/PersistentHashMap.lean b/src/Lean/Data/PersistentHashMap.lean index e9b94cab2d06..b90a38258950 100644 --- a/src/Lean/Data/PersistentHashMap.lean +++ b/src/Lean/Data/PersistentHashMap.lean @@ -226,8 +226,10 @@ partial def eraseAux [BEq α] : Node α β → USize → α → Node α β × Bo | n@(Node.collision keys vals heq), _, k => match keys.indexOf? k with | some idx => - let ⟨keys', keq⟩ := keys.eraseIdx' idx - let ⟨vals', veq⟩ := vals.eraseIdx' (Eq.ndrec idx heq) + let keys' := keys.feraseIdx idx + have keq := keys.size_feraseIdx idx + let vals' := vals.feraseIdx (Eq.ndrec idx heq) + have veq := vals.size_feraseIdx (Eq.ndrec idx heq) have : keys.size - 1 = vals.size - 1 := by rw [heq] (Node.collision keys' vals' (keq.trans (this.trans veq.symm)), true) | none => (n, false)