{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Hs.A ( Apple (..), U
            , AB (..), AI, AF
            , P2 (..), P3 (..), P4 (..)
            , hs2, hs3, hs4
            ) where

import           Control.Monad     (forM, zipWithM_)
import           Data.Int          (Int64)
import           Data.List.Split   (chunksOf)
import           Data.Word         (Word8)
import           Foreign.Ptr       (Ptr, castPtr, plusPtr)
import           Foreign.Storable  (Storable (..))
import           Prettyprinter     (Doc, Pretty (..), align, brackets, concatWith, hardline, space, (<+>))
import           Prettyprinter.Ext

type AI = Apple Int64; type AF = Apple Double
type U a = Ptr (Apple a)

data AB = F | T deriving AB -> AB -> Bool
(AB -> AB -> Bool) -> (AB -> AB -> Bool) -> Eq AB
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AB -> AB -> Bool
== :: AB -> AB -> Bool
$c/= :: AB -> AB -> Bool
/= :: AB -> AB -> Bool
Eq

instance Pretty AB where pretty :: forall ann. AB -> Doc ann
pretty AB
T=Doc ann
"#t"; pretty AB
F=Doc ann
"#f"

instance Show AB where show :: AB -> String
show=Doc Any -> String
forall a. Show a => a -> String
show(Doc Any -> String) -> (AB -> Doc Any) -> AB -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.AB -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. AB -> Doc ann
pretty

data Apple a = AA !Int64 [Int64] [a] deriving (Apple a -> Apple a -> Bool
(Apple a -> Apple a -> Bool)
-> (Apple a -> Apple a -> Bool) -> Eq (Apple a)
forall a. Eq a => Apple a -> Apple a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Apple a -> Apple a -> Bool
== :: Apple a -> Apple a -> Bool
$c/= :: forall a. Eq a => Apple a -> Apple a -> Bool
/= :: Apple a -> Apple a -> Bool
Eq, (forall a b. (a -> b) -> Apple a -> Apple b)
-> (forall a b. a -> Apple b -> Apple a) -> Functor Apple
forall a b. a -> Apple b -> Apple a
forall a b. (a -> b) -> Apple a -> Apple b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Apple a -> Apple b
fmap :: forall a b. (a -> b) -> Apple a -> Apple b
$c<$ :: forall a b. a -> Apple b -> Apple a
<$ :: forall a b. a -> Apple b -> Apple a
Functor)

data P2 a b = P2 a b; hs2 :: P2 a b -> (a, b)
hs2 (P2 a
a b
b) = (a
a,b
b)
data P3 a b c = P3 a b c; hs3 :: P3 a b c -> (a, b, c)
hs3 (P3 a
a b
b c
c) = (a
a,b
b,c
c)
data P4 a b c d = P4 a b c d; hs4 :: P4 a b c d -> (a, b, c, d)
hs4 (P4 a
a b
b c
c d
d) = (a
a,b
b,c
c,d
d)

instance Storable AB where
    sizeOf :: AB -> Int
sizeOf AB
_ = Int
1
    peek :: Ptr AB -> IO AB
peek Ptr AB
p = (\Word8
b -> case Word8
b of Word8
1 -> AB
T; Word8
0 -> AB
F) (Word8 -> AB) -> IO Word8 -> IO AB
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr Word8 -> IO Word8
forall a. Storable a => Ptr a -> IO a
peek (Ptr AB -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr AB
p :: Ptr Word8)
    poke :: Ptr AB -> AB -> IO ()
poke Ptr AB
p AB
F = Ptr Word8 -> Word8 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr AB -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr AB
p :: Ptr Word8) Word8
0
    poke Ptr AB
p AB
T = Ptr Word8 -> Word8 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr AB -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr AB
p :: Ptr Word8) Word8
1

instance (Storable a, Storable b) => Storable (P2 a b) where
    sizeOf :: P2 a b -> Int
sizeOf P2 a b
_ = a -> Int
forall a. Storable a => a -> Int
sizeOf(a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
+b -> Int
forall a. Storable a => a -> Int
sizeOf(b
forall a. HasCallStack => a
undefined::b)
    peek :: Ptr (P2 a b) -> IO (P2 a b)
peek Ptr (P2 a b)
p = a -> b -> P2 a b
forall a b. a -> b -> P2 a b
P2 (a -> b -> P2 a b) -> IO a -> IO (b -> P2 a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P2 a b) -> Ptr a
forall a b. Ptr a -> Ptr b
castPtr Ptr (P2 a b)
p) IO (b -> P2 a b) -> IO b -> IO (P2 a b)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr b -> IO b
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P2 a b)
p Ptr (P2 a b) -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf(a
forall a. HasCallStack => a
undefined::a))

instance (Storable a, Storable b, Storable c) => Storable (P3 a b c) where
    sizeOf :: P3 a b c -> Int
sizeOf P3 a b c
_ = a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
+b -> Int
forall a. Storable a => a -> Int
sizeOf (b
forall a. HasCallStack => a
undefined::b)Int -> Int -> Int
forall a. Num a => a -> a -> a
+c -> Int
forall a. Storable a => a -> Int
sizeOf (c
forall a. HasCallStack => a
undefined::c)
    peek :: Ptr (P3 a b c) -> IO (P3 a b c)
peek Ptr (P3 a b c)
p = a -> b -> c -> P3 a b c
forall a b c. a -> b -> c -> P3 a b c
P3 (a -> b -> c -> P3 a b c) -> IO a -> IO (b -> c -> P3 a b c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P3 a b c) -> Ptr a
forall a b. Ptr a -> Ptr b
castPtr Ptr (P3 a b c)
p) IO (b -> c -> P3 a b c) -> IO b -> IO (c -> P3 a b c)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr b -> IO b
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P3 a b c)
p Ptr (P3 a b c) -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined::a)) IO (c -> P3 a b c) -> IO c -> IO (P3 a b c)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr c -> IO c
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P3 a b c)
p Ptr (P3 a b c) -> Int -> Ptr c
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
+b -> Int
forall a. Storable a => a -> Int
sizeOf (b
forall a. HasCallStack => a
undefined::b)))

instance (Storable a, Storable b, Storable c, Storable d) => Storable (P4 a b c d) where
    sizeOf :: P4 a b c d -> Int
sizeOf P4 a b c d
_ = a -> Int
forall a. Storable a => a -> Int
sizeOf(a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
+b -> Int
forall a. Storable a => a -> Int
sizeOf(b
forall a. HasCallStack => a
undefined::b)Int -> Int -> Int
forall a. Num a => a -> a -> a
+c -> Int
forall a. Storable a => a -> Int
sizeOf(c
forall a. HasCallStack => a
undefined::c)Int -> Int -> Int
forall a. Num a => a -> a -> a
+d -> Int
forall a. Storable a => a -> Int
sizeOf(d
forall a. HasCallStack => a
undefined::d)
    peek :: Ptr (P4 a b c d) -> IO (P4 a b c d)
peek Ptr (P4 a b c d)
p = a -> b -> c -> d -> P4 a b c d
forall a b c d. a -> b -> c -> d -> P4 a b c d
P4 (a -> b -> c -> d -> P4 a b c d)
-> IO a -> IO (b -> c -> d -> P4 a b c d)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P4 a b c d) -> Ptr a
forall a b. Ptr a -> Ptr b
castPtr Ptr (P4 a b c d)
p) IO (b -> c -> d -> P4 a b c d) -> IO b -> IO (c -> d -> P4 a b c d)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr b -> IO b
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P4 a b c d)
p Ptr (P4 a b c d) -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf(a
forall a. HasCallStack => a
undefined::a)) IO (c -> d -> P4 a b c d) -> IO c -> IO (d -> P4 a b c d)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr c -> IO c
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P4 a b c d)
p Ptr (P4 a b c d) -> Int -> Ptr c
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (a -> Int
forall a. Storable a => a -> Int
sizeOf(a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
+b -> Int
forall a. Storable a => a -> Int
sizeOf(b
forall a. HasCallStack => a
undefined::b))) IO (d -> P4 a b c d) -> IO d -> IO (P4 a b c d)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr d -> IO d
forall a. Storable a => Ptr a -> IO a
peek (Ptr (P4 a b c d)
p Ptr (P4 a b c d) -> Int -> Ptr d
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (a -> Int
forall a. Storable a => a -> Int
sizeOf(a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
+b -> Int
forall a. Storable a => a -> Int
sizeOf(b
forall a. HasCallStack => a
undefined::b)Int -> Int -> Int
forall a. Num a => a -> a -> a
+c -> Int
forall a. Storable a => a -> Int
sizeOf(c
forall a. HasCallStack => a
undefined::c)))

pE :: Pretty a => [Int64] -> [a] -> Doc ann
pE :: forall a ann. Pretty a => [Int64] -> [a] -> Doc ann
pE [Int64
_, Int64
n] [a]
xs = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
brackets (Doc ann
forall ann. Doc ann
space Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann -> Doc ann) -> [Doc ann] -> Doc ann
forall (t :: * -> *) ann.
Foldable t =>
(Doc ann -> Doc ann -> Doc ann) -> t (Doc ann) -> Doc ann
concatWith (\Doc ann
x Doc ann
y -> Doc ann
x Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
hardline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
", " Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
y) ([a] -> Doc ann
forall ann. [a] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty([a] -> Doc ann) -> [[a]] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>Int -> [a] -> [[a]]
forall e. Int -> [e] -> [[e]]
chunksOf (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
n) [a]
xs) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
space))
pE [Int64]
_ [a]
xs      = [a] -> Doc ann
forall ann. [a] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [a]
xs

instance Pretty a => Pretty (Apple a) where
    pretty :: forall ann. Apple a -> Doc ann
pretty (AA Int64
_ [Int64]
dims [a]
xs) = Doc ann
"Arr" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann -> [Doc ann] -> Doc ann
forall ann. Doc ann -> [Doc ann] -> Doc ann
tupledBy Doc ann
"×" (Int64 -> Doc ann
forall ann. Int64 -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Int64 -> Doc ann) -> [Int64] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int64]
dims) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Int64] -> [a] -> Doc ann
forall a ann. Pretty a => [Int64] -> [a] -> Doc ann
pE [Int64]
dims [a]
xs

instance Pretty a => Show (Apple a) where show :: Apple a -> String
show=Doc Any -> String
forall a. Show a => a -> String
show(Doc Any -> String) -> (Apple a -> Doc Any) -> Apple a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Apple a -> Doc Any
forall a ann. Pretty a => a -> Doc ann
forall ann. Apple a -> Doc ann
pretty

instance (Pretty a, Pretty b) => Pretty (P2 a b) where
    pretty :: forall ann. P2 a b -> Doc ann
pretty (P2 a
x b
y) = Doc ann -> [Doc ann] -> Doc ann
forall ann. Doc ann -> [Doc ann] -> Doc ann
tupledBy Doc ann
"*" [a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
x, b -> Doc ann
forall ann. b -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty b
y]

instance (Pretty a, Pretty b, Pretty c) => Pretty (P3 a b c) where
    pretty :: forall ann. P3 a b c -> Doc ann
pretty (P3 a
x b
y c
z) = Doc ann -> [Doc ann] -> Doc ann
forall ann. Doc ann -> [Doc ann] -> Doc ann
tupledBy Doc ann
"*" [a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
x, b -> Doc ann
forall ann. b -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty b
y, c -> Doc ann
forall ann. c -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty c
z]

instance (Pretty a, Pretty b, Pretty c, Pretty d) => Pretty (P4 a b c d) where
    pretty :: forall ann. P4 a b c d -> Doc ann
pretty (P4 a
x b
y c
z d
w) = Doc ann -> [Doc ann] -> Doc ann
forall ann. Doc ann -> [Doc ann] -> Doc ann
tupledBy Doc ann
"*" [a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
x, b -> Doc ann
forall ann. b -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty b
y, c -> Doc ann
forall ann. c -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty c
z, d -> Doc ann
forall ann. d -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty d
w]

instance Storable a => Storable (Apple a) where
    sizeOf :: Apple a -> Int
sizeOf (AA Int64
rnk [Int64]
dims [a]
_) = Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
rnkInt -> Int -> Int
forall a. Num a => a -> a -> a
+(a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int64] -> Int64
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int64]
dims))
    poke :: Ptr (Apple a) -> Apple a -> IO ()
poke Ptr (Apple a)
p (AA Int64
rnk [Int64]
dims [a]
xs) = do
        Ptr Int64 -> Int64 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr (Apple a) -> Ptr Int64
forall a b. Ptr a -> Ptr b
castPtr Ptr (Apple a)
p) Int64
rnk
        (Int64 -> Int -> IO ()) -> [Int64] -> [Int] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (\Int64
i Int
o -> Ptr Int64 -> Int64 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr (Apple a)
p Ptr (Apple a) -> Int -> Ptr Int64
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
o)) Int64
i) [Int64]
dims [Int
0..]
        let datOffs :: Int
datOffs = Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
rnk
        (a -> Int -> IO ()) -> [a] -> [Int] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (\a
x Int
o -> Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr (Apple a)
p Ptr (Apple a) -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
datOffsInt -> Int -> Int
forall a. Num a => a -> a -> a
+a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
o)) a
x) [a]
xs [Int
0..]
    peek :: Ptr (Apple a) -> IO (Apple a)
peek Ptr (Apple a)
p = do
        rnk <- Ptr Int64 -> IO Int64
forall a. Storable a => Ptr a -> IO a
peek (Ptr (Apple a) -> Ptr Int64
forall a b. Ptr a -> Ptr b
castPtr Ptr (Apple a)
p)
        dims <- forM [1..fromIntegral rnk] $ \Int
o -> Ptr Int64 -> IO Int64
forall a. Storable a => Ptr a -> IO a
peek (Ptr Int64 -> IO Int64) -> Ptr Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Ptr (Apple a)
p Ptr (Apple a) -> Int -> Ptr Int64
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
o)
        let datOffs = Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
rnk
        xs <- forM [1..fromIntegral (product dims)] $ \Int
o -> Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek (Ptr a -> IO a) -> Ptr a -> IO a
forall a b. (a -> b) -> a -> b
$ Ptr (Apple a)
p Ptr (Apple a) -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
datOffsInt -> Int -> Int
forall a. Num a => a -> a -> a
+a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined::a)Int -> Int -> Int
forall a. Num a => a -> a -> a
*(Int
oInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
        pure $ AA rnk dims xs