From 312abefdec082e790746c11b9c0bfecf1d35b030 Mon Sep 17 00:00:00 2001 From: "F. G. Dorais" Date: Tue, 28 May 2024 13:20:10 -0400 Subject: [PATCH] feat: add `push` and `pop` --- Batteries/Data/DArray.lean | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/Batteries/Data/DArray.lean b/Batteries/Data/DArray.lean index 065a839705..8b0ea9df82 100644 --- a/Batteries/Data/DArray.lean +++ b/Batteries/Data/DArray.lean @@ -66,6 +66,13 @@ private unsafe def umodifyFImpl [Functor f] (a : DArray n α) (i : USize) (h : i let a := unsafeCast <| a.data.uset i (unsafeCast ()) lcProof usetImpl a i h <$> t v +private unsafe def pushImpl (a : DArray n α) (v : β) : + DArray (n+1) fun i => if h : i.val < n then α ⟨i.val, h⟩ else β := + unsafeCast <| a.data.push <| unsafeCast v + +private unsafe def popImpl (a : DArray (n+1) α) : DArray n fun i => α i.castSucc := + unsafeCast <| a.data.pop + private unsafe def copyImpl (a : DArray n α) : DArray n α := unsafeCast <| a.data.extract 0 n @@ -136,6 +143,17 @@ protected abbrev umodify (a : DArray n α) (i : USize) (h : i.toNat < n) @[implemented_by copyImpl] protected def copy (a : DArray n α) : DArray n α := mk a.get +/-- Push an element onto the end of a `DArray`. `O(1)` if exclusive else `O(n)`. -/ +@[implemented_by pushImpl] +protected def push (a : DArray n α) (v : β) : + DArray (n+1) fun i => if h : i.val < n then α ⟨i.val, h⟩ else β := + mk fun i => if h : i.val < n then dif_pos h ▸ a.get ⟨i.val, h⟩ else dif_neg h ▸ v + +/-- Delete the last item of a `DArray`. `O(1)` if exclusive else `O(n)`. -/ +@[implemented_by popImpl] +protected def pop (a : DArray (n+1) α) : DArray n fun i => α i.castSucc := + mk fun i => a.get i.castSucc + @[ext] protected theorem ext : {a b : DArray n α} → (∀ i, a.get i = b.get i) → a = b | mk _, mk _, h => congrArg _ <| funext fun i => h i