diff --git a/Batteries/Data/DArray.lean b/Batteries/Data/DArray.lean index 8f0458da05..09a00fe705 100644 --- a/Batteries/Data/DArray.lean +++ b/Batteries/Data/DArray.lean @@ -49,6 +49,9 @@ private unsafe def ugetImpl (a : DArray n α) (i : USize) (h : i.toNat < n) : α private unsafe def setImpl (a : DArray n α) (i) (v : α i) : DArray n α := unsafeCast <| a.data.set ⟨i.val, lcProof⟩ <| unsafeCast v +private unsafe def usetImpl (a : DArray n α) (i : USize) (h : i.toNat < n) (v : α ⟨i.toNat, h⟩) : + DArray n α := unsafeCast <| a.data.uset i (unsafeCast v) lcProof + private unsafe def copyImpl (a : DArray n α) : DArray n α := unsafeCast <| a.data.extract 0 n @@ -84,6 +87,14 @@ attribute [implemented_by casesOnImpl] DArray.casesOn protected def set (a : DArray n α) (i : Fin n) (v : α i) : DArray n α := mk fun j => if h : i = j then h ▸ v else a.get j +/-- +Sets the `DArray` item at index `i : USize`. +Slightly faster than `set` and `O(1)` if exclusive else `O(n)`. +-/ +@[implemented_by usetImpl] +protected def uset (a : DArray n α) (i : USize) (h : i.toNat < n) (v : α ⟨i.toNat, h⟩) := + a.set ⟨i.toNat, h⟩ v + @[simp, inherit_doc DArray.set] protected abbrev setN (a : DArray n α) (i) (h : i < n := by get_elem_tactic) (v : α ⟨i, h⟩) := a.set ⟨i, h⟩ v @@ -102,10 +113,6 @@ theorem get_mk (i : Fin n) : DArray.get (.mk init) i = init i := rfl theorem set_mk {α : Fin n → Type _} {init : (i : Fin n) → α i} (i : Fin n) (v : α i) : DArray.set (.mk init) i v = .mk fun j => if h : i = j then h ▸ v else init j := rfl -@[simp] -theorem uget_eq_get (a : DArray n α) (i : USize) (h : i.toNat < n) : - a.uget i h = a.get ⟨i.toNat, h⟩ := rfl - @[simp] theorem get_set (a : DArray n α) (i : Fin n) (v : α i) : (a.set i v).get i = v := by simp only [DArray.get, DArray.set, dif_pos] @@ -121,5 +128,13 @@ theorem set_set (a : DArray n α) (i : Fin n) (v w : α i) : (a.set i v).set i w else rw [get_set_ne _ _ h, get_set_ne _ _ h, get_set_ne _ _ h] +@[simp] +theorem uget_eq_get (a : DArray n α) (i : USize) (h : i.toNat < n) : + a.uget i h = a.get ⟨i.toNat, h⟩ := rfl + +@[simp] +theorem uset_eq_set (a : DArray n α) (i : USize) (h : i.toNat < n) (v : α ⟨i.toNat, h⟩) : + a.uset i h v = a.set ⟨i.toNat, h⟩ v := rfl + @[simp] theorem copy_eq (a : DArray n α) : a.copy = a := rfl