{-# LANGUAGE TypeOperators #-}
module Numeric.BLAS.Slice where

import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.Array.Comfort.Boxed as BoxedArray
import Data.Array.Comfort.Shape ((::+)((::+)))

import qualified Data.Traversable as Trav
import qualified Data.List as List
import Data.Map (Map)
import Data.Bool.HT (if')

import Text.Printf (printf)


{- $setup
>>> :set -XTypeOperators
>>>
>>> import qualified Numeric.BLAS.Slice as Slice
>>> import Test.Slice (shapeInt)
>>>
>>> import qualified Data.Array.Comfort.Boxed as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> import qualified Data.Map as Map
>>> import Data.Array.Comfort.Shape ((::+)((::+)))
>>> import Data.Array.Comfort.Boxed ((!))
>>>
>>> import Control.Applicative (liftA2, liftA3, pure)
>>>
>>> import qualified Test.QuickCheck as QC
>>>
>>> genShape :: QC.Gen (Shape.Range Int)
>>> genShape =
>>>    liftA2
>>>       (\m n -> Shape.Range (min m n) (max m n))
>>>       QC.arbitrary QC.arbitrary
>>>
>>> genAppend :: QC.Gen (Shape.Range Int ::+ Shape.Range Int)
>>> genAppend = liftA2 (::+) genShape genShape
>>>
>>> genSlice :: sh -> QC.Gen (Slice.T sh)
>>> genSlice sh =
>>>    liftA3 Slice.Cons (QC.choose (0,100)) (QC.choose (1,100)) (pure sh)
>>>
>>> genSlice2 :: shA -> shB -> QC.Gen (Slice.T shA, Slice.T shB)
>>> genSlice2 shA shB = do
>>>    s <- QC.choose (0,100)
>>>    k <- QC.choose (1,100)
>>>    return (Slice.Cons s k shA, Slice.Cons s k shB)
>>>
>>> type ShapeInt = Shape.ZeroBased Int
>>>
>>> mapShape :: (shA -> shB) -> Slice.T shA -> Slice.T shB
>>> mapShape f (Slice.Cons s k sh) = Slice.Cons s k (f sh)
>>>
>>> toShapeInt :: (Shape.C sh) => Slice.T sh -> Slice.T ShapeInt
>>> toShapeInt = mapShape (shapeInt . Shape.size)
-}


data T sh = Cons {forall sh. T sh -> Int
start, forall sh. T sh -> Int
skip :: Int, forall sh. T sh -> sh
shape :: sh}
   deriving (T sh -> T sh -> Bool
(T sh -> T sh -> Bool) -> (T sh -> T sh -> Bool) -> Eq (T sh)
forall sh. Eq sh => T sh -> T sh -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall sh. Eq sh => T sh -> T sh -> Bool
== :: T sh -> T sh -> Bool
$c/= :: forall sh. Eq sh => T sh -> T sh -> Bool
/= :: T sh -> T sh -> Bool
Eq, Int -> T sh -> ShowS
[T sh] -> ShowS
T sh -> [Char]
(Int -> T sh -> ShowS)
-> (T sh -> [Char]) -> ([T sh] -> ShowS) -> Show (T sh)
forall sh. Show sh => Int -> T sh -> ShowS
forall sh. Show sh => [T sh] -> ShowS
forall sh. Show sh => T sh -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall sh. Show sh => Int -> T sh -> ShowS
showsPrec :: Int -> T sh -> ShowS
$cshow :: forall sh. Show sh => T sh -> [Char]
show :: T sh -> [Char]
$cshowList :: forall sh. Show sh => [T sh] -> ShowS
showList :: [T sh] -> ShowS
Show)

fromShape :: (Shape.C sh) => sh -> T sh
fromShape :: forall sh. C sh => sh -> T sh
fromShape = Int -> Int -> sh -> T sh
forall sh. Int -> Int -> sh -> T sh
Cons Int
0 Int
1


row ::
   (Shape.Indexed sh0, Shape.C sh1) => Shape.Index sh0 -> T (sh0,sh1) -> T sh1
row :: forall sh0 sh1.
(Indexed sh0, C sh1) =>
Index sh0 -> T (sh0, sh1) -> T sh1
row Index sh0
ix0 (Cons Int
s Int
k (sh0
sh0,sh1
sh1)) =
   Int -> Int -> sh1 -> T sh1
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* sh0 -> Index sh0 -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh0
sh0 Index sh0
ix0 Int -> Int -> Int
forall a. Num a => a -> a -> a
* sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1) Int
k sh1
sh1

column ::
   (Shape.C sh0, Shape.Indexed sh1) => Shape.Index sh1 -> T (sh0,sh1) -> T sh0
column :: forall sh0 sh1.
(C sh0, Indexed sh1) =>
Index sh1 -> T (sh0, sh1) -> T sh0
column Index sh1
ix1 (Cons Int
s Int
k (sh0
sh0,sh1
sh1)) =
   let (Int
size1, Index sh1 -> Int
offset1) = sh1 -> (Int, Index sh1 -> Int)
forall sh. Indexed sh => sh -> (Int, Index sh -> Int)
Shape.sizeOffset sh1
sh1
   in Int -> Int -> sh0 -> T sh0
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Index sh1 -> Int
offset1 Index sh1
ix1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size1) sh0
sh0

{- |
prop> QC.forAll (QC.choose (1,100)) $ \numRows -> QC.forAll (QC.choose (0,100)) $ \numColumns -> QC.forAll (genSlice (shapeInt numRows, shapeInt numColumns)) $ \slice -> QC.forAll (QC.elements $ Shape.indices $ shapeInt numRows) $ \ix -> Slice.row ix slice == Slice.rowArray slice ! ix
-}
rowArray ::
   (Shape.Indexed sh0, Shape.C sh1) =>
   T (sh0,sh1) -> BoxedArray.Array sh0 (T sh1)
rowArray :: forall sh0 sh1.
(Indexed sh0, C sh1) =>
T (sh0, sh1) -> Array sh0 (T sh1)
rowArray (Cons Int
s Int
k (sh0
sh0,sh1
sh1)) =
   let step :: Int
step = sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k
   in sh0 -> [T sh1] -> Array sh0 (T sh1)
forall sh a. C sh => sh -> [a] -> Array sh a
BoxedArray.fromList sh0
sh0 ([T sh1] -> Array sh0 (T sh1)) -> [T sh1] -> Array sh0 (T sh1)
forall a b. (a -> b) -> a -> b
$
      (Int -> T sh1) -> [Int] -> [T sh1]
forall a b. (a -> b) -> [a] -> [b]
List.map (\Int
si -> Int -> Int -> sh1 -> T sh1
forall sh. Int -> Int -> sh -> T sh
Cons Int
si Int
k sh1
sh1) ([Int] -> [T sh1]) -> [Int] -> [T sh1]
forall a b. (a -> b) -> a -> b
$
      Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take (sh0 -> Int
forall sh. C sh => sh -> Int
Shape.size sh0
sh0) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int
stepInt -> Int -> Int
forall a. Num a => a -> a -> a
+) Int
s

{- |
prop> QC.forAll (QC.choose (0,100)) $ \numRows -> QC.forAll (QC.choose (1,100)) $ \numColumns -> QC.forAll (genSlice (shapeInt numRows, shapeInt numColumns)) $ \slice -> QC.forAll (QC.elements $ Shape.indices $ shapeInt numColumns) $ \ix -> Slice.column ix slice == Slice.columnArray slice ! ix
-}
columnArray ::
   (Shape.C sh0, Shape.Indexed sh1) =>
   T (sh0,sh1) -> BoxedArray.Array sh1 (T sh0)
columnArray :: forall sh0 sh1.
(C sh0, Indexed sh1) =>
T (sh0, sh1) -> Array sh1 (T sh0)
columnArray (Cons Int
s Int
k (sh0
sh0,sh1
sh1)) =
   let step :: Int
step = sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k
   in sh1 -> [T sh0] -> Array sh1 (T sh0)
forall sh a. C sh => sh -> [a] -> Array sh a
BoxedArray.fromList sh1
sh1 ([T sh0] -> Array sh1 (T sh0)) -> [T sh0] -> Array sh1 (T sh0)
forall a b. (a -> b) -> a -> b
$
      (Int -> T sh0) -> [Int] -> [T sh0]
forall a b. (a -> b) -> [a] -> [b]
List.map (\Int
si -> Int -> Int -> sh0 -> T sh0
forall sh. Int -> Int -> sh -> T sh
Cons Int
si Int
step sh0
sh0) ([Int] -> [T sh0]) -> [Int] -> [T sh0]
forall a b. (a -> b) -> a -> b
$
      Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take (sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+) Int
s


topSubmatrix ::
   (Shape.C sh, Shape.C sh0, Shape.C sh1) =>
   T (sh0::+sh1, sh) -> T (sh0,sh)
topSubmatrix :: forall sh sh0 sh1.
(C sh, C sh0, C sh1) =>
T (sh0 ::+ sh1, sh) -> T (sh0, sh)
topSubmatrix (Cons Int
s Int
k (sh0
sh0::+sh1
_sh1, sh
sh)) =
   Int -> Int -> (sh0, sh) -> T (sh0, sh)
forall sh. Int -> Int -> sh -> T sh
Cons Int
s Int
k (sh0
sh0,sh
sh)

bottomSubmatrix ::
   (Shape.C sh, Shape.C sh0, Shape.C sh1) =>
   T (sh0::+sh1, sh) -> T (sh1,sh)
bottomSubmatrix :: forall sh sh0 sh1.
(C sh, C sh0, C sh1) =>
T (sh0 ::+ sh1, sh) -> T (sh1, sh)
bottomSubmatrix (Cons Int
s Int
k (sh0
sh0::+sh1
sh1, sh
sh)) =
   Int -> Int -> (sh1, sh) -> T (sh1, sh)
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* sh0 -> Int
forall sh. C sh => sh -> Int
Shape.size sh0
sh0) Int
k (sh1
sh1,sh
sh)


diagonal :: (Shape.C sh) => T (Shape.Square sh) -> T sh
diagonal :: forall sh. C sh => T (Square sh) -> T sh
diagonal (Cons Int
s Int
k (Shape.Square sh
sh)) =
   Int -> Int -> sh -> T sh
forall sh. Int -> Int -> sh -> T sh
Cons Int
s (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) sh
sh

cartesianFromSquare :: T (Shape.Square sh) -> T (sh,sh)
cartesianFromSquare :: forall sh. T (Square sh) -> T (sh, sh)
cartesianFromSquare (Cons Int
s Int
k (Shape.Square sh
sh)) = Int -> Int -> (sh, sh) -> T (sh, sh)
forall sh. Int -> Int -> sh -> T sh
Cons Int
s Int
k (sh
sh,sh
sh)

squareRow ::
   (Shape.Indexed sh) => Shape.Index sh -> T (Shape.Square sh) -> T sh
squareRow :: forall sh. Indexed sh => Index sh -> T (Square sh) -> T sh
squareRow Index sh
ix0 = Index sh -> T (sh, sh) -> T sh
forall sh0 sh1.
(Indexed sh0, C sh1) =>
Index sh0 -> T (sh0, sh1) -> T sh1
row Index sh
ix0 (T (sh, sh) -> T sh)
-> (T (Square sh) -> T (sh, sh)) -> T (Square sh) -> T sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T (Square sh) -> T (sh, sh)
forall sh. T (Square sh) -> T (sh, sh)
cartesianFromSquare

squareColumn ::
   (Shape.Indexed sh) => Shape.Index sh -> T (Shape.Square sh) -> T sh
squareColumn :: forall sh. Indexed sh => Index sh -> T (Square sh) -> T sh
squareColumn Index sh
ix1 = Index sh -> T (sh, sh) -> T sh
forall sh0 sh1.
(C sh0, Indexed sh1) =>
Index sh1 -> T (sh0, sh1) -> T sh0
column Index sh
ix1 (T (sh, sh) -> T sh)
-> (T (Square sh) -> T (sh, sh)) -> T (Square sh) -> T sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T (Square sh) -> T (sh, sh)
forall sh. T (Square sh) -> T (sh, sh)
cartesianFromSquare



plane12 ::
   (Shape.Indexed sh0, Shape.C sh1, Shape.C sh2) =>
   Shape.Index sh0 -> T (sh0,sh1,sh2) -> T (sh1,sh2)
plane12 :: forall sh0 sh1 sh2.
(Indexed sh0, C sh1, C sh2) =>
Index sh0 -> T (sh0, sh1, sh2) -> T (sh1, sh2)
plane12 Index sh0
ix0 (Cons Int
s Int
k (sh0
sh0,sh1
sh1,sh2
sh2)) =
   Int -> Int -> (sh1, sh2) -> T (sh1, sh2)
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* sh0 -> Index sh0 -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh0
sh0 Index sh0
ix0 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (sh1, sh2) -> Int
forall sh. C sh => sh -> Int
Shape.size (sh1
sh1,sh2
sh2)) Int
k (sh1
sh1,sh2
sh2)

plane01 ::
   (Shape.C sh0, Shape.C sh1, Shape.Indexed sh2) =>
   Shape.Index sh2 -> T (sh0,sh1,sh2) -> T (sh0,sh1)
plane01 :: forall sh0 sh1 sh2.
(C sh0, C sh1, Indexed sh2) =>
Index sh2 -> T (sh0, sh1, sh2) -> T (sh0, sh1)
plane01 Index sh2
ix2 (Cons Int
s Int
k (sh0
sh0,sh1
sh1,sh2
sh2)) =
   let (Int
size2, Index sh2 -> Int
offset2) = sh2 -> (Int, Index sh2 -> Int)
forall sh. Indexed sh => sh -> (Int, Index sh -> Int)
Shape.sizeOffset sh2
sh2
   in Int -> Int -> (sh0, sh1) -> T (sh0, sh1)
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Index sh2 -> Int
offset2 Index sh2
ix2) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size2) (sh0
sh0,sh1
sh1)

column2of3 ::
   (Shape.Indexed sh0, Shape.Indexed sh1, Shape.C sh2) =>
   Shape.Index sh0 -> Shape.Index sh1 -> T (sh0,sh1,sh2) -> T sh2
column2of3 :: forall sh0 sh1 sh2.
(Indexed sh0, Indexed sh1, C sh2) =>
Index sh0 -> Index sh1 -> T (sh0, sh1, sh2) -> T sh2
column2of3 Index sh0
ix0 Index sh1
ix1 = Index sh1 -> T (sh1, sh2) -> T sh2
forall sh0 sh1.
(Indexed sh0, C sh1) =>
Index sh0 -> T (sh0, sh1) -> T sh1
row Index sh1
ix1 (T (sh1, sh2) -> T sh2)
-> (T (sh0, sh1, sh2) -> T (sh1, sh2))
-> T (sh0, sh1, sh2)
-> T sh2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index sh0 -> T (sh0, sh1, sh2) -> T (sh1, sh2)
forall sh0 sh1 sh2.
(Indexed sh0, C sh1, C sh2) =>
Index sh0 -> T (sh0, sh1, sh2) -> T (sh1, sh2)
plane12 Index sh0
ix0

column1of3 ::
   (Shape.Indexed sh0, Shape.C sh1, Shape.Indexed sh2) =>
   Shape.Index sh0 -> Shape.Index sh2 -> T (sh0,sh1,sh2) -> T sh1
column1of3 :: forall sh0 sh1 sh2.
(Indexed sh0, C sh1, Indexed sh2) =>
Index sh0 -> Index sh2 -> T (sh0, sh1, sh2) -> T sh1
column1of3 Index sh0
ix0 Index sh2
ix2 = Index sh2 -> T (sh1, sh2) -> T sh1
forall sh0 sh1.
(C sh0, Indexed sh1) =>
Index sh1 -> T (sh0, sh1) -> T sh0
column Index sh2
ix2 (T (sh1, sh2) -> T sh1)
-> (T (sh0, sh1, sh2) -> T (sh1, sh2))
-> T (sh0, sh1, sh2)
-> T sh1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index sh0 -> T (sh0, sh1, sh2) -> T (sh1, sh2)
forall sh0 sh1 sh2.
(Indexed sh0, C sh1, C sh2) =>
Index sh0 -> T (sh0, sh1, sh2) -> T (sh1, sh2)
plane12 Index sh0
ix0

column0of3 ::
   (Shape.C sh0, Shape.Indexed sh1, Shape.Indexed sh2) =>
   Shape.Index sh1 -> Shape.Index sh2 -> T (sh0,sh1,sh2) -> T sh0
column0of3 :: forall sh0 sh1 sh2.
(C sh0, Indexed sh1, Indexed sh2) =>
Index sh1 -> Index sh2 -> T (sh0, sh1, sh2) -> T sh0
column0of3 Index sh1
ix1 Index sh2
ix2 = Index sh1 -> T (sh0, sh1) -> T sh0
forall sh0 sh1.
(C sh0, Indexed sh1) =>
Index sh1 -> T (sh0, sh1) -> T sh0
column Index sh1
ix1 (T (sh0, sh1) -> T sh0)
-> (T (sh0, sh1, sh2) -> T (sh0, sh1))
-> T (sh0, sh1, sh2)
-> T sh0
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index sh2 -> T (sh0, sh1, sh2) -> T (sh0, sh1)
forall sh0 sh1 sh2.
(C sh0, C sh1, Indexed sh2) =>
Index sh2 -> T (sh0, sh1, sh2) -> T (sh0, sh1)
plane01 Index sh2
ix2


left :: (Shape.C sh0, Shape.C sh1) => T (sh0::+sh1) -> T sh0
left :: forall sh0 sh1. (C sh0, C sh1) => T (sh0 ::+ sh1) -> T sh0
left (Cons Int
s Int
k (sh0
sh0::+sh1
_sh1)) = Int -> Int -> sh0 -> T sh0
forall sh. Int -> Int -> sh -> T sh
Cons Int
s Int
k sh0
sh0

right :: (Shape.C sh0, Shape.C sh1) => T (sh0::+sh1) -> T sh1
right :: forall sh0 sh1. (C sh0, C sh1) => T (sh0 ::+ sh1) -> T sh1
right (Cons Int
s Int
k (sh0
sh0::+sh1
sh1)) = Int -> Int -> sh1 -> T sh1
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* sh0 -> Int
forall sh. C sh => sh -> Int
Shape.size sh0
sh0) Int
k sh1
sh1


type ShapeInt = Shape.ZeroBased Int

{- |
prop> QC.forAll (genSlice =<< genAppend) $ \slice -> Slice.take (case Slice.shape slice of (sh::+_) -> Shape.size sh) (toShapeInt slice) == toShapeInt (Slice.left slice)
-}
take :: Int -> T ShapeInt -> T ShapeInt
take :: Int -> T ShapeInt -> T ShapeInt
take Int
m (Cons Int
s Int
k (Shape.ZeroBased Int
n)) =
   if Int
mInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0
      then [Char] -> T ShapeInt
forall a. HasCallStack => [Char] -> a
error ([Char] -> T ShapeInt) -> [Char] -> T ShapeInt
forall a b. (a -> b) -> a -> b
$ [Char] -> Int -> [Char]
forall r. PrintfType r => [Char] -> r
printf [Char]
"Slice.take: negative size %d" Int
m
      else Int -> Int -> ShapeInt -> T ShapeInt
forall sh. Int -> Int -> sh -> T sh
Cons Int
s Int
k (ShapeInt -> T ShapeInt) -> ShapeInt -> T ShapeInt
forall a b. (a -> b) -> a -> b
$ Int -> ShapeInt
forall n. n -> ZeroBased n
Shape.ZeroBased (Int -> ShapeInt) -> Int -> ShapeInt
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m Int
n

{- |
prop> QC.forAll (genSlice =<< genAppend) $ \slice -> Slice.drop (case Slice.shape slice of (sh::+_) -> Shape.size sh) (toShapeInt slice) == toShapeInt (Slice.right slice)
-}
drop :: Int -> T ShapeInt -> T ShapeInt
drop :: Int -> T ShapeInt -> T ShapeInt
drop Int
m (Cons Int
s Int
k (Shape.ZeroBased Int
n)) =
   if Int
mInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0
      then [Char] -> T ShapeInt
forall a. HasCallStack => [Char] -> a
error ([Char] -> T ShapeInt) -> [Char] -> T ShapeInt
forall a b. (a -> b) -> a -> b
$ [Char] -> Int -> [Char]
forall r. PrintfType r => [Char] -> r
printf [Char]
"Slice.drop: negative size %d" Int
m
      else Int -> Int -> ShapeInt -> T ShapeInt
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m) Int
k (ShapeInt -> T ShapeInt) -> ShapeInt -> T ShapeInt
forall a b. (a -> b) -> a -> b
$ Int -> ShapeInt
forall n. n -> ZeroBased n
Shape.ZeroBased (Int -> ShapeInt) -> Int -> ShapeInt
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
m

{- |
@Slice.sub start size@

prop> \(QC.NonNegative n) -> QC.forAll (genSlice $ shapeInt n) $ \slice (QC.NonNegative start) (QC.NonNegative size) -> Slice.sub start size slice == Slice.take size (Slice.drop start slice)
-}
sub :: Int -> Int -> T ShapeInt -> T ShapeInt
sub :: Int -> Int -> T ShapeInt -> T ShapeInt
sub Int
start_ Int
size (Cons Int
s Int
k (Shape.ZeroBased Int
n)) =
   Bool -> T ShapeInt -> T ShapeInt -> T ShapeInt
forall a. Bool -> a -> a -> a
if' (Int
start_Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0) ([Char] -> T ShapeInt
forall a. HasCallStack => [Char] -> a
error ([Char] -> T ShapeInt) -> [Char] -> T ShapeInt
forall a b. (a -> b) -> a -> b
$ [Char] -> Int -> [Char]
forall r. PrintfType r => [Char] -> r
printf [Char]
"Slice.sub: negative start %d" Int
start_) (T ShapeInt -> T ShapeInt) -> T ShapeInt -> T ShapeInt
forall a b. (a -> b) -> a -> b
$
   Bool -> T ShapeInt -> T ShapeInt -> T ShapeInt
forall a. Bool -> a -> a -> a
if' (Int
size  Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0) ([Char] -> T ShapeInt
forall a. HasCallStack => [Char] -> a
error ([Char] -> T ShapeInt) -> [Char] -> T ShapeInt
forall a b. (a -> b) -> a -> b
$ [Char] -> Int -> [Char]
forall r. PrintfType r => [Char] -> r
printf [Char]
"Slice.sub: negative size %d" Int
size) (T ShapeInt -> T ShapeInt) -> T ShapeInt -> T ShapeInt
forall a b. (a -> b) -> a -> b
$
   Int -> Int -> ShapeInt -> T ShapeInt
forall sh. Int -> Int -> sh -> T sh
Cons (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
start_) Int
k (ShapeInt -> T ShapeInt) -> ShapeInt -> T ShapeInt
forall a b. (a -> b) -> a -> b
$ Int -> ShapeInt
forall n. n -> ZeroBased n
Shape.ZeroBased (Int -> ShapeInt) -> Int -> ShapeInt
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
size (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
start_


{- |
prop> QC.forAll (fmap shapeInt $ QC.choose (0,100)) $ \shapeA -> QC.forAll (fmap shapeInt $ QC.choose (0,100)) $ \shapeB -> QC.forAll (fmap shapeInt $ QC.choose (0,100)) $ \shapeC -> QC.forAll (genSlice2 (Map.fromList $ ('a', shapeA) : ('b', shapeB) : ('c', shapeC) : []) (shapeA ::+ shapeB ::+ shapeC)) $ \(sliceMap, sliceParted) -> Slice.map sliceMap Map.! 'b' == Slice.left (Slice.right sliceParted)

prop> QC.forAll (QC.choose (0,100)) $ \numRows -> QC.forAll (QC.choose (0,100)) $ \numColumns -> let rowShape = shapeInt numRows; columnShape = shapeInt numColumns; mapShape_ = Map.fromList $ map (\k -> (k, columnShape)) (Shape.indices rowShape) in QC.forAll (genSlice2 mapShape_ (rowShape, columnShape)) $ \(sliceMap, sliceMatrix) -> Map.toAscList (Slice.map sliceMap) == Array.toAssociations (Slice.rowArray sliceMatrix)
-}
map :: (Shape.C sh) => T (Map k sh) -> Map k (T sh)
map :: forall sh k. C sh => T (Map k sh) -> Map k (T sh)
map (Cons Int
s Int
k Map k sh
m) =
   (Int, Map k (T sh)) -> Map k (T sh)
forall a b. (a, b) -> b
snd ((Int, Map k (T sh)) -> Map k (T sh))
-> (Int, Map k (T sh)) -> Map k (T sh)
forall a b. (a -> b) -> a -> b
$
   (Int -> sh -> (Int, T sh))
-> Int -> Map k sh -> (Int, Map k (T sh))
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
Trav.mapAccumL
      (\Int
offset sh
sh -> (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k, Int -> Int -> sh -> T sh
forall sh. Int -> Int -> sh -> T sh
Cons Int
offset Int
k sh
sh)) Int
s Map k sh
m