Refactor Data.NumIdr.Array.Coords

src/Data/NumIdr/Array/Coords.idr

@@ -1,13 +1,17 @@
 module Data.NumIdr.Array.Coords
 
 import Data.Either
+import Data.List
 import Data.Vect
 import Data.NP
 
 %default total
 
 
-namespace Coords
+--------------------------------------------------------------------------------
+-- Array coordinate types
+--------------------------------------------------------------------------------
+
 
 ||| A type-safe coordinate system for an array. The coordinates are
 ||| 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)
 
 
-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
 data CRange : Nat -> Type where
@@ -63,9 +87,10 @@ namespace CoordsRange
   StartBound : Fin (S n) -> CRange n
   EndBound : 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 ...
-
 public export
 (...) : Fin (S n) -> Fin (S n) -> CRange n
 (...) = Bounds
@@ -76,18 +101,20 @@ namespace CoordsRange
 CoordsRange = NP CRange
 
 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
 newRank : {s : _} -> CoordsRange s -> Nat
 newRank [] = 0
-  newRank (r :: rs) = case cRangeToBounds r of
+newRank (r :: rs) = case cRangeToList r of
                       Left _ => newRank rs
                       Right _ => S (newRank rs)
 
@@ -95,16 +122,17 @@ namespace CoordsRange
 public export
 newShape : {s : _} -> (rs : CoordsRange s) -> Vect (newRank rs) Nat
 newShape [] = []
-  newShape (r :: rs) with (cRangeToBounds r)
+newShape (r :: rs) with (cRangeToList r)
   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 [] [] = []
-  getNewPos (r :: rs) (i :: is) with (cRangeToBounds r)
+getNewPos (r :: rs) (i :: is) with (cRangeToList r)
   _ | 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
 getCoordsList : {s : Vect rk Nat} -> (rs : CoordsRange s) -> List (Vect rk Nat, Vect (newRank rs) Nat)
@@ -112,6 +140,4 @@ namespace CoordsRange
   where
     go : {0 rk : _} -> {s : Vect rk Nat} -> CoordsRange s -> List (Vect rk Nat)
     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 |]