79 lines
2.4 KiB
Idris
79 lines
2.4 KiB
Idris
module Data.NumIdr.PrimArray.Linked
|
|
|
|
import Data.Vect
|
|
import Data.DPair
|
|
import Data.NumIdr.Array.Rep
|
|
import Data.NumIdr.Array.Coords
|
|
|
|
%default total
|
|
|
|
|
|
export
|
|
constant : (s : Vect rk Nat) -> a -> Vects s a
|
|
constant [] x = x
|
|
constant (d :: s) xs = replicate d (constant s xs)
|
|
|
|
export
|
|
index : Coords s -> Vects s a -> a
|
|
index [] x = x
|
|
index (c :: cs) xs = index cs (index c xs)
|
|
|
|
export
|
|
update : Coords s -> (a -> a) -> Vects s a -> Vects s a
|
|
update [] f v = f v
|
|
update (i :: is) f v = updateAt i (update is f) v
|
|
|
|
export %unsafe
|
|
assertFin : Nat -> Fin n
|
|
assertFin n = natToFinLt n @{believe_me Oh}
|
|
|
|
export
|
|
indexRange : {s : _} -> (rs : CoordsRange s) -> Vects s a -> Vects (newShape rs) a
|
|
indexRange [] v = v
|
|
indexRange (r :: rs) v with (cRangeToList r)
|
|
_ | Left i = indexRange rs (Vect.index (assertFin i) v)
|
|
_ | Right is = assert_total $
|
|
case toVect _ (map (\i => indexRange rs (Vect.index (assertFin i) v)) is) of
|
|
Just v => v
|
|
|
|
export
|
|
indexSetRange : {s : _} -> (rs : CoordsRange s) -> Vects (newShape rs) a -> Vects s a -> Vects s a
|
|
indexSetRange {s=[]} [] rv _ = rv
|
|
indexSetRange {s=_::_} (r :: rs) rv v with (cRangeToList r)
|
|
_ | Left i = updateAt (assertFin i) (indexSetRange rs rv) v
|
|
_ | Right is = foldl (\v,i => updateAt (assertFin i) (indexSetRange rs (Vect.index (assertFin i) rv)) v) v is
|
|
|
|
|
|
export
|
|
fromFunction : {s : _} -> (Coords s -> a) -> Vects s a
|
|
fromFunction {s = []} f = f []
|
|
fromFunction {s = d :: s} f = tabulate (\i => fromFunction (f . (i::)))
|
|
|
|
export
|
|
fromFunctionNB : {s : _} -> (Vect rk Nat -> a) -> Vects {rk} s a
|
|
fromFunctionNB {s = []} f = f []
|
|
fromFunctionNB {s = d :: s} f = tabulate' {n=d} (\i => fromFunctionNB (f . (i::)))
|
|
where
|
|
tabulate' : forall a. {n : _} -> (Nat -> a) -> Vect n a
|
|
tabulate' {n=Z} f = []
|
|
tabulate' {n=S _} f = f Z :: tabulate' (f . S)
|
|
|
|
export
|
|
mapVects : {s : _} -> (a -> b) -> Vects s a -> Vects s b
|
|
mapVects {s = []} f x = f x
|
|
mapVects {s = _ :: _} f v = Prelude.map (mapVects f) v
|
|
|
|
export
|
|
foldl : {s : _} -> (b -> a -> b) -> b -> Vects s a -> b
|
|
foldl {s=[]} f z v = f z v
|
|
foldl {s=_::_} f z v = Prelude.foldl (Linked.foldl f) z v
|
|
|
|
export
|
|
foldr : {s : _} -> (a -> b -> b) -> b -> Vects s a -> b
|
|
foldr {s=[]} f z v = f v z
|
|
foldr {s=_::_} f z v = Prelude.foldr (flip $ Linked.foldr f) z v
|
|
|
|
export
|
|
traverse : {s : _} -> Applicative f => (a -> f b) -> Vects s a -> f (Vects s b)
|
|
traverse {s=[]} f v = f v
|
|
traverse {s=_::_} f v = Prelude.traverse (Linked.traverse f) v
|