Skip to content

Commit

Permalink
feat: add map and toArray
Browse files Browse the repository at this point in the history
  • Loading branch information
fgdorais committed May 30, 2024
1 parent bf6f528 commit 517399b
Showing 1 changed file with 18 additions and 0 deletions.
18 changes: 18 additions & 0 deletions Batteries/Data/DArray/Basic.lean
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,14 @@ private unsafe def popImpl (a : DArray (n+1) α) : DArray n fun i => α i.castSu
private unsafe def copyImpl (a : DArray n α) : DArray n α :=
unsafeCast <| a.data.extract 0 n

private unsafe def toArrayImpl (a : DArray n fun _ => α) : Array α :=
unsafeCast a

@[specialize]
private unsafe def mapImpl (f : {i : Fin n} → α i → β i) (a : DArray n α) : DArray n β :=
let f := fun i x => (unsafeCast (f (i:=i.cast lcProof) (unsafeCast x)) : NonScalar)
unsafeCast <| a.data.mapIdx f

private unsafe def foldlMImpl [Monad m] (a : DArray n α) (f : β → {i : Fin n} → α i → m β)
(init : β) : m β :=
if n < USize.size then
Expand Down Expand Up @@ -185,6 +193,16 @@ protected def push (a : DArray n α) (v : β) :
protected def pop (a : DArray (n+1) α) : DArray n fun i => α i.castSucc :=
mk fun i => a.get i.castSucc

/-- Cast a dependent array with constant types to an array. `O(1)` if exclusive else `O(n)`. -/
@[implemented_by toArrayImpl]
protected def toArray (a : DArray n fun _ => α) : Array α :=
.ofFn fun i => a.get i

/-- Applies `f` to each element of a dependent array, returns the array of results. -/
@[implemented_by mapImpl]
protected def map (f : {i : Fin n} → α i → β i) (a : DArray n α) : DArray n β :=
mk fun i => f (a.get i)

/--
Folds a monadic function over a `DArray` from left to right:
```
Expand Down

0 comments on commit 517399b

Please sign in to comment.