{-# 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)
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
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
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
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
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
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_
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