module TypeMachine.Type (
    Type (..),
    getField,
    hasField,
    typeToDec,
    decToType,
    reifyType,
) where

import Data.Bifunctor (first)
import Data.Functor ((<&>))
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Language.Haskell.TH.Syntax hiding (Type, reifyType)
import qualified Language.Haskell.TH.Syntax as TH (Type)
import TypeMachine.TH.Internal.Utils

-- | Data structure to easily manipulate Template Haskell's 'Language.Haskell.TH.Dec.DataD'
data Type = Type
    { Type -> Name
name :: Name
    -- ^ Name of the data type
    , Type -> Map String BangType
fields :: Map String BangType
    -- ^ Fields of the data type
    , Type -> [(String, Maybe Kind)]
typeParams :: [(String, Maybe Kind)]
    -- ^ Type parameter of the ADT
    }
    deriving (Int -> Type -> ShowS
[Type] -> ShowS
Type -> String
(Int -> Type -> ShowS)
-> (Type -> String) -> ([Type] -> ShowS) -> Show Type
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Type -> ShowS
showsPrec :: Int -> Type -> ShowS
$cshow :: Type -> String
show :: Type -> String
$cshowList :: [Type] -> ShowS
showList :: [Type] -> ShowS
Show, Type -> Type -> Bool
(Type -> Type -> Bool) -> (Type -> Type -> Bool) -> Eq Type
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Type -> Type -> Bool
== :: Type -> Type -> Bool
$c/= :: Type -> Type -> Bool
/= :: Type -> Type -> Bool
Eq)

getField :: String -> Type -> Maybe (Bang, TH.Type)
getField :: String -> Type -> Maybe BangType
getField String
fieldName Type
ty = String -> Map String BangType -> Maybe BangType
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
fieldName (Type -> Map String BangType
fields Type
ty)

hasField :: String -> Type -> Bool
hasField :: String -> Type -> Bool
hasField String
n Type
t = String -> Map String BangType -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member String
n (Map String BangType -> Bool) -> Map String BangType -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Map String BangType
fields Type
t

-- | Turns a 'Type' back to a Template Haskell 'Language.Haskell.TH.Dec'
typeToDec :: Type -> Dec
typeToDec :: Type -> Dec
typeToDec (Type Name
n Map String BangType
fs [(String, Maybe Kind)]
tp) =
    Cxt
-> Name
-> [TyVarBndr BndrVis]
-> Maybe Kind
-> [Con]
-> [DerivClause]
-> Dec
DataD
        []
        Name
n
        ((Name, Maybe Kind) -> TyVarBndr BndrVis
tpToTyVarBndrs ((Name, Maybe Kind) -> TyVarBndr BndrVis)
-> ((String, Maybe Kind) -> (Name, Maybe Kind))
-> (String, Maybe Kind)
-> TyVarBndr BndrVis
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> Name) -> (String, Maybe Kind) -> (Name, Maybe Kind)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first String -> Name
mkName ((String, Maybe Kind) -> TyVarBndr BndrVis)
-> [(String, Maybe Kind)] -> [TyVarBndr BndrVis]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(String, Maybe Kind)]
tp)
        Maybe Kind
forall a. Maybe a
Nothing
        [Name -> [VarBangType] -> Con
RecC (String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
n) ([VarBangType] -> Con) -> [VarBangType] -> Con
forall a b. (a -> b) -> a -> b
$ Map String BangType -> [VarBangType]
fieldsToVbt Map String BangType
fs]
        []
  where
    tpToTyVarBndrs :: (Name, Maybe Kind) -> TyVarBndr BndrVis
tpToTyVarBndrs (Name
varName, Maybe Kind
Nothing) = Name -> BndrVis -> TyVarBndr BndrVis
forall flag. Name -> flag -> TyVarBndr flag
PlainTV Name
varName BndrVis
BndrInvis
    tpToTyVarBndrs (Name
varName, Just Kind
k) = Name -> BndrVis -> Kind -> TyVarBndr BndrVis
forall flag. Name -> flag -> Kind -> TyVarBndr flag
KindedTV Name
varName BndrVis
BndrInvis Kind
k

-- | Transform a Template Haskell 'Dec' into a TypeMachine's type
--
-- For this to succeed, the input type must have exactly one record constructor
decToType :: (Quasi m) => Dec -> m Type
decToType :: forall (m :: * -> *). Quasi m => Dec -> m Type
decToType Dec
dec = do
    (tyName, tyBndrs, cons) <- case Dec
dec of
        (TySynD Name
_ [TyVarBndr BndrVis]
_ (ConT Name
tyconsName)) -> do
            tyconsInfo <- Q Info -> m Info
forall (m :: * -> *) a. Quasi m => Q a -> m a
runQ (Q Info -> m Info) -> Q Info -> m Info
forall a b. (a -> b) -> a -> b
$ Name -> Q Info
reify Name
tyconsName
            case tyconsInfo of
                TyConI Dec
dec' -> Dec -> m (Name, [TyVarBndr BndrVis], [Con])
forall {m :: * -> *}.
MonadFail m =>
Dec -> m (Name, [TyVarBndr BndrVis], [Con])
getTyInfo Dec
dec'
                Info
_ -> String -> m (Name, [TyVarBndr BndrVis], [Con])
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
failMsg
        Dec
_ -> Dec -> m (Name, [TyVarBndr BndrVis], [Con])
forall {m :: * -> *}.
MonadFail m =>
Dec -> m (Name, [TyVarBndr BndrVis], [Con])
getTyInfo Dec
dec
    vbt <- getRecordConstructorVars cons
    let tparams =
            [TyVarBndr BndrVis]
tyBndrs [TyVarBndr BndrVis]
-> (TyVarBndr BndrVis -> (Name, Maybe Kind))
-> [(Name, Maybe Kind)]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
                PlainTV Name
n BndrVis
_ -> (Name
n, Maybe Kind
forall a. Maybe a
Nothing)
                KindedTV Name
n BndrVis
_ Kind
k -> (Name
n, Kind -> Maybe Kind
forall a. a -> Maybe a
Just Kind
k)
    return $ Type tyName (vbtToFields vbt) (first nameBase <$> tparams)
  where
    failMsg :: String
failMsg = String
"Unsupported data type: Expected record type with exactly one constructor"
    getTyInfo :: Dec -> m (Name, [TyVarBndr BndrVis], [Con])
getTyInfo Dec
decl = case Dec
decl of
        (DataD Cxt
_ Name
tyName [TyVarBndr BndrVis]
tybndrs Maybe Kind
_ [Con]
cons [DerivClause]
_) -> (Name, [TyVarBndr BndrVis], [Con])
-> m (Name, [TyVarBndr BndrVis], [Con])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
tyName, [TyVarBndr BndrVis]
tybndrs, [Con]
cons)
        (NewtypeD Cxt
_ Name
tyName [TyVarBndr BndrVis]
tybndrs Maybe Kind
_ Con
con [DerivClause]
_) -> (Name, [TyVarBndr BndrVis], [Con])
-> m (Name, [TyVarBndr BndrVis], [Con])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
tyName, [TyVarBndr BndrVis]
tybndrs, [Con
con])
        Dec
_ -> String -> m (Name, [TyVarBndr BndrVis], [Con])
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
failMsg

-- | Wrapper around the TH's 'reify' function. Fails if the type is not a datatype declaration
reifyType :: Name -> Q Type
reifyType :: Name -> Q Type
reifyType Name
n = do
    res <- Name -> Q Info
reify Name
n
    case res of
        TyConI Dec
ty -> Dec -> Q Type
forall (m :: * -> *). Quasi m => Dec -> m Type
decToType Dec
ty
        Info
_ -> String -> Q Type
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Invalid Name. Expected Datatype declaration"

{-# INLINE fieldsToVbt #-}
fieldsToVbt :: Map String BangType -> [VarBangType]
fieldsToVbt :: Map String BangType -> [VarBangType]
fieldsToVbt = (String -> BangType -> [VarBangType] -> [VarBangType])
-> [VarBangType] -> Map String BangType -> [VarBangType]
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey (\String
key (Bang
b, Kind
t) [VarBangType]
list -> (String -> Name
mkName String
key, Bang
b, Kind
t) VarBangType -> [VarBangType] -> [VarBangType]
forall a. a -> [a] -> [a]
: [VarBangType]
list) []

{-# INLINE vbtToFields #-}
vbtToFields :: [VarBangType] -> Map String BangType
vbtToFields :: [VarBangType] -> Map String BangType
vbtToFields = [(String, BangType)] -> Map String BangType
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(String, BangType)] -> Map String BangType)
-> ([VarBangType] -> [(String, BangType)])
-> [VarBangType]
-> Map String BangType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarBangType -> (String, BangType))
-> [VarBangType] -> [(String, BangType)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Name
n, Bang
b, Kind
t) -> (Name -> String
nameBase Name
n, (Bang
b, Kind
t)))