module Data.Constructor.Extract.Internal (
    -- * Main data type
    DataAndCon (..),
    conTypeArgNames,
    dataAndConFromName,
    conNamesFromTypeName,
) where

import Data.Generics
import Data.Maybe
import Language.Haskell.TH

data DataAndCon = MkDataAndCon
    { DataAndCon -> Name
dataName :: Name
    , DataAndCon -> [Name]
dataTypeArgNames :: [Name]
    , DataAndCon -> Name
conName :: Name
    , DataAndCon -> [BangType]
conArgs :: [BangType]
    }
    deriving (DataAndCon -> DataAndCon -> Bool
(DataAndCon -> DataAndCon -> Bool)
-> (DataAndCon -> DataAndCon -> Bool) -> Eq DataAndCon
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DataAndCon -> DataAndCon -> Bool
== :: DataAndCon -> DataAndCon -> Bool
$c/= :: DataAndCon -> DataAndCon -> Bool
/= :: DataAndCon -> DataAndCon -> Bool
Eq, Int -> DataAndCon -> ShowS
[DataAndCon] -> ShowS
DataAndCon -> String
(Int -> DataAndCon -> ShowS)
-> (DataAndCon -> String)
-> ([DataAndCon] -> ShowS)
-> Show DataAndCon
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DataAndCon -> ShowS
showsPrec :: Int -> DataAndCon -> ShowS
$cshow :: DataAndCon -> String
show :: DataAndCon -> String
$cshowList :: [DataAndCon] -> ShowS
showList :: [DataAndCon] -> ShowS
Show)

-- | Build 'DataAndCon' using the given 'Dec' (must be a 'DataD') ad the target constructor's name
fromDataD :: (MonadFail m) => Dec -> Name -> m DataAndCon
fromDataD :: forall (m :: * -> *). MonadFail m => Dec -> Name -> m DataAndCon
fromDataD (DataD Cxt
_ Name
dataName [TyVarBndr BndrVis]
tyVarBnd Maybe Type
_ [Con]
cons [DerivClause]
_) Name
expectedConName = do
    let dataTypeArgNames :: [Name]
dataTypeArgNames = TyVarBndr BndrVis -> Name
forall {flag}. TyVarBndr flag -> Name
tyVarBndToTyName (TyVarBndr BndrVis -> Name) -> [TyVarBndr BndrVis] -> [Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TyVarBndr BndrVis]
tyVarBnd
    (conName, conArgs) <-
        m (Name, [BangType])
-> ((Name, [BangType]) -> m (Name, [BangType]))
-> Maybe (Name, [BangType])
-> m (Name, [BangType])
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
            (String -> m (Name, [BangType])
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m (Name, [BangType])) -> String -> m (Name, [BangType])
forall a b. (a -> b) -> a -> b
$ String
"Could not find constructor with name " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
expectedConName)
            (Name, [BangType]) -> m (Name, [BangType])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Maybe (Name, [BangType]) -> m (Name, [BangType]))
-> Maybe (Name, [BangType]) -> m (Name, [BangType])
forall a b. (a -> b) -> a -> b
$ [(Name, [BangType])] -> Maybe (Name, [BangType])
forall {a}. [a] -> Maybe a
safeHead
            ([(Name, [BangType])] -> Maybe (Name, [BangType]))
-> [(Name, [BangType])] -> Maybe (Name, [BangType])
forall a b. (a -> b) -> a -> b
$ (Con -> Maybe (Name, [BangType])) -> [Con] -> [(Name, [BangType])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Con -> Maybe (Name, [BangType])
getConNameAndArgs [Con]
cons
    return MkDataAndCon{..}
  where
    getConNameAndArgs :: Con -> Maybe (Name, [BangType])
getConNameAndArgs (NormalC Name
n [BangType]
bt)
        | Name
n Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
expectedConName = (Name, [BangType]) -> Maybe (Name, [BangType])
forall a. a -> Maybe a
Just (Name
n, [BangType]
bt)
    getConNameAndArgs Con
_ = Maybe (Name, [BangType])
forall a. Maybe a
Nothing
    tyVarBndToTyName :: TyVarBndr flag -> Name
tyVarBndToTyName = \case
        PlainTV Name
t flag
_ -> Name
t
        KindedTV Name
t flag
_ Type
_ -> Name
t
    safeHead :: [a] -> Maybe a
safeHead = \case
        [] -> Maybe a
forall a. Maybe a
Nothing
        (a
a : [a]
_) -> a -> Maybe a
forall a. a -> Maybe a
Just a
a
fromDataD Dec
x Name
_ = String -> m DataAndCon
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m DataAndCon) -> String -> m DataAndCon
forall a b. (a -> b) -> a -> b
$ String
"Expected data declaration, got: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Dec -> String
forall a. Show a => a -> String
show Dec
x

-- | Get 'DataAndCon' from the target constructor's name
dataAndConFromName :: Name -> Q DataAndCon
dataAndConFromName :: Name -> Q DataAndCon
dataAndConFromName Name
name = do
    info <- Name -> Q Info
reify Name
name
    case info of
        DataConI Name
_ Type
_ Name
p ->
            Name -> Q Info
reify Name
p Q Info -> (Info -> Q DataAndCon) -> Q DataAndCon
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                TyConI d :: Dec
d@(DataD{}) -> Dec -> Name -> Q DataAndCon
forall (m :: * -> *). MonadFail m => Dec -> Name -> m DataAndCon
fromDataD Dec
d Name
name
                Info
x -> String -> Q DataAndCon
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Expected data definition, got: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Info -> String
forall a. Show a => a -> String
show Info
x)
        Info
x -> String -> Q DataAndCon
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Expected a data constructor, got: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Info -> String
forall a. Show a => a -> String
show Info
x)

conNamesFromTypeName :: Name -> Q [Name]
conNamesFromTypeName :: Name -> Q [Name]
conNamesFromTypeName Name
rawDataName = do
    let strDataName :: String
strDataName = Name -> String
nameBase Name
rawDataName
    dataName <-
        String -> Q (Maybe Name)
lookupTypeName String
strDataName Q (Maybe Name) -> (Maybe Name -> Q Name) -> Q Name
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Just Name
n -> Name -> Q Name
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Name
n
            Maybe Name
Nothing -> String -> Q Name
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Name) -> String -> Q Name
forall a b. (a -> b) -> a -> b
$ String
"Could not resolve type name from " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
strDataName
    reify dataName >>= \case
        TyConI (DataD Cxt
_ Name
_ [TyVarBndr BndrVis]
_ Maybe Type
_ [Con]
cons [DerivClause]
_) -> [Name] -> Q [Name]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Name] -> Q [Name]) -> [Name] -> Q [Name]
forall a b. (a -> b) -> a -> b
$ Con -> [Name]
conNames (Con -> [Name]) -> [Con] -> [Name]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
`concatMap` [Con]
cons
        Info
e -> String -> Q [Name]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q [Name]) -> String -> Q [Name]
forall a b. (a -> b) -> a -> b
$ String
"Expected a data declaration, got: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Info -> String
forall a. Show a => a -> String
show Info
e
  where
    conNames :: Con -> [Name]
conNames = \case
        NormalC Name
n [BangType]
_ -> [Name
n]
        RecC Name
n [VarBangType]
_ -> [Name
n]
        InfixC BangType
_ Name
n BangType
_ -> [Name
n]
        ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
c -> Con -> [Name]
conNames Con
c
        GadtC [Name]
ns [BangType]
_ Type
_ -> [Name]
ns
        RecGadtC [Name]
ns [VarBangType]
_ Type
_ -> [Name]
ns

conTypeArgNames :: DataAndCon -> [Name]
conTypeArgNames :: DataAndCon -> [Name]
conTypeArgNames DataAndCon
dc = ([Name] -> [Name] -> [Name]) -> GenericQ [Name] -> GenericQ [Name]
forall r. (r -> r -> r) -> GenericQ r -> GenericQ r
everything [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
(++) ([] [Name] -> (Type -> [Name]) -> a -> [Name]
forall a b r. (Typeable a, Typeable b) => r -> (b -> r) -> a -> r
`mkQ` Type -> [Name]
getVarName) ([BangType] -> [Name]) -> [BangType] -> [Name]
forall a b. (a -> b) -> a -> b
$ DataAndCon -> [BangType]
conArgs DataAndCon
dc
  where
    getVarName :: Type -> [Name]
getVarName = \case
        VarT Name
n -> [Name
n]
        Type
_ -> []