Refactor Data.NumIdr.Array.Coords

This commit is contained in:
Kiana Sheibani 2022-08-27 19:11:13 -04:00
parent 8f1eef25dc
commit 1d6b9d3be9
Signed by: toki
GPG key ID: 6CB106C25E86A9F7

View file

@ -1,13 +1,17 @@
module Data.NumIdr.Array.Coords module Data.NumIdr.Array.Coords
import Data.Either import Data.Either
import Data.List
import Data.Vect import Data.Vect
import Data.NP import Data.NP
%default total %default total
namespace Coords --------------------------------------------------------------------------------
-- Array coordinate types
--------------------------------------------------------------------------------
||| A type-safe coordinate system for an array. The coordinates are ||| A type-safe coordinate system for an array. The coordinates are
||| values of `Fin dim`, where `dim` is the dimension of each axis. ||| values of `Fin dim`, where `dim` is the dimension of each axis.
@ -54,7 +58,27 @@ namespace Coords
getLocation sts is = getLocation' sts (toNB is) getLocation sts is = getLocation' sts (toNB is)
namespace CoordsRange --------------------------------------------------------------------------------
-- Array coordinate types
--------------------------------------------------------------------------------
-- A Nat-based range function with better semantics
public export
range : Nat -> Nat -> List Nat
range x y = if x < y then assert_total $ takeBefore (>= y) (countFrom x S)
else []
-- helpful theorems for working with ranges
export
rangeLen : (x,y : Nat) -> length (range x y) = minus y x
rangeLen x y = believe_me $ Refl {x = minus y x}
export
rangeLenZ : (x : Nat) -> length (range 0 x) = x
rangeLenZ x = rangeLen 0 x `trans` minusZeroRight x
public export public export
data CRange : Nat -> Type where data CRange : Nat -> Type where
@ -63,9 +87,10 @@ namespace CoordsRange
StartBound : Fin (S n) -> CRange n StartBound : Fin (S n) -> CRange n
EndBound : Fin (S n) -> CRange n EndBound : Fin (S n) -> CRange n
Bounds : Fin (S n) -> Fin (S n) -> CRange n Bounds : Fin (S n) -> Fin (S n) -> CRange n
Indices : List (Fin n) -> CRange n
Filter : (Fin n -> Bool) -> CRange n
infix 0 ... infix 0 ...
public export public export
(...) : Fin (S n) -> Fin (S n) -> CRange n (...) : Fin (S n) -> Fin (S n) -> CRange n
(...) = Bounds (...) = Bounds
@ -76,18 +101,20 @@ namespace CoordsRange
CoordsRange = NP CRange CoordsRange = NP CRange
public export public export
cRangeToBounds : {n : Nat} -> CRange n -> Either Nat (Nat, Nat) cRangeToList : {n : Nat} -> CRange n -> Either Nat (List Nat)
cRangeToBounds (One x) = Left (cast x) cRangeToList (One x) = Left (cast x)
cRangeToBounds All = Right (0, n) cRangeToList All = Right $ range 0 n
cRangeToBounds (StartBound x) = Right (cast x, n) cRangeToList (StartBound x) = Right $ range (cast x) n
cRangeToBounds (EndBound x) = Right (0, cast x) cRangeToList (EndBound x) = Right $ range 0 (cast x)
cRangeToBounds (Bounds x y) = Right (cast x, cast y) cRangeToList (Bounds x y) = Right $ range (cast x) (cast y)
cRangeToList (Indices xs) = Right $ map cast xs
cRangeToList (Filter p) = Right $ map cast $ filter p $ toList Fin.range
public export public export
newRank : {s : _} -> CoordsRange s -> Nat newRank : {s : _} -> CoordsRange s -> Nat
newRank [] = 0 newRank [] = 0
newRank (r :: rs) = case cRangeToBounds r of newRank (r :: rs) = case cRangeToList r of
Left _ => newRank rs Left _ => newRank rs
Right _ => S (newRank rs) Right _ => S (newRank rs)
@ -95,16 +122,17 @@ namespace CoordsRange
public export public export
newShape : {s : _} -> (rs : CoordsRange s) -> Vect (newRank rs) Nat newShape : {s : _} -> (rs : CoordsRange s) -> Vect (newRank rs) Nat
newShape [] = [] newShape [] = []
newShape (r :: rs) with (cRangeToBounds r) newShape (r :: rs) with (cRangeToList r)
newShape (r :: rs) | Left _ = newShape rs newShape (r :: rs) | Left _ = newShape rs
newShape (r :: rs) | Right (x,y) = minus y x :: newShape rs newShape (r :: rs) | Right xs = length xs :: newShape rs
getNewPos : {s : _} -> (rs : CoordsRange {rk} s) -> Vect rk Nat -> Vect (newRank rs) Nat getNewPos : {s : _} -> (rs : CoordsRange {rk} s) -> Vect rk Nat -> Vect (newRank rs) Nat
getNewPos [] [] = [] getNewPos [] [] = []
getNewPos (r :: rs) (i :: is) with (cRangeToBounds r) getNewPos (r :: rs) (i :: is) with (cRangeToList r)
_ | Left _ = getNewPos rs is _ | Left _ = getNewPos rs is
_ | Right (x, _) = minus i x :: getNewPos rs is _ | Right xs = cast (assert_total $ case findIndex (==i) xs of Just x => x)
:: getNewPos rs is
export export
getCoordsList : {s : Vect rk Nat} -> (rs : CoordsRange s) -> List (Vect rk Nat, Vect (newRank rs) Nat) getCoordsList : {s : Vect rk Nat} -> (rs : CoordsRange s) -> List (Vect rk Nat, Vect (newRank rs) Nat)
@ -112,6 +140,4 @@ namespace CoordsRange
where where
go : {0 rk : _} -> {s : Vect rk Nat} -> CoordsRange s -> List (Vect rk Nat) go : {0 rk : _} -> {s : Vect rk Nat} -> CoordsRange s -> List (Vect rk Nat)
go [] = pure [] go [] = pure []
go (r :: rs) = [| (case cRangeToBounds r of go (r :: rs) = [| (either pure id (cRangeToList r)) :: go rs |]
Left x => pure x
Right (x,y) => [x,S x..pred y]) :: go rs |]