From 97d1bdb538b923a11fd7258e279d4bc34fca1d6f Mon Sep 17 00:00:00 2001
From: Kiana Sheibani <kiana.a.sheibani@gmail.com>
Date: Tue, 24 May 2022 09:44:13 -0400
Subject: [PATCH] Add operators for indexing

---
 src/Data/NumIdr/Array/Array.idr  | 23 +++++++++++++++++++++++
 src/Data/NumIdr/Array/Coords.idr | 15 ++++++++++-----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/Data/NumIdr/Array/Array.idr b/src/Data/NumIdr/Array/Array.idr
index 2ea76c7..3e16731 100644
--- a/src/Data/NumIdr/Array/Array.idr
+++ b/src/Data/NumIdr/Array/Array.idr
@@ -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
diff --git a/src/Data/NumIdr/Array/Coords.idr b/src/Data/NumIdr/Array/Coords.idr
index 27258a4..954ede0 100644
--- a/src/Data/NumIdr/Array/Coords.idr
+++ b/src/Data/NumIdr/Array/Coords.idr
@@ -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)