Tweak how runtime-checked indexing works

src/Data/NumIdr/Array/Array.idr

@@ -333,7 +333,7 @@ indexUpdateRange rs f arr = indexSetRange rs (f $ arr !!.. rs) arr
 ||| coordinates are out of bounds.
 export
 indexNB : Vect rk Nat -> Array {rk} s a -> Maybe a
-indexNB is arr = if and $ map delay $ zipWith (<) is (shape arr)
+indexNB is arr = if all id $ zipWith (<) is (shape arr)
                   then Just $ index (getLocation' (strides arr) is) (getPrim arr)
                   else Nothing
 
@@ -345,42 +345,47 @@ export %inline
 (!?) : Array {rk} s a -> Vect rk Nat -> Maybe a
 arr !? is = indexNB is arr
 
-||| Update the entry at the given coordinates using the function. The array is
-||| returned unchanged if the coordinate is out of bounds.
+||| Update the entry at the given coordinates using the function. `Nothing` is
+||| returned if the coordinates are out of bounds.
 export
-indexUpdateNB : Vect rk Nat -> (a -> a) -> Array {rk} s a -> Array s a
+indexUpdateNB : Vect rk Nat -> (a -> a) -> Array {rk} s a -> Maybe (Array s a)
 indexUpdateNB is f (MkArray ord sts s arr) =
-  MkArray ord sts s (updateAt (getLocation' sts is) f arr)
+  if all id $ zipWith (<) is s
+  then Just $ MkArray ord sts s (updateAt (getLocation' sts is) f arr)
+  else Nothing
 
-||| Set the entry at the given coordinates to the given value. The array is
-||| returned unchanged if the coordinate is out of bounds.
+||| Set the entry at the given coordinates using the function. `Nothing` is
+||| returned if the coordinates are out of bounds.
 export
-indexSetNB : Vect rk Nat -> a -> Array {rk} s a -> Array s a
+indexSetNB : Vect rk Nat -> a -> Array {rk} s a -> Maybe (Array s a)
 indexSetNB is = indexUpdateNB is . const
 
 
 ||| Index the array using the given range of coordinates, returning a new array.
-||| Entries outside of the array's bounds are not included.
+||| If any of the given indices are out of bounds, then `Nothing` is returned.
 export
-indexRangeNB : (rs : Vect rk CRangeNB) -> Array s a -> Array (newShape s rs) a
+indexRangeNB : (rs : Vect rk CRangeNB) -> Array s a -> Maybe (Array (newShape s rs) a)
 indexRangeNB {s} rs arr with (viewShape arr)
   _ | Shape s =
-    let ord = getOrder arr
-        sts = calcStrides ord s'
-    in  MkArray ord sts s'
-          (unsafeFromIns (product s') $
-            map (\(is,is') => (getLocation' sts is',
-                                index (getLocation' (strides arr) is) (getPrim arr)))
-            $ getCoordsList s rs)
+    if validateShape s rs
+    then
+      let ord = getOrder arr
+          sts = calcStrides ord s'
+      in  Just $ MkArray ord sts s'
+            (unsafeFromIns (product s') $
+              map (\(is,is') => (getLocation' sts is',
+                                  index (getLocation' (strides arr) is) (getPrim arr)))
+              $ getCoordsList s rs)
+    else Nothing
   where s' : Vect ? Nat
         s' = newShape s rs
 
 ||| Index the array using the given range of coordinates, returning a new array.
-||| Entries outside of the array's bounds are not included.
+||| If any of the given indices are out of bounds, then `Nothing` is returned.
 |||
 ||| This is the operator form of `indexRangeNB`.
 export %inline
-(!?..) : Array s a -> (rs : Vect rk CRangeNB) -> Array (newShape s rs) a
+(!?..) : Array s a -> (rs : Vect rk CRangeNB) -> Maybe (Array (newShape s rs) a)
 arr !?.. rs = indexRangeNB rs arr
 
 

src/Data/NumIdr/Array/Coords.idr

@@ -121,7 +121,7 @@ namespace Strict
   cRangeToList (StartBound x) = Right $ range (cast x) n
   cRangeToList (EndBound x) = Right $ range 0 (cast x)
   cRangeToList (Bounds x y) = Right $ range (cast x) (cast y)
-  cRangeToList (Indices xs) = Right $ map cast xs
+  cRangeToList (Indices xs) = Right $ map cast $ nub xs
   cRangeToList (Filter p) = Right $ map cast $ filter p $ toList Fin.range
 
 
@@ -163,15 +163,28 @@ namespace NB
   cRangeNBToList s All = range 0 s
   cRangeNBToList s (StartBound x) = range x s
   cRangeNBToList s (EndBound x) = range 0 x
-  cRangeNBToList s (Bounds x y) = range x (minimum y s)
-  cRangeNBToList s (Indices xs) = filter (<s) xs
+  cRangeNBToList s (Bounds x y) = range x y
+  cRangeNBToList s (Indices xs) = nub xs
   cRangeNBToList s (Filter p) = filter p $ range 0 s
 
+  export
+  validateCRangeNB : Nat -> CRangeNB -> Bool
+  validateCRangeNB s All = True
+  validateCRangeNB s (StartBound x) = x < s
+  validateCRangeNB s (EndBound x) = x <= s
+  validateCRangeNB s (Bounds x y) = x < s && y <= s
+  validateCRangeNB s (Indices xs) = all (<s) xs
+  validateCRangeNB s (Filter p) = True
+
   ||| Calculate the new shape given by a coordinate range.
   export
   newShape : Vect rk Nat -> Vect rk CRangeNB -> Vect rk Nat
   newShape = zipWith (length .: cRangeNBToList)
 
+  export
+  validateShape : Vect rk Nat -> Vect rk CRangeNB -> Bool
+  validateShape = all id .: zipWith validateCRangeNB
+
   export
   getNewPos : Vect rk Nat -> Vect rk CRangeNB -> Vect rk Nat -> Vect rk Nat
   getNewPos = zipWith3 (\d,r,i => assert_total $