{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      : Streamly.Internal.Data.Unbox.TH
-- Copyright   : (c) 2023 Composewell Technologies
-- License     : BSD3-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.Unbox.TH
    ( deriveUnbox

    -- th-helpers
    , DataCon(..)
    , DataType(..)
    , reifyDataType
    ) where

--------------------------------------------------------------------------------
-- Imports
--------------------------------------------------------------------------------

import Data.Word (Word16, Word32, Word64, Word8)
import Data.Proxy (Proxy(..))
import Data.List (elemIndex)

import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Streamly.Internal.Data.Unbox

--------------------------------------------------------------------------------
-- th-utilities
--------------------------------------------------------------------------------

-- Note: We don't support template-haskell < 2.14 (GHC < 8.6)

-- The following are copied to remove the dependency on th-utilities.
-- The code has been copied from th-abstraction and th-utilities.

-- Some CPP macros in the following code are not required but are kept
-- anyway. They can be removed if deemed as a maintainance burden.

#if MIN_VERSION_template_haskell(2,17,0)
type TyVarBndr_ flag = TyVarBndr flag
#else
type TyVarBndr_ flag = TyVarBndr
#endif

-- | Case analysis for a 'TyVarBndr'. If the value is a @'PlainTV' n _@, apply
-- the first function to @n@; if it is @'KindedTV' n _ k@, apply the second
-- function to @n@ and @k@.
elimTV :: (Name -> r) -> (Name -> Kind -> r) -> TyVarBndr_ flag -> r
#if MIN_VERSION_template_haskell(2,17,0)
elimTV :: forall r flag.
(Name -> r) -> (Name -> Type -> r) -> TyVarBndr_ flag -> r
elimTV Name -> r
ptv Name -> Type -> r
_ktv (PlainTV Name
n flag
_)    = Name -> r
ptv Name
n
elimTV Name -> r
_ptv Name -> Type -> r
ktv (KindedTV Name
n flag
_ Type
k) = Name -> Type -> r
ktv Name
n Type
k
#else
elimTV ptv _ktv (PlainTV n)    = ptv n
elimTV _ptv ktv (KindedTV n k) = ktv n k
#endif

-- | Extract the type variable name from a 'TyVarBndr', ignoring the
-- kind signature if one exists.
tvName :: TyVarBndr_ flag -> Name
tvName :: forall flag. TyVarBndr_ flag -> Name
tvName = (Name -> Name) -> (Name -> Type -> Name) -> TyVarBndr_ flag -> Name
forall r flag.
(Name -> r) -> (Name -> Type -> r) -> TyVarBndr_ flag -> r
elimTV Name -> Name
forall a. a -> a
id (\Name
n Type
_ -> Name
n)

-- | Get the 'Name' of a 'TyVarBndr'
tyVarBndrName :: TyVarBndr_ flag -> Name
tyVarBndrName :: forall flag. TyVarBndr_ flag -> Name
tyVarBndrName = TyVarBndr_ flag -> Name
forall flag. TyVarBndr_ flag -> Name
tvName

-- | Simplified info about a 'DataD'. Omits deriving, strictness,
-- kind info, and whether it's @data@ or @newtype@.
data DataType = DataType
    { DataType -> Name
dtName :: Name
    , DataType -> [Name]
dtTvs :: [Name]
    , DataType -> Cxt
dtCxt :: Cxt
    , DataType -> [DataCon]
dtCons :: [DataCon]
    } deriving (DataType -> DataType -> Bool
(DataType -> DataType -> Bool)
-> (DataType -> DataType -> Bool) -> Eq DataType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DataType -> DataType -> Bool
$c/= :: DataType -> DataType -> Bool
== :: DataType -> DataType -> Bool
$c== :: DataType -> DataType -> Bool
Eq, Int -> DataType -> ShowS
[DataType] -> ShowS
DataType -> String
(Int -> DataType -> ShowS)
-> (DataType -> String) -> ([DataType] -> ShowS) -> Show DataType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DataType] -> ShowS
$cshowList :: [DataType] -> ShowS
show :: DataType -> String
$cshow :: DataType -> String
showsPrec :: Int -> DataType -> ShowS
$cshowsPrec :: Int -> DataType -> ShowS
Show, Eq DataType
Eq DataType
-> (DataType -> DataType -> Ordering)
-> (DataType -> DataType -> Bool)
-> (DataType -> DataType -> Bool)
-> (DataType -> DataType -> Bool)
-> (DataType -> DataType -> Bool)
-> (DataType -> DataType -> DataType)
-> (DataType -> DataType -> DataType)
-> Ord DataType
DataType -> DataType -> Bool
DataType -> DataType -> Ordering
DataType -> DataType -> DataType
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: DataType -> DataType -> DataType
$cmin :: DataType -> DataType -> DataType
max :: DataType -> DataType -> DataType
$cmax :: DataType -> DataType -> DataType
>= :: DataType -> DataType -> Bool
$c>= :: DataType -> DataType -> Bool
> :: DataType -> DataType -> Bool
$c> :: DataType -> DataType -> Bool
<= :: DataType -> DataType -> Bool
$c<= :: DataType -> DataType -> Bool
< :: DataType -> DataType -> Bool
$c< :: DataType -> DataType -> Bool
compare :: DataType -> DataType -> Ordering
$ccompare :: DataType -> DataType -> Ordering
Ord) --, Data, Typeable, Generic)

-- | Simplified info about a 'Con'. Omits deriving, strictness, and kind
-- info. This is much nicer than consuming 'Con' directly, because it
-- unifies all the constructors into one.
data DataCon = DataCon
    { DataCon -> Name
dcName :: Name
    , DataCon -> [Name]
dcTvs :: [Name]
    , DataCon -> Cxt
dcCxt :: Cxt
    , DataCon -> [(Maybe Name, Type)]
dcFields :: [(Maybe Name, Type)]
    } deriving (DataCon -> DataCon -> Bool
(DataCon -> DataCon -> Bool)
-> (DataCon -> DataCon -> Bool) -> Eq DataCon
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DataCon -> DataCon -> Bool
$c/= :: DataCon -> DataCon -> Bool
== :: DataCon -> DataCon -> Bool
$c== :: DataCon -> DataCon -> Bool
Eq, Int -> DataCon -> ShowS
[DataCon] -> ShowS
DataCon -> String
(Int -> DataCon -> ShowS)
-> (DataCon -> String) -> ([DataCon] -> ShowS) -> Show DataCon
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DataCon] -> ShowS
$cshowList :: [DataCon] -> ShowS
show :: DataCon -> String
$cshow :: DataCon -> String
showsPrec :: Int -> DataCon -> ShowS
$cshowsPrec :: Int -> DataCon -> ShowS
Show, Eq DataCon
Eq DataCon
-> (DataCon -> DataCon -> Ordering)
-> (DataCon -> DataCon -> Bool)
-> (DataCon -> DataCon -> Bool)
-> (DataCon -> DataCon -> Bool)
-> (DataCon -> DataCon -> Bool)
-> (DataCon -> DataCon -> DataCon)
-> (DataCon -> DataCon -> DataCon)
-> Ord DataCon
DataCon -> DataCon -> Bool
DataCon -> DataCon -> Ordering
DataCon -> DataCon -> DataCon
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: DataCon -> DataCon -> DataCon
$cmin :: DataCon -> DataCon -> DataCon
max :: DataCon -> DataCon -> DataCon
$cmax :: DataCon -> DataCon -> DataCon
>= :: DataCon -> DataCon -> Bool
$c>= :: DataCon -> DataCon -> Bool
> :: DataCon -> DataCon -> Bool
$c> :: DataCon -> DataCon -> Bool
<= :: DataCon -> DataCon -> Bool
$c<= :: DataCon -> DataCon -> Bool
< :: DataCon -> DataCon -> Bool
$c< :: DataCon -> DataCon -> Bool
compare :: DataCon -> DataCon -> Ordering
$ccompare :: DataCon -> DataCon -> Ordering
Ord) --, Data, Typeable, Generic)


-- | Convert a 'Con' to a list of 'DataCon'. The result is a list
-- because 'GadtC' and 'RecGadtC' can define multiple constructors.
conToDataCons :: Con -> [DataCon]
conToDataCons :: Con -> [DataCon]
conToDataCons = \case
    NormalC Name
name [BangType]
slots ->
        [Name -> [Name] -> Cxt -> [(Maybe Name, Type)] -> DataCon
DataCon Name
name [] [] ((BangType -> (Maybe Name, Type))
-> [BangType] -> [(Maybe Name, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Bang
_, Type
ty) -> (Maybe Name
forall a. Maybe a
Nothing, Type
ty)) [BangType]
slots)]
    RecC Name
name [VarBangType]
fields ->
        [Name -> [Name] -> Cxt -> [(Maybe Name, Type)] -> DataCon
DataCon Name
name [] [] ((VarBangType -> (Maybe Name, Type))
-> [VarBangType] -> [(Maybe Name, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n, Bang
_, Type
ty) -> (Name -> Maybe Name
forall a. a -> Maybe a
Just Name
n, Type
ty)) [VarBangType]
fields)]
    InfixC (Bang
_, Type
ty1) Name
name (Bang
_, Type
ty2) ->
        [Name -> [Name] -> Cxt -> [(Maybe Name, Type)] -> DataCon
DataCon Name
name [] [] [(Maybe Name
forall a. Maybe a
Nothing, Type
ty1), (Maybe Name
forall a. Maybe a
Nothing, Type
ty2)]]
    ForallC [TyVarBndr Specificity]
tvs Cxt
preds Con
con ->
        (DataCon -> DataCon) -> [DataCon] -> [DataCon]
forall a b. (a -> b) -> [a] -> [b]
map (\(DataCon Name
name [Name]
tvs0 Cxt
preds0 [(Maybe Name, Type)]
fields) ->
            Name -> [Name] -> Cxt -> [(Maybe Name, Type)] -> DataCon
DataCon Name
name ([Name]
tvs0 [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ (TyVarBndr Specificity -> Name)
-> [TyVarBndr Specificity] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr Specificity -> Name
forall flag. TyVarBndr_ flag -> Name
tyVarBndrName [TyVarBndr Specificity]
tvs) (Cxt
preds0 Cxt -> Cxt -> Cxt
forall a. [a] -> [a] -> [a]
++ Cxt
preds) [(Maybe Name, Type)]
fields) (Con -> [DataCon]
conToDataCons Con
con)
#if MIN_VERSION_template_haskell(2,11,0)
    GadtC [Name]
ns [BangType]
slots Type
_ ->
        (Name -> DataCon) -> [Name] -> [DataCon]
forall a b. (a -> b) -> [a] -> [b]
map (\Name
dn -> Name -> [Name] -> Cxt -> [(Maybe Name, Type)] -> DataCon
DataCon Name
dn [] [] ((BangType -> (Maybe Name, Type))
-> [BangType] -> [(Maybe Name, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Bang
_, Type
ty) -> (Maybe Name
forall a. Maybe a
Nothing, Type
ty)) [BangType]
slots)) [Name]
ns
    RecGadtC [Name]
ns [VarBangType]
fields Type
_ ->
        (Name -> DataCon) -> [Name] -> [DataCon]
forall a b. (a -> b) -> [a] -> [b]
map (\Name
dn -> Name -> [Name] -> Cxt -> [(Maybe Name, Type)] -> DataCon
DataCon Name
dn [] [] ((VarBangType -> (Maybe Name, Type))
-> [VarBangType] -> [(Maybe Name, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
fn, Bang
_, Type
ty) -> (Name -> Maybe Name
forall a. a -> Maybe a
Just Name
fn, Type
ty)) [VarBangType]
fields)) [Name]
ns
#endif

-- | Reify the given data or newtype declaration, and yields its
-- 'DataType' representation.
reifyDataType :: Name -> Q DataType
reifyDataType :: Name -> Q DataType
reifyDataType Name
name = do
    Info
info <- Name -> Q Info
reify Name
name
    case Info -> Maybe DataType
infoToDataType Info
info of
        Maybe DataType
Nothing -> String -> Q DataType
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q DataType) -> String -> Q DataType
forall a b. (a -> b) -> a -> b
$ String
"Expected to reify a datatype. Instead got:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Info -> String
forall a. Ppr a => a -> String
pprint Info
info
        Just DataType
x -> DataType -> Q DataType
forall (m :: * -> *) a. Monad m => a -> m a
return DataType
x

infoToDataType :: Info -> Maybe DataType
infoToDataType :: Info -> Maybe DataType
infoToDataType Info
info = case Info
info of
#if MIN_VERSION_template_haskell(2,11,0)
    TyConI (DataD Cxt
preds Name
name [TyVarBndr ()]
tvs Maybe Type
_kind [Con]
cons [DerivClause]
_deriving) ->
#else
    TyConI (DataD preds name tvs cons _deriving) ->
#endif
        DataType -> Maybe DataType
forall a. a -> Maybe a
Just (DataType -> Maybe DataType) -> DataType -> Maybe DataType
forall a b. (a -> b) -> a -> b
$ Name -> [Name] -> Cxt -> [DataCon] -> DataType
DataType Name
name ((TyVarBndr () -> Name) -> [TyVarBndr ()] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr () -> Name
forall flag. TyVarBndr_ flag -> Name
tyVarBndrName [TyVarBndr ()]
tvs) Cxt
preds ((Con -> [DataCon]) -> [Con] -> [DataCon]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Con -> [DataCon]
conToDataCons [Con]
cons)
#if MIN_VERSION_template_haskell(2,11,0)
    TyConI (NewtypeD Cxt
preds Name
name [TyVarBndr ()]
tvs Maybe Type
_kind Con
con [DerivClause]
_deriving) ->
#else
    TyConI (NewtypeD preds name tvs con _deriving) ->
#endif
        DataType -> Maybe DataType
forall a. a -> Maybe a
Just (DataType -> Maybe DataType) -> DataType -> Maybe DataType
forall a b. (a -> b) -> a -> b
$ Name -> [Name] -> Cxt -> [DataCon] -> DataType
DataType Name
name ((TyVarBndr () -> Name) -> [TyVarBndr ()] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr () -> Name
forall flag. TyVarBndr_ flag -> Name
tyVarBndrName [TyVarBndr ()]
tvs) Cxt
preds (Con -> [DataCon]
conToDataCons Con
con)
    Info
_ -> Maybe DataType
forall a. Maybe a
Nothing

--------------------------------------------------------------------------------
-- Helpers
--------------------------------------------------------------------------------

type Field = (Maybe Name, Type)

_arr :: Name
_arr :: Name
_arr = String -> Name
mkName String
"arr"

_tag :: Name
_tag :: Name
_tag = String -> Name
mkName String
"tag"

_initialOffset :: Name
_initialOffset :: Name
_initialOffset = String -> Name
mkName String
"initialOffset"

_val :: Name
_val :: Name
_val = String -> Name
mkName String
"val"

mkOffsetName :: Int -> Name
mkOffsetName :: Int -> Name
mkOffsetName Int
i = String -> Name
mkName (String
"offset" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i)

mkFieldName :: Int -> Name
mkFieldName :: Int -> Name
mkFieldName Int
i = String -> Name
mkName (String
"field" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i)

--------------------------------------------------------------------------------
-- Domain specific helpers
--------------------------------------------------------------------------------

exprGetSize :: Type -> Q Exp
exprGetSize :: Type -> Q Exp
exprGetSize Type
ty = Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE 'sizeOf) [|Proxy :: Proxy $(pure ty)|]

getTagSize :: Int -> Int
getTagSize :: Int -> Int
getTagSize Int
numConstructors
    | Int
numConstructors Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Int
0
    | Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
forall a. Bounded a => a
maxBound :: Word8) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = Int
1
    | Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16
forall a. Bounded a => a
maxBound :: Word16) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = Int
2
    | Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32
forall a. Bounded a => a
maxBound :: Word32) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = Int
4
    | Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
forall a. Bounded a => a
maxBound :: Word64) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = Int
8
    | Bool
otherwise = String -> Int
forall a. HasCallStack => String -> a
error String
"Too many constructors"

getTagType :: Int -> Name
getTagType :: Int -> Name
getTagType Int
numConstructors
    | Int
numConstructors Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = String -> Name
forall a. HasCallStack => String -> a
error String
"No tag for 1 constructor"
    | Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
forall a. Bounded a => a
maxBound :: Word8) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = ''Word8
    | Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16
forall a. Bounded a => a
maxBound :: Word16) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = ''Word16
    | Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32
forall a. Bounded a => a
maxBound :: Word32) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = ''Word32
    | Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
forall a. Bounded a => a
maxBound :: Word64) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
numConstructors = ''Word64
    | Bool
otherwise = String -> Name
forall a. HasCallStack => String -> a
error String
"Too many constructors"

mkOffsetDecls :: Int -> [Field] -> [Q Dec]
mkOffsetDecls :: Int -> [(Maybe Name, Type)] -> [Q Dec]
mkOffsetDecls Int
tagSize [(Maybe Name, Type)]
fields =
    [Q Dec] -> [Q Dec]
forall a. [a] -> [a]
init
        ((:) (Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD
                  (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP (Int -> Name
mkOffsetName Int
0))
                  (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
                       [|$(litE (IntegerL (fromIntegral tagSize))) +
                         $(varE _initialOffset)|])
                  [])
             (((Int, (Maybe Name, Type)) -> Q Dec)
-> [(Int, (Maybe Name, Type))] -> [Q Dec]
forall a b. (a -> b) -> [a] -> [b]
map (Int, (Maybe Name, Type)) -> Q Dec
forall {a}. (Int, (a, Type)) -> Q Dec
mkOffsetExpr ([Int] -> [(Maybe Name, Type)] -> [(Int, (Maybe Name, Type))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 ..] [(Maybe Name, Type)]
fields)))

    where

    mkOffsetExpr :: (Int, (a, Type)) -> Q Dec
mkOffsetExpr (Int
i, (a
_, Type
ty)) =
        Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD
            (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP (Int -> Name
mkOffsetName Int
i))
            (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|$(varE (mkOffsetName (i - 1))) + $(exprGetSize ty)|])
            []

--------------------------------------------------------------------------------
-- Size
--------------------------------------------------------------------------------

isUnitType :: [DataCon] -> Bool
isUnitType :: [DataCon] -> Bool
isUnitType [DataCon Name
_ [Name]
_ Cxt
_ []] = Bool
True
isUnitType [DataCon]
_ = Bool
False

mkSizeOfExpr :: Type -> [DataCon] -> Q Exp
mkSizeOfExpr :: Type -> [DataCon] -> Q Exp
mkSizeOfExpr Type
headTy [DataCon]
constructors =
    case [DataCon]
constructors of
        [] ->
            [|error
                  ("Attempting to get size with no constructors (" ++
                   $(lift (pprint headTy)) ++ ")")|]
        -- One constructor with no fields is a unit type. Size of a unit type is
        -- 1.
        [con :: DataCon
con@(DataCon Name
_ [Name]
_ Cxt
_ [(Maybe Name, Type)]
fields)] ->
            case [(Maybe Name, Type)]
fields of
                [] -> Lit -> Q Exp
forall (m :: * -> *). Quote m => Lit -> m Exp
litE (Integer -> Lit
IntegerL Integer
1)
                [(Maybe Name, Type)]
_ -> [|$(sizeOfConstructor con)|]
        [DataCon]
_ -> [|$(litE (IntegerL (fromIntegral tagSize))) + $(sizeOfHeadDt)|]

    where

    tagSize :: Int
tagSize = Int -> Int
getTagSize ([DataCon] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DataCon]
constructors)
    sizeOfField :: (a, Type) -> Q Exp
sizeOfField (a
_, Type
ty) = Type -> Q Exp
exprGetSize Type
ty
    sizeOfConstructor :: DataCon -> Q Exp
sizeOfConstructor (DataCon Name
_ [Name]
_ Cxt
_ [(Maybe Name, Type)]
fields) =
        Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE 'sum) ([Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
listE (((Maybe Name, Type) -> Q Exp) -> [(Maybe Name, Type)] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Name, Type) -> Q Exp
forall {a}. (a, Type) -> Q Exp
sizeOfField [(Maybe Name, Type)]
fields))
    -- The size of any Unbox type is atleast 1
    sizeOfHeadDt :: Q Exp
sizeOfHeadDt =
        Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE 'maximum) ([Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
listE ((DataCon -> Q Exp) -> [DataCon] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map DataCon -> Q Exp
sizeOfConstructor [DataCon]
constructors))

--------------------------------------------------------------------------------
-- Peek
--------------------------------------------------------------------------------

mkPeekExprOne :: Int -> DataCon -> Q Exp
mkPeekExprOne :: Int -> DataCon -> Q Exp
mkPeekExprOne Int
tagSize (DataCon Name
cname [Name]
_ Cxt
_ [(Maybe Name, Type)]
fields) =
    case [(Maybe Name, Type)]
fields of
        -- XXX Should we peek and check if the byte value is 0?
        [] -> [|pure $(conE cname)|]
        [(Maybe Name, Type)]
_ ->
            [Q Dec] -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => [m Dec] -> m Exp -> m Exp
letE
                (Int -> [(Maybe Name, Type)] -> [Q Dec]
mkOffsetDecls Int
tagSize [(Maybe Name, Type)]
fields)
                ((Q Exp -> Int -> Q Exp) -> Q Exp -> [Int] -> Q Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                     (\Q Exp
acc Int
i -> [|$(acc) <*> $(peekField i)|])
                     [|$(conE cname) <$> $(peekField 0)|]
                     [Int
1 .. ([(Maybe Name, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Maybe Name, Type)]
fields Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)])

    where

    peekField :: Int -> m Exp
peekField Int
i = [|peekAt $(varE (mkOffsetName i)) $(varE _arr)|]

mkPeekExpr :: Type -> [DataCon] -> Q Exp
mkPeekExpr :: Type -> [DataCon] -> Q Exp
mkPeekExpr Type
headTy [DataCon]
cons =
    case [DataCon]
cons of
        [] ->
            [|error
                  ("Attempting to peek type with no constructors (" ++
                   $(lift (pprint headTy)) ++ ")")|]
        [DataCon
con] -> Int -> DataCon -> Q Exp
mkPeekExprOne Int
0 DataCon
con
        [DataCon]
_ ->
            [Q Stmt] -> Q Exp
forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE
                [ Q Pat -> Q Exp -> Q Stmt
forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS
                      (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
_tag)
                      [|peekAt $(varE _initialOffset) $(varE _arr)|]
                , Q Exp -> Q Stmt
forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS
                      (Q Exp -> [Q Match] -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE
                           (Q Exp -> Q Type -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Type -> m Exp
sigE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
_tag) (Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
tagType))
                           (((Integer, DataCon) -> Q Match)
-> [(Integer, DataCon)] -> [Q Match]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, DataCon) -> Q Match
peekMatch ([Integer] -> [DataCon] -> [(Integer, DataCon)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0 ..] [DataCon]
cons) [Q Match] -> [Q Match] -> [Q Match]
forall a. [a] -> [a] -> [a]
++ [Q Match
peekErr]))
                ]

    where

    lenCons :: Int
lenCons = [DataCon] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DataCon]
cons
    tagType :: Name
tagType = Int -> Name
getTagType Int
lenCons
    tagSize :: Int
tagSize = Int -> Int
getTagSize Int
lenCons
    peekMatch :: (Integer, DataCon) -> Q Match
peekMatch (Integer
i, DataCon
con) =
        Q Pat -> Q Body -> [Q Dec] -> Q Match
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (Lit -> Q Pat
forall (m :: * -> *). Quote m => Lit -> m Pat
litP (Integer -> Lit
IntegerL Integer
i)) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Int -> DataCon -> Q Exp
mkPeekExprOne Int
tagSize DataCon
con)) []
    peekErr :: Q Match
peekErr =
        Q Pat -> Q Body -> [Q Dec] -> Q Match
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match
            Q Pat
forall (m :: * -> *). Quote m => m Pat
wildP
            (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
                 [|error
                       ("Found invalid tag while peeking (" ++
                        $(lift (pprint headTy)) ++ ")")|])
            []

--------------------------------------------------------------------------------
-- Poke
--------------------------------------------------------------------------------

mkPokeExprTag :: Name -> Int -> Q Exp
mkPokeExprTag :: Name -> Int -> Q Exp
mkPokeExprTag Name
tagType Int
tagVal = Q Exp
pokeTag

    where

    pokeTag :: Q Exp
pokeTag =
        [|pokeAt
              $(varE _initialOffset)
              $(varE _arr)
              $((sigE (litE (IntegerL (fromIntegral tagVal))) (conT tagType)))|]

mkPokeExprFields :: Int -> [Field] -> Q Exp
mkPokeExprFields :: Int -> [(Maybe Name, Type)] -> Q Exp
mkPokeExprFields Int
tagSize [(Maybe Name, Type)]
fields = do
    case [(Maybe Name, Type)]
fields of
        [] -> [|pure ()|]
        [(Maybe Name, Type)]
_ ->
            [Q Dec] -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => [m Dec] -> m Exp -> m Exp
letE
                (Int -> [(Maybe Name, Type)] -> [Q Dec]
mkOffsetDecls Int
tagSize [(Maybe Name, Type)]
fields)
                ([Q Stmt] -> Q Exp
forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE ([Q Stmt] -> Q Exp) -> [Q Stmt] -> Q Exp
forall a b. (a -> b) -> a -> b
$ (Int -> Q Stmt) -> [Int] -> [Q Stmt]
forall a b. (a -> b) -> [a] -> [b]
map (Q Exp -> Q Stmt
forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS (Q Exp -> Q Stmt) -> (Int -> Q Exp) -> Int -> Q Stmt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Q Exp
forall {m :: * -> *}. Quote m => Int -> m Exp
pokeField) [Int
0 .. (Int
numFields Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)])

    where

    numFields :: Int
numFields = [(Maybe Name, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Maybe Name, Type)]
fields
    pokeField :: Int -> m Exp
pokeField Int
i =
        [|pokeAt
              $(varE (mkOffsetName i))
              $(varE _arr)
              $(varE (mkFieldName i))|]

mkPokeMatch :: Name -> Int -> Q Exp -> Q Match
mkPokeMatch :: Name -> Int -> Q Exp -> Q Match
mkPokeMatch Name
cname Int
numFields Q Exp
exp0 =
    Q Pat -> Q Body -> [Q Dec] -> Q Match
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match
        (Name -> [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
cname ((Name -> Q Pat) -> [Name] -> [Q Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP ((Int -> Name) -> [Int] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Name
mkFieldName [Int
0 .. (Int
numFields Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)])))
        (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB Q Exp
exp0)
        []

mkPokeExpr :: Type -> [DataCon] -> Q Exp
mkPokeExpr :: Type -> [DataCon] -> Q Exp
mkPokeExpr Type
headTy [DataCon]
cons =
    case [DataCon]
cons of
        [] ->
            [|error
                  ("Attempting to poke type with no constructors (" ++
                   $(lift (pprint headTy)) ++ ")")|]
        -- XXX We don't gaurentee encoded equivalilty for Unbox. Does it still
        -- make sense to encode a default value for unit constructor?
        [(DataCon Name
_ [Name]
_ Cxt
_ [])] -> [|pure ()|] -- mkPokeExprTag ''Word8 0
        [(DataCon Name
cname [Name]
_ Cxt
_ [(Maybe Name, Type)]
fields)] ->
            Q Exp -> [Q Match] -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE
                (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
_val)
                [Name -> Int -> Q Exp -> Q Match
mkPokeMatch Name
cname ([(Maybe Name, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Maybe Name, Type)]
fields) (Int -> [(Maybe Name, Type)] -> Q Exp
mkPokeExprFields Int
0 [(Maybe Name, Type)]
fields)]
        [DataCon]
_ ->
            Q Exp -> [Q Match] -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE
                (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
_val)
                (((Int, DataCon) -> Q Match) -> [(Int, DataCon)] -> [Q Match]
forall a b. (a -> b) -> [a] -> [b]
map (\(Int
tagVal, (DataCon Name
cname [Name]
_ Cxt
_ [(Maybe Name, Type)]
fields)) ->
                          Name -> Int -> Q Exp -> Q Match
mkPokeMatch
                              Name
cname
                              ([(Maybe Name, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Maybe Name, Type)]
fields)
                              ([Q Stmt] -> Q Exp
forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE [ Q Exp -> Q Stmt
forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS (Q Exp -> Q Stmt) -> Q Exp -> Q Stmt
forall a b. (a -> b) -> a -> b
$ Name -> Int -> Q Exp
mkPokeExprTag Name
tagType Int
tagVal
                                   , Q Exp -> Q Stmt
forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS (Q Exp -> Q Stmt) -> Q Exp -> Q Stmt
forall a b. (a -> b) -> a -> b
$ Int -> [(Maybe Name, Type)] -> Q Exp
mkPokeExprFields Int
tagSize [(Maybe Name, Type)]
fields
                                   ]))
                     ([Int] -> [DataCon] -> [(Int, DataCon)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [DataCon]
cons))

    where

    lenCons :: Int
lenCons = [DataCon] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DataCon]
cons
    tagType :: Name
tagType = Int -> Name
getTagType Int
lenCons
    tagSize :: Int
tagSize = Int -> Int
getTagSize Int
lenCons

--------------------------------------------------------------------------------
-- Main
--------------------------------------------------------------------------------

-- | A general function to derive Unbox instances where you can control which
-- Constructors of the datatype to consider and what the Context for the 'Unbox'
-- instance would be.
--
-- Consider the datatype:
-- @
-- data CustomDataType a b
--     = CDTConstructor1
--     | CDTConstructor2 Bool
--     | CDTConstructor3 Bool b
--     deriving (Show, Eq)
-- @
--
-- Usage:
-- @
-- $(deriveUnboxInternal
--       [AppT (ConT ''Unbox) (VarT (mkName "b"))]
--       (AppT
--            (AppT (ConT ''CustomDataType) (VarT (mkName "a")))
--            (VarT (mkName "b")))
--       [ DataCon 'CDTConstructor1 [] [] []
--       , DataCon 'CDTConstructor2 [] [] [(Nothing, (ConT ''Bool))]
--       , DataCon
--             'CDTConstructor3
--             []
--             []
--             [(Nothing, (ConT ''Bool)), (Nothing, (VarT (mkName "b")))]
--       ])
-- @
deriveUnboxInternal :: Type -> [DataCon] -> ([Dec] -> Q [Dec]) -> Q [Dec]
deriveUnboxInternal :: Type -> [DataCon] -> ([Dec] -> Q [Dec]) -> Q [Dec]
deriveUnboxInternal Type
headTy [DataCon]
cons [Dec] -> Q [Dec]
mkDec = do
    Exp
sizeOfMethod <- Type -> [DataCon] -> Q Exp
mkSizeOfExpr Type
headTy [DataCon]
cons
    Exp
peekMethod <- Type -> [DataCon] -> Q Exp
mkPeekExpr Type
headTy [DataCon]
cons
    Exp
pokeMethod <- Type -> [DataCon] -> Q Exp
mkPokeExpr Type
headTy [DataCon]
cons
    let methods :: [Dec]
methods =
            -- INLINE on sizeOf actually worsens some benchmarks, and improves
            -- none
            [ -- PragmaD (InlineP 'sizeOf Inline FunLike AllPhases)
              Name -> [Clause] -> Dec
FunD 'sizeOf [[Pat] -> Body -> [Dec] -> Clause
Clause [Pat
WildP] (Exp -> Body
NormalB Exp
sizeOfMethod) []]
            , Pragma -> Dec
PragmaD (Name -> Inline -> RuleMatch -> Phases -> Pragma
InlineP 'peekAt Inline
Inline RuleMatch
FunLike Phases
AllPhases)
            , Name -> [Clause] -> Dec
FunD
                  'peekAt
                  [ [Pat] -> Body -> [Dec] -> Clause
Clause
                        (if [DataCon] -> Bool
isUnitType [DataCon]
cons
                             then [Pat
WildP, Pat
WildP]
                             else [Name -> Pat
VarP Name
_initialOffset, Name -> Pat
VarP Name
_arr])
                        (Exp -> Body
NormalB Exp
peekMethod)
                        []
                  ]
            , Pragma -> Dec
PragmaD (Name -> Inline -> RuleMatch -> Phases -> Pragma
InlineP 'pokeAt Inline
Inline RuleMatch
FunLike Phases
AllPhases)
            , Name -> [Clause] -> Dec
FunD
                  'pokeAt
                  [ [Pat] -> Body -> [Dec] -> Clause
Clause
                        (if [DataCon] -> Bool
isUnitType [DataCon]
cons
                             then [Pat
WildP, Pat
WildP, Pat
WildP]
                             else [Name -> Pat
VarP Name
_initialOffset, Name -> Pat
VarP Name
_arr, Name -> Pat
VarP Name
_val])
                        (Exp -> Body
NormalB Exp
pokeMethod)
                        []
                  ]
            ]
    [Dec] -> Q [Dec]
mkDec [Dec]
methods

-- | Given an 'Unbox' instance declaration splice without the methods (e.g.
-- @[d|instance Unbox a => Unbox (Maybe a)|]@), generate an instance
-- declaration including all the type class method implementations.
--
-- Usage:
--
-- @
-- \$(deriveUnbox [d|instance Unbox a => Unbox (Maybe a)|])
-- @
deriveUnbox :: Q [Dec] -> Q [Dec]
deriveUnbox :: Q [Dec] -> Q [Dec]
deriveUnbox Q [Dec]
mDecs = do
    [Dec]
dec <- Q [Dec]
mDecs
    case [Dec]
dec of
        [InstanceD Maybe Overlap
mo Cxt
preds Type
headTyWC []] -> do
            let headTy :: Type
headTy = [Dec] -> Type -> Type
forall {a}. Ppr a => a -> Type -> Type
unwrap [Dec]
dec Type
headTyWC
                (Name
mainTyName, Cxt
subs) = [Dec] -> Type -> (Name, Cxt)
forall {p}. Ppr p => p -> Type -> (Name, Cxt)
getMainTypeName [Dec]
dec Type
headTy
            DataType
dt <- Name -> Q DataType
reifyDataType Name
mainTyName
            let tyVars :: [Name]
tyVars = DataType -> [Name]
dtTvs DataType
dt
                mapper :: Type -> Type
mapper = Cxt -> Cxt -> Type -> Type
forall {a}. Eq a => [a] -> [a] -> a -> a
mapperWith (Name -> Type
VarT (Name -> Type) -> [Name] -> Cxt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
tyVars) Cxt
subs
                cons :: [DataCon]
cons = (DataCon -> DataCon) -> [DataCon] -> [DataCon]
forall a b. (a -> b) -> [a] -> [b]
map ((Type -> Type) -> DataCon -> DataCon
modifyConVariables Type -> Type
mapper) (DataType -> [DataCon]
dtCons DataType
dt)
            Type -> [DataCon] -> ([Dec] -> Q [Dec]) -> Q [Dec]
deriveUnboxInternal Type
headTy [DataCon]
cons (Maybe Overlap -> Cxt -> Type -> [Dec] -> Q [Dec]
forall {f :: * -> *}.
Applicative f =>
Maybe Overlap -> Cxt -> Type -> [Dec] -> f [Dec]
mkInst Maybe Overlap
mo Cxt
preds Type
headTyWC)
        [Dec]
_ -> [Dec] -> Q [Dec]
forall {a} {a}. Ppr a => a -> a
errorMessage [Dec]
dec

    where

    mapperWith :: [a] -> [a] -> a -> a
mapperWith [a]
l1 [a]
l2 a
a =
        case a -> [a] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex a
a [a]
l1 of
            Maybe Int
Nothing -> a
a
            -- XXX Capture this case and give a relavant error.
            Just Int
i -> [a]
l2 [a] -> Int -> a
forall a. [a] -> Int -> a
!! Int
i

    mapType :: (Type -> Type) -> Type -> Type
mapType Type -> Type
f (AppT Type
t1 Type
t2) = Type -> Type -> Type
AppT ((Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
t1) ((Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
t2)
    mapType Type -> Type
f (InfixT Type
t1 Name
n Type
t2) = Type -> Name -> Type -> Type
InfixT ((Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
t1) Name
n ((Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
t2)
    mapType Type -> Type
f (UInfixT Type
t1 Name
n Type
t2) = Type -> Name -> Type -> Type
UInfixT ((Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
t1) Name
n ((Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
t2)
    mapType Type -> Type
f (ParensT Type
t) = Type -> Type
ParensT ((Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
t)
    mapType Type -> Type
f Type
v = Type -> Type
f Type
v

    modifyConVariables :: (Type -> Type) -> DataCon -> DataCon
modifyConVariables Type -> Type
f DataCon
con =
        DataCon
con { dcFields :: [(Maybe Name, Type)]
dcFields = ((Maybe Name, Type) -> (Maybe Name, Type))
-> [(Maybe Name, Type)] -> [(Maybe Name, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Maybe Name
a, Type
b) -> (Maybe Name
a, (Type -> Type) -> Type -> Type
mapType Type -> Type
f Type
b)) (DataCon -> [(Maybe Name, Type)]
dcFields DataCon
con) }

    mkInst :: Maybe Overlap -> Cxt -> Type -> [Dec] -> f [Dec]
mkInst Maybe Overlap
mo Cxt
preds Type
headTyWC [Dec]
methods =
        [Dec] -> f [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
mo Cxt
preds Type
headTyWC [Dec]
methods]

    errorMessage :: a -> a
errorMessage a
dec =
        String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
            [ String
"Error: deriveUnbox:"
            , String
""
            , String
">> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Ppr a => a -> String
pprint a
dec
            , String
""
            , String
"The supplied declaration not a valid instance declaration."
            , String
"Provide a valid Haskell instance declaration without a body."
            , String
""
            , String
"Examples:"
            , String
"instance Unbox (Proxy a)"
            , String
"instance Unbox a => Unbox (Identity a)"
            , String
"instance Unbox (TableT Identity)"
            ]

    unwrap :: a -> Type -> Type
unwrap a
_ (AppT (ConT Name
_) Type
r) = Type
r
    unwrap a
dec Type
_ = a -> Type
forall {a} {a}. Ppr a => a -> a
errorMessage a
dec

    getMainTypeName :: p -> Type -> (Name, Cxt)
getMainTypeName p
dec = Cxt -> Type -> (Name, Cxt)
go []

        where

        go :: Cxt -> Type -> (Name, Cxt)
go Cxt
xs (ConT Name
nm) = (Name
nm, Cxt
xs)
        go Cxt
xs (AppT Type
l Type
r) = Cxt -> Type -> (Name, Cxt)
go (Type
rType -> Cxt -> Cxt
forall a. a -> [a] -> [a]
:Cxt
xs) Type
l
        go Cxt
_ Type
_ = p -> (Name, Cxt)
forall {a} {a}. Ppr a => a -> a
errorMessage p
dec