{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Store.TH.Internal
    (
    -- * TH functions for generating Store instances
      deriveManyStoreFromStorable
    , deriveTupleStoreInstance
    , deriveGenericInstance
    , deriveGenericInstanceFromName
    , deriveManyStorePrimVector
    , deriveManyStoreUnboxVector
    , deriveStore
    , makeStore
    -- * Misc utilties used in Store test
    , getAllInstanceTypes1
    , isMonoType
    ) where

import           Control.Applicative
import           Data.Complex ()
import           Data.Generics.Aliases (extT, mkQ, extQ)
import           Data.Generics.Schemes (listify, everywhere, something)
import           Data.List (find)
import qualified Data.Map as M
import           Data.Maybe (fromMaybe, isJust, mapMaybe)
import           Data.Primitive.ByteArray
import           Data.Primitive.Types
import           Data.Store.Core
import           Data.Store.Impl
import qualified Data.Text as T
import           Data.Traversable (forM)
import qualified Data.Vector.Primitive as PV
import qualified Data.Vector.Unboxed as UV
import           Data.Word
import           Foreign.Storable (Storable)
import           GHC.Types (Int(..))
import           Language.Haskell.TH
import           Language.Haskell.TH.ReifyMany.Internal (TypeclassInstance(..), getInstances, unAppsT)
import           Language.Haskell.TH.Syntax (lift)
import           Prelude
import           Safe (headMay)
import           TH.Derive (Deriver(..))
import           TH.ReifySimple
import           TH.Utilities (expectTyCon1, dequalify, plainInstanceD, appsT)

instance Deriver (Store a) where
    runDeriver :: Proxy (Store a) -> [Type] -> Type -> Q [Dec]
runDeriver Proxy (Store a)
_ [Type]
preds Type
ty = do
        Type
argTy <- Name -> Type -> Q Type
expectTyCon1 ''Store Type
ty
        DataType
dt <- Type -> Q DataType
reifyDataTypeSubstituted Type
argTy
        (forall a. a -> [a] -> [a]
:[]) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Type] -> Type -> [DataCon] -> Q Dec
deriveStore [Type]
preds Type
argTy (DataType -> [DataCon]
dtCons DataType
dt)

-- | Given the name of a type, generate a Store instance for it,
-- assuming that all type variables also need to be Store instances.
--
-- Note that when used with datatypes that require type variables, the
-- ScopedTypeVariables extension is required.
makeStore :: Name -> Q [Dec]
makeStore :: Name -> Q [Dec]
makeStore Name
name = do
    DataType
dt <- Name -> Q DataType
reifyDataType Name
name
    let preds :: [Type]
preds = forall a b. (a -> b) -> [a] -> [b]
map (Type -> Type
storePred forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Type
VarT) (DataType -> [Name]
dtTvs DataType
dt)
        argTy :: Type
argTy = Type -> [Type] -> Type
appsT (Name -> Type
ConT Name
name) (forall a b. (a -> b) -> [a] -> [b]
map Name -> Type
VarT (DataType -> [Name]
dtTvs DataType
dt))
    (forall a. a -> [a] -> [a]
:[]) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Type] -> Type -> [DataCon] -> Q Dec
deriveStore [Type]
preds Type
argTy (DataType -> [DataCon]
dtCons DataType
dt)

deriveStore :: Cxt -> Type -> [DataCon] -> Q Dec
deriveStore :: [Type] -> Type -> [DataCon] -> Q Dec
deriveStore [Type]
preds Type
headTy [DataCon]
cons0 =
    [Type] -> Type -> Exp -> Exp -> Exp -> Dec
makeStoreInstance [Type]
preds Type
headTy
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Q Exp
sizeExpr
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Q Exp
peekExpr
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Q Exp
pokeExpr
  where
    cons :: [(Name, [(Name, Type)])]
    cons :: [(Name, [(Name, Type)])]
cons =
      [ ( DataCon -> Name
dcName DataCon
dc
        , [ (String -> Name
mkName (String
"c" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
ixc forall a. [a] -> [a] -> [a]
++ String
"f" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
ixf), Type
ty)
          | Int
ixf <- [Int]
ints
          | (Maybe Name
_, Type
ty) <- DataCon -> [(Maybe Name, Type)]
dcFields DataCon
dc
          ]
        )
      | Int
ixc <- [Int]
ints
      | DataCon
dc <- [DataCon]
cons0
      ]
    -- NOTE: tag code duplicated in th-storable.
    (Name
tagType, Int
_, Int
tagSize) =
        forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error String
"Too many constructors") forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\(Name
_, Int
maxN, Int
_) -> Int
maxN forall a. Ord a => a -> a -> Bool
>= forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Name, [(Name, Type)])]
cons) [(Name, Int, Int)]
tagTypes
    tagTypes :: [(Name, Int, Int)]
    tagTypes :: [(Name, Int, Int)]
tagTypes =
        [ ('(), Int
1, Int
0)
        , (''Word8, forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Bounded a => a
maxBound :: Word8), Int
1)
        , (''Word16, forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Bounded a => a
maxBound :: Word16), Int
2)
        , (''Word32, forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Bounded a => a
maxBound :: Word32), Int
4)
        , (''Word64, forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Bounded a => a
maxBound :: Word64), Int
8)
        ]
    fName :: a -> Name
fName a
ix = String -> Name
mkName (String
"f" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show a
ix)
    ints :: [Int]
ints = [Int
0..] :: [Int]
    fNames :: [Name]
fNames = forall a b. (a -> b) -> [a] -> [b]
map forall {a}. Show a => a -> Name
fName [Int]
ints
    sizeNames :: [Name]
sizeNames = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name, [(Name, Type)])
_ -> String -> Name
mkName forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"sz" forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show) [(Name, [(Name, Type)])]
cons [Int]
ints
    tagName :: Name
tagName = String -> Name
mkName String
"tag"
    valName :: Name
valName = String -> Name
mkName String
"val"
    sizeExpr :: Q Exp
sizeExpr
        -- Maximum size of GHC tuples
        | forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Name, [(Name, Type)])]
cons forall a. Ord a => a -> a -> Bool
<= Int
62 =
            forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (forall (m :: * -> *). Quote m => [m Exp] -> m Exp
tupE (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a b. (a -> b) -> [a] -> [b]
map (Name, Type) -> Q Exp
sizeAtType forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, [(Name, Type)])]
cons))
                  (case [(Name, [(Name, Type)])]
cons of
                     -- Avoid overlapping matches when the case expression is ()
                     [] -> [Q Match
matchConstSize]
                     [(Name, [(Name, Type)])
c] | forall (t :: * -> *) a. Foldable t => t a -> Bool
null (forall a b. (a, b) -> b
snd (Name, [(Name, Type)])
c) -> [Q Match
matchConstSize]
                     [(Name, [(Name, Type)])]
_ -> [Q Match
matchConstSize, Q Match
matchVarSize])
        | Bool
otherwise = Q Exp
varSizeExpr
      where
        sizeAtType :: (Name, Type) -> ExpQ
        sizeAtType :: (Name, Type) -> Q Exp
sizeAtType (Name
_, Type
ty) = [| size :: Size $(return ty) |]
        matchConstSize :: MatchQ
        matchConstSize :: Q Match
matchConstSize = do
            let sz0 :: Exp
sz0 = Name -> Exp
VarE (String -> Name
mkName String
"sz0")
                sizeDecls :: [Q Dec]
sizeDecls =
                    if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Name]
sizeNames
                        then [forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP (String -> Name
mkName String
"sz0")) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [| 0 |]) []]
                        else forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> (Name, [(Name, Type)]) -> Q Dec
constSizeDec [Name]
sizeNames [(Name, [(Name, Type)])]
cons
            Exp
sameSizeExpr <-
                case [Name]
sizeNames of
                    (Name
_ : [Name]
tailSizeNames) ->
                        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Q Exp
l Q Exp
r -> [| $(l) && $(r) |]) [| True |] forall a b. (a -> b) -> a -> b
$
                        forall a b. (a -> b) -> [a] -> [b]
map (\Name
szn -> [| $(return sz0) == $(varE szn) |]) [Name]
tailSizeNames
                    [] -> [| True |]
            Exp
result <- [| ConstSize (tagSize + $(return sz0)) |]
            forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => [m Pat] -> m Pat
tupP (forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n, Type
_) -> forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP 'ConstSize [forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n])
                             (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd [(Name, [(Name, Type)])]
cons)))
                  (forall (m :: * -> *). Quote m => [m (Guard, Exp)] -> m Body
guardedB [forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Guard
NormalG Exp
sameSizeExpr, Exp
result)])
                  [Q Dec]
sizeDecls
        constSizeDec :: Name -> (Name, [(Name, Type)]) -> DecQ
        constSizeDec :: Name -> (Name, [(Name, Type)]) -> Q Dec
constSizeDec Name
szn (Name
_, []) =
            forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
szn) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [| 0 |]) []
        constSizeDec Name
szn (Name
_, [(Name, Type)]
fields) =
            forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
szn) Q Body
body []
          where
            body :: Q Body
body = forall (m :: * -> *). Quote m => m Exp -> m Body
normalB forall a b. (a -> b) -> a -> b
$
                forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 (\Q Exp
l Q Exp
r -> [| $(l) + $(r) |]) forall a b. (a -> b) -> a -> b
$
                forall a b. (a -> b) -> [a] -> [b]
map (\(Name
sizeName, Type
_) -> forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
sizeName) [(Name, Type)]
fields
        matchVarSize :: MatchQ
        matchVarSize :: Q Match
matchVarSize = do
            forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => [m Pat] -> m Pat
tupP (forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n, Type
_) -> forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n) (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd [(Name, [(Name, Type)])]
cons)))
                  (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB Q Exp
varSizeExpr)
                  []
        varSizeExpr :: ExpQ
        varSizeExpr :: Q Exp
varSizeExpr =
            [| VarSize $ \x -> tagSize + $(caseE [| x |] (map matchVar cons)) |]
        matchVar :: (Name, [(Name, Type)]) -> MatchQ
        matchVar :: (Name, [(Name, Type)]) -> Q Match
matchVar (Name
cname, []) =
            forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
cname []) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [| 0 |]) []
        matchVar (Name
cname, [(Name, Type)]
fields) =
            forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
cname (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name, Type)
_ Name
fn -> forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
fn) [(Name, Type)]
fields [Name]
fNames))
                  Q Body
body
                  []
          where
            body :: Q Body
body = forall (m :: * -> *). Quote m => m Exp -> m Body
normalB forall a b. (a -> b) -> a -> b
$
                forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 (\Q Exp
l Q Exp
r -> [| $(l) + $(r) |])
                (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name
sizeName, Type
_) Name
fn -> [| getSizeWith $(varE sizeName) $(varE fn) |])
                         [(Name, Type)]
fields
                         [Name]
fNames)
    -- Choose a tag size large enough for this constructor count.
    -- Expression used for the definition of peek.
    peekExpr :: Q Exp
peekExpr = case [(Name, [(Name, Type)])]
cons of
        [] -> [| error ("Attempting to peek type with no constructors (" ++ $(lift (show headTy)) ++ ")") |]
        [(Name, [(Name, Type)])
con] -> forall {m :: * -> *} {b}. Quote m => (Name, [(Name, b)]) -> m Exp
peekCon (Name, [(Name, Type)])
con
        [(Name, [(Name, Type)])]
_ -> forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE
            [ forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
tagName) [| peek |]
            , forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS (forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (forall (m :: * -> *). Quote m => m Exp -> m Type -> m Exp
sigE (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
tagName) (forall (m :: * -> *). Quote m => Name -> m Type
conT Name
tagType))
                      (forall a b. (a -> b) -> [a] -> [b]
map forall {m :: * -> *} {b}.
Quote m =>
(Integer, (Name, [(Name, b)])) -> m Match
peekMatch (forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0..] [(Name, [(Name, Type)])]
cons) forall a. [a] -> [a] -> [a]
++ [Q Match
peekErr]))
            ]
    peekMatch :: (Integer, (Name, [(Name, b)])) -> m Match
peekMatch (Integer
ix, (Name, [(Name, b)])
con) = forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Lit -> m Pat
litP (Integer -> Lit
IntegerL Integer
ix)) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (forall {m :: * -> *} {b}. Quote m => (Name, [(Name, b)]) -> m Exp
peekCon (Name, [(Name, b)])
con)) []
    peekErr :: Q Match
peekErr = forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match forall (m :: * -> *). Quote m => m Pat
wildP (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
        [| peekException $ T.pack $ "Found invalid tag while peeking (" ++ $(lift (show headTy)) ++ ")" |]) []
    peekCon :: (Name, [(Name, b)]) -> m Exp
peekCon (Name
cname, [(Name, b)]
fields) =
        case [(Name, b)]
fields of
            [] -> [| pure $(conE cname) |]
            [(Name, b)]
_ -> forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE forall a b. (a -> b) -> a -> b
$
                forall a b. (a -> b) -> [a] -> [b]
map (\(Name
fn, b
_) -> forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
fn) [| peek |]) [(Name, b)]
fields forall a. [a] -> [a] -> [a]
++
               [forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
varE 'return) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => [m Exp] -> m Exp
appsE forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
cname forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (\(Name
fn, b
_) -> forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
fn) [(Name, b)]
fields]
    pokeExpr :: Q Exp
pokeExpr = forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE [forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
valName] forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
valName) forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> (Name, [(Name, Type)]) -> Q Match
pokeCon [Int
0..] [(Name, [(Name, Type)])]
cons
    pokeCon :: Int -> (Name, [(Name, Type)]) -> MatchQ
    pokeCon :: Int -> (Name, [(Name, Type)]) -> Q Match
pokeCon Int
ix (Name
cname, [(Name, Type)]
fields) =
        forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
cname (forall a b. (a -> b) -> [a] -> [b]
map (\(Name
fn, Type
_) -> forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
fn) [(Name, Type)]
fields)) Q Body
body []
      where
        body :: Q Body
body = forall (m :: * -> *). Quote m => m Exp -> m Body
normalB forall a b. (a -> b) -> a -> b
$
            case [(Name, [(Name, Type)])]
cons of
                ((Name, [(Name, Type)])
_:(Name, [(Name, Type)])
_:[(Name, [(Name, Type)])]
_) -> forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE (forall {m :: * -> *} {t}. (Quote m, Lift t) => t -> m Stmt
pokeTag Int
ix forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall {m :: * -> *} {b}. Quote m => (Name, b) -> m Stmt
pokeField [(Name, Type)]
fields)
                [(Name, [(Name, Type)])]
_ -> forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE (forall a b. (a -> b) -> [a] -> [b]
map forall {m :: * -> *} {b}. Quote m => (Name, b) -> m Stmt
pokeField [(Name, Type)]
fields)
    pokeTag :: t -> m Stmt
pokeTag t
ix = forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS [| poke (ix :: $(conT tagType)) |]
    pokeField :: (Name, b) -> m Stmt
pokeField (Name
fn, b
_) = forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS [| poke $(varE fn) |]

{- What the generated code looks like

data Foo = Foo Int Double Float

instance Store Foo where
    size =
        case (size :: Size Int, size :: Size Double, size :: Size Float) of
            (ConstSize c0f0, ConstSize c0f1, ConstSize c0f2) -> ConstSize (0 + sz0)
              where
                sz0 = c0f0 + c0f1 + c0f2
            (c0f0, c0f1, c0f2)
                VarSize $ \(Foo f0 f1 f2) -> 0 +
                    getSizeWith c0f0 f0 + getSizeWith c0f1 f1 + getSizeWith c0f2 f2
    peek = do
        f0 <- peek
        f1 <- peek
        f2 <- peek
        return (Foo f0 f1 f2)
    poke (Foo f0 f1 f2) = do
        poke f0
        poke f1
        poke f2

data Bar = Bar Int | Baz Double

instance Store Bar where
    size =
        case (size :: Size Int, size :: Size Double) of
            (ConstSize c0f0, ConstSize c1f0) | sz0 == sz1 -> ConstSize (1 + sz0)
              where
                sz0 = c0f0
                sz1 = c1f0
            (c0f0, c1f0) -> VarSize $ \x -> 1 +
                case x of
                    Bar f0 -> getSizeWith c0f0 f0
                    Baz f0 -> getSizeWith c1f0 f0
    peek = do
        tag <- peek
        case (tag :: Word8) of
            0 -> do
                f0 <- peek
                return (Bar f0)
            1 -> do
                f0 <- peek
                return (Baz f0)
            _ -> peekException "Found invalid tag while peeking (Bar)"
    poke (Bar f0) = do
        poke 0
        poke f0
    poke (Bar f0) = do
        poke 1
        poke f0
-}

------------------------------------------------------------------------
-- Generic

deriveTupleStoreInstance :: Int -> Dec
deriveTupleStoreInstance :: Int -> Dec
deriveTupleStoreInstance Int
n =
    [Type] -> Type -> Dec
deriveGenericInstance (forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
storePred [Type]
tvs)
                          (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Type -> Type -> Type
AppT (Int -> Type
TupleT Int
n forall a. a -> [a] -> [a]
: [Type]
tvs))
  where
    tvs :: [Type]
tvs = forall a. Int -> [a] -> [a]
take Int
n (forall a b. (a -> b) -> [a] -> [b]
map (Name -> Type
VarT forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Name
mkName forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. a -> [a] -> [a]
:[])) [Char
'a'..Char
'z'])

deriveGenericInstance :: Cxt -> Type -> Dec
deriveGenericInstance :: [Type] -> Type -> Dec
deriveGenericInstance [Type]
cs Type
ty = [Type] -> Type -> [Dec] -> Dec
plainInstanceD [Type]
cs (Type -> Type -> Type
AppT (Name -> Type
ConT ''Store) Type
ty) []

deriveGenericInstanceFromName :: Name -> Q Dec
deriveGenericInstanceFromName :: Name -> Q Dec
deriveGenericInstanceFromName Name
n = do
    [Type]
tvs <- forall a b. (a -> b) -> [a] -> [b]
map Name -> Type
VarT forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataType -> [Name]
dtTvs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q DataType
reifyDataType Name
n
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Type] -> Type -> Dec
deriveGenericInstance (forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
storePred [Type]
tvs) (Type -> [Type] -> Type
appsT (Name -> Type
ConT Name
n) [Type]
tvs)

------------------------------------------------------------------------
-- Storable

-- TODO: Generate inline pragmas? Probably not necessary

deriveManyStoreFromStorable :: (Type -> Bool) -> Q [Dec]
deriveManyStoreFromStorable :: (Type -> Bool) -> Q [Dec]
deriveManyStoreFromStorable Type -> Bool
p = do
    Map [Type] TypeclassInstance
storables <- forall a. Map [Type] [a] -> Map [Type] a
postprocess forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeclassInstance] -> Map [Type] [TypeclassInstance]
instancesMap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q [TypeclassInstance]
getInstances ''Storable
    Map [Type] TypeclassInstance
stores <- forall a. Map [Type] [a] -> Map [Type] a
postprocess forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeclassInstance] -> Map [Type] [TypeclassInstance]
instancesMap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q [TypeclassInstance]
getInstances ''Store
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
M.elems forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe (Map [Type] TypeclassInstance
storables forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` Map [Type] TypeclassInstance
stores) forall a b. (a -> b) -> a -> b
$
        \(TypeclassInstance [Type]
cs Type
ty [Dec]
_) ->
        let argTy :: Type
argTy = forall a. [a] -> a
head (forall a. [a] -> [a]
tail (Type -> [Type]
unAppsT Type
ty))
            tyNameLit :: Exp
tyNameLit = Lit -> Exp
LitE (String -> Lit
StringL (forall a. Ppr a => a -> String
pprint Type
ty)) in
        if Type -> Bool
p Type
argTy Bool -> Bool -> Bool
&& Bool -> Bool
not ([Type] -> Bool
superclassHasStorable [Type]
cs)
            then forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ [Type] -> Type -> Exp -> Exp -> Exp -> Dec
makeStoreInstance [Type]
cs Type
argTy
                (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'sizeStorableTy) Exp
tyNameLit)
                (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'peekStorableTy) Exp
tyNameLit)
                (Name -> Exp
VarE 'pokeStorable)
            else forall a. Maybe a
Nothing

-- See #143. Often Storable superclass constraints should instead be
-- Store constraints, so instead it just warns for these.
superclassHasStorable :: Cxt -> Bool
superclassHasStorable :: [Type] -> Bool
superclassHasStorable = forall a. Maybe a -> Bool
isJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. GenericQ (Maybe u) -> GenericQ (Maybe u)
something (forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
mkQ forall a. Maybe a
Nothing Type -> Maybe ()
justStorable forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` String -> Maybe ()
ignoreStrings)
  where
    justStorable :: Type -> Maybe ()
    justStorable :: Type -> Maybe ()
justStorable (ConT Name
n) | Name
n forall a. Eq a => a -> a -> Bool
== ''Storable = forall a. a -> Maybe a
Just ()
    justStorable Type
_ = forall a. Maybe a
Nothing
    ignoreStrings :: String -> Maybe ()
    ignoreStrings :: String -> Maybe ()
ignoreStrings String
_ = forall a. Maybe a
Nothing

------------------------------------------------------------------------
-- Vector

deriveManyStorePrimVector :: Q [Dec]
deriveManyStorePrimVector :: Q [Dec]
deriveManyStorePrimVector = do
    Map [Type] TypeclassInstance
prims <- forall a. Map [Type] [a] -> Map [Type] a
postprocess forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeclassInstance] -> Map [Type] [TypeclassInstance]
instancesMap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q [TypeclassInstance]
getInstances ''PV.Prim
    Map [Type] TypeclassInstance
stores <- forall a. Map [Type] [a] -> Map [Type] a
postprocess forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeclassInstance] -> Map [Type] [TypeclassInstance]
instancesMap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q [TypeclassInstance]
getInstances ''Store
    let primInsts :: Map [Type] TypeclassInstance
primInsts =
            forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys (forall a b. (a -> b) -> [a] -> [b]
map (Type -> Type -> Type
AppT (Name -> Type
ConT ''PV.Vector))) Map [Type] TypeclassInstance
prims
            forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference`
            Map [Type] TypeclassInstance
stores
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall k a. Map k a -> [(k, a)]
M.toList Map [Type] TypeclassInstance
primInsts) forall a b. (a -> b) -> a -> b
$ \([Type], TypeclassInstance)
primInst -> case ([Type], TypeclassInstance)
primInst of
        ([Type
_], TypeclassInstance [Type]
cs Type
ty [Dec]
_) -> do
            let argTy :: Type
argTy = forall a. [a] -> a
head (forall a. [a] -> [a]
tail (Type -> [Type]
unAppsT Type
ty))
            Exp
sizeExpr <- [|
                VarSize $ \x ->
                    I# $(primSizeOfExpr (ConT ''Int)) +
                    I# $(primSizeOfExpr argTy) * PV.length x
                |]
            Exp
peekExpr <- [| do
                len <- peek
                let sz = I# $(primSizeOfExpr argTy)
                array <- peekToByteArray $(lift ("Primitive Vector (" ++ pprint argTy ++ ")"))
                                         (len * sz)
                return (PV.Vector 0 len array)
                |]
            Exp
pokeExpr <- [| \(PV.Vector offset len (ByteArray array)) -> do
                let sz = I# $(primSizeOfExpr argTy)
                poke len
                pokeFromByteArray array (offset * sz) (len * sz)
                |]
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Type] -> Type -> Exp -> Exp -> Exp -> Dec
makeStoreInstance [Type]
cs (Type -> Type -> Type
AppT (Name -> Type
ConT ''PV.Vector) Type
argTy) Exp
sizeExpr Exp
peekExpr Exp
pokeExpr
        ([Type], TypeclassInstance)
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Invariant violated in derivemanyStorePrimVector"


primSizeOfExpr :: Type -> ExpQ
primSizeOfExpr :: Type -> Q Exp
primSizeOfExpr Type
ty = [| $(varE 'sizeOf#) (error "sizeOf# evaluated its argument" :: $(return ty)) |]

deriveManyStoreUnboxVector :: Q [Dec]
deriveManyStoreUnboxVector :: Q [Dec]
deriveManyStoreUnboxVector = do
    [([Type], Type, [DataCon])]
unboxes <- Q [([Type], Type, [DataCon])]
getUnboxInfo
    Map [Type] TypeclassInstance
stores <- forall a. Map [Type] [a] -> Map [Type] a
postprocess forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeclassInstance] -> Map [Type] [TypeclassInstance]
instancesMap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q [TypeclassInstance]
getInstances ''Store
    Map [Type] TypeclassInstance
unboxInstances <- forall a. Map [Type] [a] -> Map [Type] a
postprocess forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeclassInstance] -> Map [Type] [TypeclassInstance]
instancesMap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q [TypeclassInstance]
getInstances ''UV.Unbox
    let dataFamilyDecls :: Map [Type] ([Type], [DataCon])
dataFamilyDecls =
            forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (forall a b. (a -> b) -> [a] -> [b]
map (\([Type]
preds, Type
ty, [DataCon]
cons) -> ([Type -> Type -> Type
AppT (Name -> Type
ConT ''UV.Vector) Type
ty], ([Type]
preds, [DataCon]
cons))) [([Type], Type, [DataCon])]
unboxes)
            forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference`
            Map [Type] TypeclassInstance
stores
#if MIN_VERSION_template_haskell(2,10,0)
        substituteConstraint :: Type -> Type
substituteConstraint (AppT (ConT Name
n) Type
arg)
          | Name
n forall a. Eq a => a -> a -> Bool
== ''UV.Unbox = Type -> Type -> Type
AppT (Name -> Type
ConT ''Store) (Type -> Type -> Type
AppT (Name -> Type
ConT ''UV.Vector) Type
arg)
#else
        substituteConstraint (ClassP n [arg])
          | n == ''UV.Unbox = ClassP ''Store [AppT (ConT ''UV.Vector) arg]
#endif
        substituteConstraint Type
x = Type
x
    -- TODO: ideally this would use a variant of 'deriveStore' which
    -- assumes VarSize.
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall k a. Map k a -> [(k, a)]
M.toList Map [Type] ([Type], [DataCon])
dataFamilyDecls) forall a b. (a -> b) -> a -> b
$ \case
        ([Type
ty], ([Type]
_, [DataCon]
cons)) -> do
            let headTy :: Type
headTy = Type -> Type
getTyHead (Type -> [Type]
unAppsT Type
ty forall a. [a] -> Int -> a
!! Int
1)
            ([Type]
preds, Type
ty') <- case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup [Type
headTy] Map [Type] TypeclassInstance
unboxInstances of
              Maybe TypeclassInstance
Nothing -> do
                String -> Q ()
reportWarning forall a b. (a -> b) -> a -> b
$ String
"No Unbox instance found for " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> String
pprint Type
headTy
                forall (m :: * -> *) a. Monad m => a -> m a
return ([], Type
ty)
              Just (TypeclassInstance [Type]
cs (AppT Type
_ Type
ty') [Dec]
_) ->
                forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
substituteConstraint [Type]
cs, Type -> Type -> Type
AppT (Name -> Type
ConT ''UV.Vector) Type
ty')
              Just TypeclassInstance
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Impossible case"
            [Type] -> Type -> [DataCon] -> Q Dec
deriveStore [Type]
preds Type
ty' [DataCon]
cons
        ([Type], ([Type], [DataCon]))
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"impossible case in deriveManyStoreUnboxVector"

-- TODO: Add something for this purpose to TH.ReifyDataType

getUnboxInfo :: Q [(Cxt, Type, [DataCon])]
getUnboxInfo :: Q [([Type], Type, [DataCon])]
getUnboxInfo = do
    FamilyI Dec
_ [Dec]
insts <- Name -> Q Info
reify ''UV.Vector
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. (a -> b) -> [a] -> [b]
map ((forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhere (forall a. a -> a
id forall a b.
(Typeable a, Typeable b) =>
(a -> a) -> (b -> b) -> a -> a
`extT` Type -> Type
dequalVarT)) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Dec -> Maybe ([Type], Type, [DataCon])
go [Dec]
insts)
  where
#if MIN_VERSION_template_haskell(2,15,0)
    go :: Dec -> Maybe ([Type], Type, [DataCon])
go (NewtypeInstD [Type]
preds Maybe [TyVarBndr ()]
_ Type
lhs Maybe Type
_ Con
con [DerivClause]
_)
      | [Type
_, Type
ty] <- Type -> [Type]
unAppsT Type
lhs
      = [Type] -> Type -> [Con] -> Maybe ([Type], Type, [DataCon])
toResult [Type]
preds Type
ty [Con
con]
    go (DataInstD [Type]
preds Maybe [TyVarBndr ()]
_ Type
lhs Maybe Type
_ [Con]
cons [DerivClause]
_)
      | [Type
_, Type
ty] <- Type -> [Type]
unAppsT Type
lhs
      = [Type] -> Type -> [Con] -> Maybe ([Type], Type, [DataCon])
toResult [Type]
preds Type
ty [Con]
cons
#elif MIN_VERSION_template_haskell(2,11,0)
    go (NewtypeInstD preds _ [ty] _ con _) = toResult preds ty [con]
    go (DataInstD preds _ [ty] _ cons _) = toResult preds ty cons
#else
    go (NewtypeInstD preds _ [ty] con _) = toResult preds ty [con]
    go (DataInstD preds _ [ty] cons _) = toResult preds ty cons
#endif
    go Dec
x = forall a. HasCallStack => String -> a
error (String
"Unexpected result from reifying Unboxed Vector instances: " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> String
pprint Dec
x)
    toResult :: Cxt -> Type -> [Con] -> Maybe (Cxt, Type, [DataCon])
    toResult :: [Type] -> Type -> [Con] -> Maybe ([Type], Type, [DataCon])
toResult [Type]
_ Type
_ [NormalC Name
conName [BangType]
_]
      | Name -> String
nameBase Name
conName forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
skippedUnboxConstructors = forall a. Maybe a
Nothing
    toResult [Type]
preds Type
ty [Con]
cons
      = forall a. a -> Maybe a
Just ([Type]
preds, Type
ty, forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Con -> [DataCon]
conToDataCons [Con]
cons)
    dequalVarT :: Type -> Type
    dequalVarT :: Type -> Type
dequalVarT (VarT Name
n) = Name -> Type
VarT (Name -> Name
dequalify Name
n)
    dequalVarT Type
ty = Type
ty

-- See issue #174
skippedUnboxConstructors :: [String]
skippedUnboxConstructors :: [String]
skippedUnboxConstructors = [String
"MV_UnboxAs", String
"V_UnboxAs", String
"MV_UnboxViaPrim", String
"V_UnboxViaPrim"]

------------------------------------------------------------------------
-- Utilities

-- Filters out overlapping instances and instances with more than one
-- type arg (should be impossible).
postprocess :: M.Map [Type] [a] -> M.Map [Type] a
postprocess :: forall a. Map [Type] [a] -> Map [Type] a
postprocess =
    forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
M.mapMaybeWithKey forall a b. (a -> b) -> a -> b
$ \[Type]
tys [a]
xs ->
        case ([Type]
tys, [a]
xs) of
            ([Type
_ty], [a
x]) -> forall a. a -> Maybe a
Just a
x
            ([Type], [a])
_ -> forall a. Maybe a
Nothing

makeStoreInstance :: Cxt -> Type -> Exp -> Exp -> Exp -> Dec
makeStoreInstance :: [Type] -> Type -> Exp -> Exp -> Exp -> Dec
makeStoreInstance [Type]
cs Type
ty Exp
sizeExpr Exp
peekExpr Exp
pokeExpr =
    [Type] -> Type -> [Dec] -> Dec
plainInstanceD
        [Type]
cs
        (Type -> Type -> Type
AppT (Name -> Type
ConT ''Store) Type
ty)
        [ Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP 'size) (Exp -> Body
NormalB Exp
sizeExpr) []
        , Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP 'peek) (Exp -> Body
NormalB Exp
peekExpr) []
        , Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP 'poke) (Exp -> Body
NormalB Exp
pokeExpr) []
        ]

-- TODO: either generate random types that satisfy instances with
-- variables in them, or have a check that there's at least a manual
-- check for polymorphic instances.

getAllInstanceTypes :: Name -> Q [[Type]]
getAllInstanceTypes :: Name -> Q [[Type]]
getAllInstanceTypes Name
n =
    forall a b. (a -> b) -> [a] -> [b]
map (\(TypeclassInstance [Type]
_ Type
ty [Dec]
_) -> forall a. Int -> [a] -> [a]
drop Int
1 (Type -> [Type]
unAppsT Type
ty)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    Name -> Q [TypeclassInstance]
getInstances Name
n

getAllInstanceTypes1 :: Name -> Q [Type]
getAllInstanceTypes1 :: Name -> Q [Type]
getAllInstanceTypes1 Name
n =
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error String
"getAllMonoInstances1 expected only one type argument") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> Maybe a
headMay))
         (Name -> Q [[Type]]
getAllInstanceTypes Name
n)

isMonoType :: Type -> Bool
isMonoType :: Type -> Bool
isMonoType = forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r. Typeable r => (r -> Bool) -> GenericQ [r]
listify Type -> Bool
isVarT

isVarT :: Type -> Bool
isVarT :: Type -> Bool
isVarT VarT{} = Bool
True
isVarT Type
_ = Bool
False

-- TOOD: move these to th-reify-many

-- | Get a map from the 'getTyHead' type of instances to
-- 'TypeclassInstance'.
instancesMap :: [TypeclassInstance] -> M.Map [Type] [TypeclassInstance]
instancesMap :: [TypeclassInstance] -> Map [Type] [TypeclassInstance]
instancesMap =
    forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith forall a. [a] -> [a] -> [a]
(++) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall a b. (a -> b) -> [a] -> [b]
map (\TypeclassInstance
ti -> (forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
getTyHead (TypeclassInstance -> [Type]
instanceArgTypes TypeclassInstance
ti), [TypeclassInstance
ti]))

instanceArgTypes :: TypeclassInstance -> [Type]
instanceArgTypes :: TypeclassInstance -> [Type]
instanceArgTypes (TypeclassInstance [Type]
_ Type
ty [Dec]
_) = forall a. Int -> [a] -> [a]
drop Int
1 (Type -> [Type]
unAppsT Type
ty)

getTyHead :: Type -> Type
getTyHead :: Type -> Type
getTyHead (SigT Type
x Type
_) = Type -> Type
getTyHead Type
x
getTyHead (ForallT [TyVarBndr Specificity]
_ [Type]
_ Type
x) = Type -> Type
getTyHead Type
x
getTyHead (AppT Type
l Type
_) = Type -> Type
getTyHead Type
l
getTyHead Type
x = Type
x

storePred :: Type -> Pred
storePred :: Type -> Type
storePred Type
ty =
#if MIN_VERSION_template_haskell(2,10,0)
        Type -> Type -> Type
AppT (Name -> Type
ConT ''Store) Type
ty
#else
        ClassP ''Store [ty]
#endif