Add operators for indexing
This commit is contained in:
parent
861f1e29f2
commit
97d1bdb538
|
@ -201,12 +201,31 @@ array v = MkArray COrder (calcStrides COrder s) s (fromList $ collapse v)
|
|||
-- Indexing
|
||||
--------------------------------------------------------------------------------
|
||||
|
||||
infix 2 !!
|
||||
infix 2 !?
|
||||
infixl 3 !!..
|
||||
infix 3 !?..
|
||||
|
||||
|
||||
||| Index the array using the given `Coords` object.
|
||||
export
|
||||
index : Coords s -> Array s a -> a
|
||||
index is arr = index (getLocation (strides arr) is) (getPrim arr)
|
||||
|
||||
export
|
||||
(!!) : Array s a -> Coords s -> a
|
||||
(!!) = flip index
|
||||
|
||||
|
||||
export
|
||||
indexMaybe : Vect rk Nat -> Array {rk} s a -> Maybe a
|
||||
indexMaybe is arr = safeIndex (getLocation' (strides arr) is) (getPrim arr)
|
||||
|
||||
export
|
||||
(!?) : Array {rk} s a -> Vect rk Nat -> Maybe a
|
||||
(!?) = flip indexMaybe
|
||||
|
||||
|
||||
||| Index the array using the given `CoordsRange` object.
|
||||
export
|
||||
indexRange : (rs : CoordsRange s) -> Array s a -> Array (newShape rs) a
|
||||
|
@ -218,6 +237,10 @@ indexRange rs arr = let ord = getOrder arr
|
|||
map (\(is,is') => (getLocation' sts is', index (getLocation' (strides arr) is) (getPrim arr))) $
|
||||
getCoordsList {s = shape arr} $ rewrite sym $ shapeEq arr in rs)
|
||||
|
||||
export
|
||||
(!!..) : Array s a -> (rs : CoordsRange s) -> Array (newShape rs) a
|
||||
arr !!.. rs = indexRange rs arr
|
||||
|
||||
|
||||
--------------------------------------------------------------------------------
|
||||
-- Operations on arrays
|
||||
|
|
|
@ -64,12 +64,20 @@ namespace CoordsRange
|
|||
EndBound : Fin (S n) -> CRange n
|
||||
Bounds : Fin (S n) -> Fin (S n) -> CRange n
|
||||
|
||||
infix 0 ...
|
||||
|
||||
public export
|
||||
(...) : Fin (S n) -> Fin (S n) -> CRange n
|
||||
(...) = Bounds
|
||||
|
||||
|
||||
public export
|
||||
data CoordsRange : Vect rk Nat -> Type where
|
||||
Nil : CoordsRange []
|
||||
(::) : CRange d -> CoordsRange s -> CoordsRange (d :: s)
|
||||
|
||||
|
||||
public export
|
||||
cRangeToBounds : {n : Nat} -> CRange n -> Either Nat (Nat, Nat)
|
||||
cRangeToBounds (One x) = Left (cast x)
|
||||
cRangeToBounds All = Right (0, n)
|
||||
|
@ -78,18 +86,15 @@ namespace CoordsRange
|
|||
cRangeToBounds (Bounds x y) = Right (cast x, cast y)
|
||||
|
||||
|
||||
public export
|
||||
newRank : {s : _} -> CoordsRange s -> Nat
|
||||
newRank [] = 0
|
||||
newRank (r :: rs) = case cRangeToBounds r of
|
||||
Left _ => newRank rs
|
||||
Right _ => S (newRank rs)
|
||||
|
||||
export
|
||||
newDim : {d : _} -> (r : CRange d) -> Maybe Nat
|
||||
newDim = map (uncurry $ flip minus) . eitherToMaybe . cRangeToBounds
|
||||
|
||||
||| Calculate the new shape given by a coordinate range.
|
||||
export
|
||||
public export
|
||||
newShape : {s : _} -> (rs : CoordsRange s) -> Vect (newRank rs) Nat
|
||||
newShape [] = []
|
||||
newShape (r :: rs) with (cRangeToBounds r)
|
||||
|
|
Loading…
Reference in a new issue