{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_HADDOCK hide #-}
module Data.Array.Accelerate.Representation.Array
where
import Data.Array.Accelerate.Array.Data
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.Representation.Elt
import Data.Array.Accelerate.Representation.Shape hiding ( zip )
import Data.Array.Accelerate.Representation.Type
import Data.List ( intersperse )
import Data.Maybe ( isJust )
import Formatting
import Language.Haskell.TH.Extra
import System.IO.Unsafe
import Text.Show ( showListWith )
import Prelude hiding ( (!!) )
import qualified Data.Vector.Unboxed as U
data Array sh e where
Array :: sh
-> ArrayData e
-> Array sh e
type Segments = Vector
type Scalar = Array DIM0
type Vector = Array DIM1
type Matrix = Array DIM2
data ArrayR a where
ArrayR :: { forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape :: ShapeR sh
, forall sh e. ArrayR (Array sh e) -> TypeR e
arrayRtype :: TypeR e
}
-> ArrayR (Array sh e)
type ArraysR = TupR ArrayR
instance Show (ArrayR a) where
show :: ArrayR a -> String
show (ArrayR ShapeR sh
shR TypeR e
eR) = String
"Array DIM" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (ShapeR sh -> Int
forall sh. ShapeR sh -> Int
rank ShapeR sh
shR) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ TypeR e -> String
forall a. Show a => a -> String
show TypeR e
eR
formatArrayR :: Format r (ArrayR a -> r)
formatArrayR :: forall r a. Format r (ArrayR a -> r)
formatArrayR = (ArrayR a -> Builder) -> Format r (ArrayR a -> r)
forall a r. (a -> Builder) -> Format r (a -> r)
later ((ArrayR a -> Builder) -> Format r (ArrayR a -> r))
-> (ArrayR a -> Builder) -> Format r (ArrayR a -> r)
forall a b. (a -> b) -> a -> b
$ \case
ArrayR ShapeR sh
shR TypeR e
eR -> Format Builder (Int -> TypeR e -> Builder)
-> Int -> TypeR e -> Builder
forall a. Format Builder a -> a
bformat (Format (Int -> TypeR e -> Builder) (Int -> TypeR e -> Builder)
"Array DIM" Format (Int -> TypeR e -> Builder) (Int -> TypeR e -> Builder)
-> Format Builder (Int -> TypeR e -> Builder)
-> Format Builder (Int -> TypeR e -> Builder)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format (TypeR e -> Builder) (Int -> TypeR e -> Builder)
forall a r. Integral a => Format r (a -> r)
int Format (TypeR e -> Builder) (Int -> TypeR e -> Builder)
-> Format Builder (TypeR e -> Builder)
-> Format Builder (Int -> TypeR e -> Builder)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format (TypeR e -> Builder) (TypeR e -> Builder)
" " Format (TypeR e -> Builder) (TypeR e -> Builder)
-> Format Builder (TypeR e -> Builder)
-> Format Builder (TypeR e -> Builder)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format Builder (TypeR e -> Builder)
forall r a. Format r (TypeR a -> r)
formatTypeR) (ShapeR sh -> Int
forall sh. ShapeR sh -> Int
rank ShapeR sh
shR) TypeR e
eR
formatArraysR :: Format r (TupR ArrayR e -> r)
formatArraysR :: forall r e. Format r (TupR ArrayR e -> r)
formatArraysR = (TupR ArrayR e -> Builder) -> Format r (TupR ArrayR e -> r)
forall a r. (a -> Builder) -> Format r (a -> r)
later ((TupR ArrayR e -> Builder) -> Format r (TupR ArrayR e -> r))
-> (TupR ArrayR e -> Builder) -> Format r (TupR ArrayR e -> r)
forall a b. (a -> b) -> a -> b
$ \case
TupR ArrayR e
TupRunit -> Builder
"()"
TupRsingle ArrayR e
aR -> Format Builder (ArrayR e -> Builder) -> ArrayR e -> Builder
forall a. Format Builder a -> a
bformat Format Builder (ArrayR e -> Builder)
forall r a. Format r (ArrayR a -> r)
formatArrayR ArrayR e
aR
TupRpair TupR ArrayR a1
aR1 TupR ArrayR b
aR2 -> Format Builder (TupR ArrayR a1 -> TupR ArrayR b -> Builder)
-> TupR ArrayR a1 -> TupR ArrayR b -> Builder
forall a. Format Builder a -> a
bformat (Format Builder (TupR ArrayR a1 -> TupR ArrayR b -> Builder)
-> Format Builder (TupR ArrayR a1 -> TupR ArrayR b -> Builder)
forall r a. Format r a -> Format r a
parenthesised (Format
(TupR ArrayR b -> Builder)
(TupR ArrayR a1 -> TupR ArrayR b -> Builder)
forall r e. Format r (TupR ArrayR e -> r)
formatArraysR Format
(TupR ArrayR b -> Builder)
(TupR ArrayR a1 -> TupR ArrayR b -> Builder)
-> Format Builder (TupR ArrayR b -> Builder)
-> Format Builder (TupR ArrayR a1 -> TupR ArrayR b -> Builder)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format (TupR ArrayR b -> Builder) (TupR ArrayR b -> Builder)
"," Format (TupR ArrayR b -> Builder) (TupR ArrayR b -> Builder)
-> Format Builder (TupR ArrayR b -> Builder)
-> Format Builder (TupR ArrayR b -> Builder)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format Builder (TupR ArrayR b -> Builder)
forall r e. Format r (TupR ArrayR e -> r)
formatArraysR)) TupR ArrayR a1
aR1 TupR ArrayR b
aR2
showArraysR :: ArraysR a -> ShowS
showArraysR :: forall a. ArraysR a -> ShowS
showArraysR = ArraysR a -> ShowS
forall a. Show a => a -> ShowS
shows
arraysRarray :: ShapeR sh -> TypeR e -> ArraysR (Array sh e)
arraysRarray :: forall sh e. ShapeR sh -> TypeR e -> ArraysR (Array sh e)
arraysRarray ShapeR sh
shR TypeR e
eR = ArrayR (Array sh e) -> TupR ArrayR (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shR TypeR e
eR)
arraysRpair :: ArrayR a -> ArrayR b -> ArraysR (((), a), b)
arraysRpair :: forall a b. ArrayR a -> ArrayR b -> ArraysR (((), a), b)
arraysRpair ArrayR a
a ArrayR b
b = TupR ArrayR ()
forall (s :: * -> *). TupR s ()
TupRunit TupR ArrayR () -> TupR ArrayR a -> TupR ArrayR ((), a)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ArrayR a -> TupR ArrayR a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR a
a TupR ArrayR ((), a) -> TupR ArrayR b -> TupR ArrayR (((), a), b)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ArrayR b -> TupR ArrayR b
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR b
b
allocateArray :: ArrayR (Array sh e) -> sh -> IO (Array sh e)
allocateArray :: forall sh e. ArrayR (Array sh e) -> sh -> IO (Array sh e)
allocateArray (ArrayR ShapeR sh
shR TypeR e
eR) sh
sh = do
GArrayDataR UniqueArray e
adata <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
eR (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shR sh
sh
sh)
Array sh e -> IO (Array sh e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array sh e -> IO (Array sh e)) -> Array sh e -> IO (Array sh e)
forall a b. (a -> b) -> a -> b
$! sh -> GArrayDataR UniqueArray e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh GArrayDataR UniqueArray e
adata
fromFunction :: ArrayR (Array sh e) -> sh -> (sh -> e) -> Array sh e
fromFunction :: forall sh e. ArrayR (Array sh e) -> sh -> (sh -> e) -> Array sh e
fromFunction ArrayR (Array sh e)
repr sh
sh sh -> e
f = IO (Array sh e) -> Array sh e
forall a. IO a -> a
unsafePerformIO (IO (Array sh e) -> Array sh e) -> IO (Array sh e) -> Array sh e
forall a b. (a -> b) -> a -> b
$! ArrayR (Array sh e) -> sh -> (sh -> IO e) -> IO (Array sh e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> IO e) -> IO (Array sh e)
fromFunctionM ArrayR (Array sh e)
repr sh
sh (e -> IO e
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (e -> IO e) -> (sh -> e) -> sh -> IO e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. sh -> e
f)
fromFunctionM :: ArrayR (Array sh e) -> sh -> (sh -> IO e) -> IO (Array sh e)
fromFunctionM :: forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> IO e) -> IO (Array sh e)
fromFunctionM (ArrayR ShapeR sh
shR TypeR e
eR) sh
sh sh -> IO e
f = do
let !n :: Int
n = ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shR sh
sh
sh
GArrayDataR UniqueArray e
arr <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
eR Int
n
let write :: Int -> IO ()
write !Int
i
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Bool
otherwise = do
e
v <- sh -> IO e
f (ShapeR sh -> sh -> Int -> sh
forall sh. HasCallStack => ShapeR sh -> sh -> Int -> sh
fromIndex ShapeR sh
ShapeR sh
shR sh
sh Int
i)
TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
eR GArrayDataR UniqueArray e
MutableArrayData e
arr Int
i e
e
v
Int -> IO ()
write (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
Int -> IO ()
write Int
0
Array sh e -> IO (Array sh e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array sh e -> IO (Array sh e)) -> Array sh e -> IO (Array sh e)
forall a b. (a -> b) -> a -> b
$! GArrayDataR UniqueArray e
arr GArrayDataR UniqueArray e -> Array sh e -> Array sh e
forall a b. a -> b -> b
`seq` sh -> GArrayDataR UniqueArray e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh GArrayDataR UniqueArray e
arr
fromList :: forall sh e. ArrayR (Array sh e) -> sh -> [e] -> Array sh e
fromList :: forall sh e. ArrayR (Array sh e) -> sh -> [e] -> Array sh e
fromList (ArrayR ShapeR sh
shR TypeR e
eR) sh
sh [e]
xs = ArrayData e
adata ArrayData e -> Array sh e -> Array sh e
forall a b. a -> b -> b
`seq` sh -> ArrayData e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh ArrayData e
adata
where
!n :: Int
n = ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shR sh
sh
sh
(ArrayData e
adata, e
_) = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @e (IO (ArrayData e, e) -> (ArrayData e, e))
-> IO (ArrayData e, e) -> (ArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
ArrayData e
arr <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
eR Int
n
let go :: Int -> [e] -> IO ()
go !Int
i [e]
_ | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
go !Int
i (e
v:[e]
vs) = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
eR ArrayData e
MutableArrayData e
arr Int
i e
v IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> [e] -> IO ()
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [e]
vs
go Int
_ [] = String -> IO ()
forall a. HasCallStack => String -> a
error String
"Data.Array.Accelerate.fromList: not enough input data"
Int -> [e] -> IO ()
go Int
0 [e]
[e]
xs
(ArrayData e, e) -> IO (ArrayData e, e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayData e
arr, e
forall a. HasCallStack => a
undefined)
toList :: ArrayR (Array sh e) -> Array sh e -> [e]
toList :: forall sh e. ArrayR (Array sh e) -> Array sh e -> [e]
toList (ArrayR ShapeR sh
shR TypeR e
eR) (Array sh
sh ArrayData e
adata) = Int -> [e]
go Int
0
where
!n :: Int
n = ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shR sh
sh
sh
go :: Int -> [e]
go !Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = []
| Bool
otherwise = TypeR e -> ArrayData e -> Int -> e
forall e. TupR ScalarType e -> ArrayData e -> Int -> e
indexArrayData TypeR e
eR ArrayData e
ArrayData e
adata Int
i e -> [e] -> [e]
forall a. a -> [a] -> [a]
: Int -> [e]
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
concatVectors :: forall e. TypeR e -> [Vector e] -> Vector e
concatVectors :: forall e. TypeR e -> [Vector e] -> Vector e
concatVectors TypeR e
tR [Vector e]
vs = ArrayData e
adata ArrayData e -> Vector e -> Vector e
forall a b. a -> b -> b
`seq` DIM1 -> ArrayData e -> Vector e
forall sh e. sh -> ArrayData e -> Array sh e
Array ((), Int
len) ArrayData e
adata
where
offsets :: [Int]
offsets = (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ((Vector e -> Int) -> [Vector e] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (ShapeR DIM1 -> DIM1 -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR DIM1
dim1 (DIM1 -> Int) -> (Vector e -> DIM1) -> Vector e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector e -> DIM1
forall sh e. Array sh e -> sh
shape) [Vector e]
vs)
len :: Int
len = [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
offsets
(ArrayData e
adata, e
_) = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @e (IO (ArrayData e, e) -> (ArrayData e, e))
-> IO (ArrayData e, e) -> (ArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
ArrayData e
arr <- TypeR e -> Int -> IO (ArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tR Int
len
[IO ()] -> IO ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ [ TypeR e -> ArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tR ArrayData e
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) (TypeR e -> ArrayData e -> Int -> e
forall e. TupR ScalarType e -> ArrayData e -> Int -> e
indexArrayData TypeR e
tR ArrayData e
ad Int
i)
| (Array ((), Int
n) ArrayData e
ad, Int
k) <- [Vector e]
vs [Vector e] -> [Int] -> [(Vector e, Int)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Int]
offsets
, Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ]
(ArrayData e, e) -> IO (ArrayData e, e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayData e
arr, e
forall a. HasCallStack => a
undefined)
shape :: Array sh e -> sh
shape :: forall sh e. Array sh e -> sh
shape (Array sh
sh ArrayData e
_) = sh
sh
reshape :: HasCallStack => ShapeR sh -> sh -> ShapeR sh' -> Array sh' e -> Array sh e
reshape :: forall sh sh' e.
HasCallStack =>
ShapeR sh -> sh -> ShapeR sh' -> Array sh' e -> Array sh e
reshape ShapeR sh
shR sh
sh ShapeR sh'
shR' (Array sh'
sh' ArrayData e
adata)
= Builder -> Bool -> Array sh e -> Array sh e
forall a. HasCallStack => Builder -> Bool -> a -> a
boundsCheck Builder
"shape mismatch" (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shR sh
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ShapeR sh' -> sh' -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh'
shR' sh'
sh')
(Array sh e -> Array sh e) -> Array sh e -> Array sh e
forall a b. (a -> b) -> a -> b
$ sh -> ArrayData e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh ArrayData e
adata
(!) :: (ArrayR (Array sh e), Array sh e) -> sh -> e
! :: forall sh e. (ArrayR (Array sh e), Array sh e) -> sh -> e
(!) = (ArrayR (Array sh e) -> Array sh e -> sh -> e)
-> (ArrayR (Array sh e), Array sh e) -> sh -> e
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ArrayR (Array sh e) -> Array sh e -> sh -> e
forall sh e. ArrayR (Array sh e) -> Array sh e -> sh -> e
indexArray
(!!) :: (TypeR e, Array sh e) -> Int -> e
!! :: forall e sh. (TypeR e, Array sh e) -> Int -> e
(!!) = (TypeR e -> Array sh e -> Int -> e)
-> (TypeR e, Array sh e) -> Int -> e
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TypeR e -> Array sh e -> Int -> e
forall e sh. TypeR e -> Array sh e -> Int -> e
linearIndexArray
indexArray :: ArrayR (Array sh e) -> Array sh e -> sh -> e
indexArray :: forall sh e. ArrayR (Array sh e) -> Array sh e -> sh -> e
indexArray (ArrayR ShapeR sh
shR TypeR e
adR) (Array sh
sh ArrayData e
adata) sh
ix = TupR ScalarType e -> ArrayData e -> Int -> e
forall e. TupR ScalarType e -> ArrayData e -> Int -> e
indexArrayData TupR ScalarType e
TypeR e
adR ArrayData e
adata (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shR sh
sh
sh sh
sh
ix)
linearIndexArray :: TypeR e -> Array sh e -> Int -> e
linearIndexArray :: forall e sh. TypeR e -> Array sh e -> Int -> e
linearIndexArray TypeR e
adR (Array sh
_ ArrayData e
adata) = TypeR e -> ArrayData e -> Int -> e
forall e. TupR ScalarType e -> ArrayData e -> Int -> e
indexArrayData TypeR e
adR ArrayData e
adata
showArray :: (e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
showArray :: forall e sh.
(e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
showArray e -> ShowS
f arrR :: ArrayR (Array sh e)
arrR@(ArrayR ShapeR sh
shR TypeR e
_) arr :: Array sh e
arr@(Array sh
sh ArrayData e
_) = case ShapeR sh
shR of
ShapeR sh
ShapeRz -> String
"Scalar Z " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
xs
ShapeRsnoc ShapeR sh1
ShapeRz -> String
"Vector (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
shapeString String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
xs
ShapeRsnoc (ShapeRsnoc ShapeR sh1
ShapeRz) -> String
"Matrix (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
shapeString String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (e -> ShowS) -> ArrayR (Array DIM2 e) -> Array DIM2 e -> String
forall e.
(e -> ShowS) -> ArrayR (Array DIM2 e) -> Array DIM2 e -> String
showMatrix e -> ShowS
f ArrayR (Array sh e)
ArrayR (Array DIM2 e)
arrR Array sh e
Array DIM2 e
arr
ShapeR sh
_ -> String
"Array (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
shapeString String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
xs
where
shapeString :: String
shapeString = ShapeR sh -> sh -> String
forall sh. ShapeR sh -> sh -> String
showShape ShapeR sh
shR sh
sh
sh
xs :: String
xs = (e -> ShowS) -> [e] -> ShowS
forall a. (a -> ShowS) -> [a] -> ShowS
showListWith e -> ShowS
f (ArrayR (Array sh e) -> Array sh e -> [e]
forall sh e. ArrayR (Array sh e) -> Array sh e -> [e]
toList ArrayR (Array sh e)
arrR Array sh e
arr) String
""
showArrayShort :: Int -> (e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
showArrayShort :: forall e sh.
Int -> (e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
showArrayShort Int
n e -> ShowS
f ArrayR (Array sh e)
arrR Array sh e
arr = Char
'[' Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> [e] -> String
go Int
0 (ArrayR (Array sh e) -> Array sh e -> [e]
forall sh e. ArrayR (Array sh e) -> Array sh e -> [e]
toList ArrayR (Array sh e)
arrR Array sh e
arr)
where
go :: Int -> [e] -> String
go Int
_ [] = String
"]"
go Int
i (e
x:[e]
xs)
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = String
" ..]"
| Bool
otherwise = Char
',' Char -> ShowS
forall a. a -> [a] -> [a]
: e -> ShowS
f e
x (Int -> [e] -> String
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [e]
xs)
showMatrix :: (e -> ShowS) -> ArrayR (Array DIM2 e) -> Array DIM2 e -> String
showMatrix :: forall e.
(e -> ShowS) -> ArrayR (Array DIM2 e) -> Array DIM2 e -> String
showMatrix e -> ShowS
f (ArrayR ShapeR sh
_ TypeR e
arrR) arr :: Array DIM2 e
arr@(Array DIM2
sh ArrayData e
_)
| Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = String
"[]"
| Bool
otherwise = String
"\n [" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> Int -> String
ppMat Int
0 Int
0
where
(((), Int
rows), Int
cols) = DIM2
sh
lengths :: Vector Int
lengths = Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
U.generate (Int
rowsInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
cols) (\Int
i -> String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (e -> ShowS
f (TypeR e -> Array DIM2 e -> Int -> e
forall e sh. TypeR e -> Array sh e -> Int -> e
linearIndexArray TypeR e
TypeR e
arrR Array DIM2 e
arr Int
i) String
""))
widths :: Vector Int
widths = Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
U.generate Int
cols (\Int
c -> Vector Int -> Int
forall a. (Unbox a, Ord a) => Vector a -> a
U.maximum (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
U.generate Int
rows (\Int
r -> Vector Int
lengths Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.! (Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
colsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
c))))
ppMat :: Int -> Int -> String
ppMat :: Int -> Int -> String
ppMat !Int
r !Int
c | Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
cols = Int -> Int -> String
ppMat (Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
0
ppMat !Int
r !Int
c =
let
!i :: Int
i = Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
colsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
c
!l :: Int
l = Vector Int
lengths Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.! Int
i
!w :: Int
w = Vector Int
widths Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.! Int
c
!pad :: Int
pad = Int
1
cell :: String
cell = Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
wInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
pad) Char
' ' String -> ShowS
forall a. [a] -> [a] -> [a]
++ e -> ShowS
f (TypeR e -> Array DIM2 e -> Int -> e
forall e sh. TypeR e -> Array sh e -> Int -> e
linearIndexArray TypeR e
TypeR e
arrR Array DIM2 e
arr Int
i) String
""
before :: String
before
| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = String
"\n "
| Bool
otherwise = String
""
after :: String
after
| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
rowsInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 Bool -> Bool -> Bool
&& Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
colsInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 = String
"]"
| Bool
otherwise = Char
',' Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> Int -> String
ppMat Int
r (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
in
String
before String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
cell String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
after
showArrays :: ArraysR arrs -> arrs -> String
showArrays :: forall arrs. ArraysR arrs -> arrs -> String
showArrays ArraysR arrs
repr arrs
arrs = ArraysR arrs -> arrs -> ShowS
forall arrs. ArraysR arrs -> arrs -> ShowS
showsArrays ArraysR arrs
repr arrs
arrs String
""
showsArrays :: ArraysR arrs -> arrs -> ShowS
showsArrays :: forall arrs. ArraysR arrs -> arrs -> ShowS
showsArrays ArraysR arrs
repr arrs
arrs = Int -> ArraysR arrs -> arrs -> ShowS
forall a. Int -> ArraysR a -> a -> ShowS
go Int
0 ArraysR arrs
repr arrs
arrs
where
go :: Int -> ArraysR a -> a -> ShowS
go :: forall a. Int -> ArraysR a -> a -> ShowS
go Int
_ TupR ArrayR a
TupRunit ()
= String -> ShowS
showString String
"()"
go Int
level TupR ArrayR a
repr' a
arrs'
| Just [Int -> ShowS]
tuple <- TupR ArrayR a -> a -> Maybe [Int -> ShowS]
forall a. ArraysR a -> a -> Maybe [Int -> ShowS]
extractTuple TupR ArrayR a
repr' a
arrs'
= let
constructor :: String
constructor = Char
'T' Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> String
forall a. Show a => a -> String
show ([Int -> ShowS] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int -> ShowS]
tuple)
level' :: Int
level' = Int
level Int -> Int -> Int
forall a. Num a => a -> a -> a
+ String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
constructor Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
content :: [ShowS]
content = ShowS -> [ShowS] -> [ShowS]
forall a. a -> [a] -> [a]
intersperse ((Char
'\n' Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ShowS
indent Int
level') ([ShowS] -> [ShowS]) -> [ShowS] -> [ShowS]
forall a b. (a -> b) -> a -> b
$ ((Int -> ShowS) -> ShowS) -> [Int -> ShowS] -> [ShowS]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> ShowS) -> Int -> ShowS
forall a b. (a -> b) -> a -> b
$ Int
level') [Int -> ShowS]
tuple
in
String -> ShowS
showString String
constructor ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
' ' Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ShowS -> ShowS -> ShowS) -> ShowS -> [ShowS] -> ShowS
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((ShowS -> ShowS -> ShowS) -> ShowS -> ShowS -> ShowS
forall a b c. (a -> b -> c) -> b -> a -> c
flip ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.)) ShowS
forall a. a -> a
id [ShowS]
content
go Int
level (TupRpair TupR ArrayR a1
repr1 TupR ArrayR b
repr2) (a1
arrs1, b
arrs2)
= String -> ShowS
showString String
"( " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> TupR ArrayR a1 -> a1 -> ShowS
forall a. Int -> ArraysR a -> a -> ShowS
go (Int
level Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) TupR ArrayR a1
repr1 a1
arrs1 ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
"\n"
ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ShowS
indent Int
level ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
", " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> TupR ArrayR b -> b -> ShowS
forall a. Int -> ArraysR a -> a -> ShowS
go (Int
level Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) TupR ArrayR b
repr2 b
arrs2 ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
"\n"
ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ShowS
indent Int
level ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
")"
go Int
level (TupRsingle r :: ArrayR a
r@ArrayR{}) a
arr
= String -> ShowS
showString (String -> ShowS) -> String -> ShowS
forall a b. (a -> b) -> a -> b
$ Int -> ShowS
indents Int
level ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ (e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
forall e sh.
(e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
showArray (TypeR e -> e -> ShowS
forall e. TypeR e -> e -> ShowS
showsElt (TypeR e -> e -> ShowS) -> TypeR e -> e -> ShowS
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> TypeR e
forall sh e. ArrayR (Array sh e) -> TypeR e
arrayRtype ArrayR a
ArrayR (Array sh e)
r) ArrayR a
ArrayR (Array sh e)
r a
Array sh e
arr
indent :: Int -> ShowS
indent :: Int -> ShowS
indent Int
0 String
str = String
str
indent Int
n String
str = Char
' ' Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> ShowS
indent (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) String
str
indents :: Int -> String -> String
indents :: Int -> ShowS
indents Int
_ [] = []
indents Int
level (Char
'\n' : String
xs) = Char
'\n' Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> ShowS
indent Int
level (Int -> ShowS
indents Int
level String
xs)
indents Int
level (Char
x : String
xs) = Char
x Char -> ShowS
forall a. a -> [a] -> [a]
: Int -> ShowS
indents Int
level String
xs
extractTuple :: ArraysR a -> a -> Maybe [Int -> ShowS]
extractTuple :: forall a. ArraysR a -> a -> Maybe [Int -> ShowS]
extractTuple TupR ArrayR a
TupRunit () = [Int -> ShowS] -> Maybe [Int -> ShowS]
forall a. a -> Maybe a
Just []
extractTuple (TupRpair TupR ArrayR a1
rs TupR ArrayR b
r) (a1
as, b
a) = (Int -> ShowS
current (Int -> ShowS) -> [Int -> ShowS] -> [Int -> ShowS]
forall a. a -> [a] -> [a]
:) ([Int -> ShowS] -> [Int -> ShowS])
-> Maybe [Int -> ShowS] -> Maybe [Int -> ShowS]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TupR ArrayR a1 -> a1 -> Maybe [Int -> ShowS]
forall a. ArraysR a -> a -> Maybe [Int -> ShowS]
extractTuple TupR ArrayR a1
rs a1
as
where
current :: Int -> ShowS
current Int
level
| TupR ArrayR b -> b -> Bool
forall a. ArraysR a -> a -> Bool
needsParens TupR ArrayR b
r b
a = String -> ShowS
showString String
"( " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> TupR ArrayR b -> b -> ShowS
forall a. Int -> ArraysR a -> a -> ShowS
go (Int
level Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) TupR ArrayR b
r b
a ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
" )"
| Bool
otherwise = Int -> TupR ArrayR b -> b -> ShowS
forall a. Int -> ArraysR a -> a -> ShowS
go Int
level TupR ArrayR b
r b
a
extractTuple TupR ArrayR a
_ a
_ = Maybe [Int -> ShowS]
forall a. Maybe a
Nothing
needsParens :: ArraysR a -> a -> Bool
needsParens :: forall a. ArraysR a -> a -> Bool
needsParens TupR ArrayR a
TupRunit a
_ = Bool
False
needsParens repr' :: TupR ArrayR a
repr'@(TupRpair TupR ArrayR a1
_ TupR ArrayR b
_) a
as = Maybe [Int -> ShowS] -> Bool
forall a. Maybe a -> Bool
isJust (Maybe [Int -> ShowS] -> Bool) -> Maybe [Int -> ShowS] -> Bool
forall a b. (a -> b) -> a -> b
$ TupR ArrayR a -> a -> Maybe [Int -> ShowS]
forall a. ArraysR a -> a -> Maybe [Int -> ShowS]
extractTuple TupR ArrayR a
repr' a
as
needsParens TupR ArrayR a
_ a
_ = Bool
True
reduceRank :: ArrayR (Array (sh, Int) e) -> ArrayR (Array sh e)
reduceRank :: forall sh e. ArrayR (Array (sh, Int) e) -> ArrayR (Array sh e)
reduceRank (ArrayR (ShapeRsnoc ShapeR sh1
shR) TypeR e
aeR) = ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh1
shR TypeR e
TypeR e
aeR
rnfArray :: ArrayR a -> a -> ()
rnfArray :: forall a. ArrayR a -> a -> ()
rnfArray (ArrayR ShapeR sh
shR TypeR e
adR) (Array sh
sh ArrayData e
ad) = ShapeR sh -> sh -> ()
forall sh. ShapeR sh -> sh -> ()
rnfShape ShapeR sh
shR sh
sh () -> () -> ()
forall a b. a -> b -> b
`seq` TypeR e -> ArrayData e -> ()
forall e. TupR ScalarType e -> ArrayData e -> ()
rnfArrayData TypeR e
adR ArrayData e
ad
rnfArrayR :: ArrayR arr -> ()
rnfArrayR :: forall arr. ArrayR arr -> ()
rnfArrayR (ArrayR ShapeR sh
shR TypeR e
tR) = ShapeR sh -> ()
forall sh. ShapeR sh -> ()
rnfShapeR ShapeR sh
shR () -> () -> ()
forall a b. a -> b -> b
`seq` (forall b. ScalarType b -> ()) -> TypeR e -> ()
forall (s :: * -> *) a. (forall b. s b -> ()) -> TupR s a -> ()
rnfTupR ScalarType b -> ()
forall b. ScalarType b -> ()
rnfScalarType TypeR e
tR
rnfArraysR :: ArraysR arrs -> arrs -> ()
rnfArraysR :: forall arrs. ArraysR arrs -> arrs -> ()
rnfArraysR TupR ArrayR arrs
TupRunit () = ()
rnfArraysR (TupRsingle ArrayR arrs
arrR) arrs
arr = ArrayR arrs -> arrs -> ()
forall a. ArrayR a -> a -> ()
rnfArray ArrayR arrs
arrR arrs
arr
rnfArraysR (TupRpair TupR ArrayR a1
aR1 TupR ArrayR b
aR2) (a1
a1,b
a2) = TupR ArrayR a1 -> a1 -> ()
forall arrs. ArraysR arrs -> arrs -> ()
rnfArraysR TupR ArrayR a1
aR1 a1
a1 () -> () -> ()
forall a b. a -> b -> b
`seq` TupR ArrayR b -> b -> ()
forall arrs. ArraysR arrs -> arrs -> ()
rnfArraysR TupR ArrayR b
aR2 b
a2
liftArrayR :: ArrayR a -> CodeQ (ArrayR a)
liftArrayR :: forall a. ArrayR a -> CodeQ (ArrayR a)
liftArrayR (ArrayR ShapeR sh
shR TypeR e
tR) = [|| ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR $$(ShapeR sh -> CodeQ (ShapeR sh)
forall sh. ShapeR sh -> CodeQ (ShapeR sh)
liftShapeR ShapeR sh
shR) $$(TypeR e -> CodeQ (TypeR e)
forall t. TypeR t -> CodeQ (TypeR t)
liftTypeR TypeR e
tR) ||]
liftArraysR :: ArraysR arrs -> CodeQ (ArraysR arrs)
liftArraysR :: forall arrs. ArraysR arrs -> CodeQ (ArraysR arrs)
liftArraysR TupR ArrayR arrs
TupRunit = [|| TupR s ()
forall (s :: * -> *). TupR s ()
TupRunit ||]
liftArraysR (TupRsingle ArrayR arrs
repr) = [|| s a -> TupR s a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle $$(ArrayR arrs -> CodeQ (ArrayR arrs)
forall a. ArrayR a -> CodeQ (ArrayR a)
liftArrayR ArrayR arrs
repr) ||]
liftArraysR (TupRpair TupR ArrayR a1
a TupR ArrayR b
b) = [|| TupR s a1 -> TupR s b -> TupR s (a1, b)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
TupRpair $$(TupR ArrayR a1 -> CodeQ (TupR ArrayR a1)
forall arrs. ArraysR arrs -> CodeQ (ArraysR arrs)
liftArraysR TupR ArrayR a1
a) $$(TupR ArrayR b -> CodeQ (TupR ArrayR b)
forall arrs. ArraysR arrs -> CodeQ (ArraysR arrs)
liftArraysR TupR ArrayR b
b) ||]
liftArray :: forall sh e. ArrayR (Array sh e) -> Array sh e -> CodeQ (Array sh e)
liftArray :: forall sh e.
ArrayR (Array sh e) -> Array sh e -> CodeQ (Array sh e)
liftArray (ArrayR ShapeR sh
shR TypeR e
adR) (Array sh
sh ArrayData e
adata) =
[|| sh -> ArrayData e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array $$(TypeR sh -> sh -> CodeQ sh
forall t. TypeR t -> t -> CodeQ t
liftElt (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shR) sh
sh
sh) $$(Int -> TypeR e -> ArrayData e -> CodeQ (ArrayData e)
forall e. Int -> TypeR e -> ArrayData e -> CodeQ (ArrayData e)
liftArrayData Int
sz TypeR e
adR ArrayData e
ArrayData e
adata) ||] Code Q (Array sh e) -> Q Type -> Code Q (Array sh e)
forall t. CodeQ t -> Q Type -> CodeQ t
`at` [t| Array $(TypeR sh -> Q Type
forall t. TypeR t -> Q Type
liftTypeQ (ShapeR sh -> TypeR sh
forall sh. ShapeR sh -> TypeR sh
shapeType ShapeR sh
shR)) $(TypeR e -> Q Type
forall t. TypeR t -> Q Type
liftTypeQ TypeR e
adR) |]
where
sz :: Int
sz :: Int
sz = ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shR sh
sh
sh
at :: CodeQ t -> Q Type -> CodeQ t
at :: forall t. CodeQ t -> Q Type -> CodeQ t
at CodeQ t
e Q Type
t = Q Exp -> CodeQ t
forall a (m :: * -> *). Quote m => m Exp -> Code m a
unsafeCodeCoerce (Q Exp -> CodeQ t) -> Q Exp -> CodeQ t
forall a b. (a -> b) -> a -> b
$ Q Exp -> Q Type -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Type -> m Exp
sigE (CodeQ t -> Q Exp
forall a (m :: * -> *). Quote m => Code m a -> m Exp
unTypeCode CodeQ t
e) Q Type
t