{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

-- | Template Haskell utilities.
module Test.Mutagen.TH.Util
  ( -- * Reification helpers
    reifyName
  , reifyTypeDef

    -- * DType helpers
  , (.==.)
  , isMaybeOf

    -- * DTyVarBndrVis helpers
  , dTyVarBndrName

    -- * DCon helpers
  , dConName
  , dConFields
  , dConFieldsTypes
  , dConFieldsNum

    -- * Pure builders
  , applyTyVars
  , mkConDExp
  , mkApplicativeDExp
  , mkListDExp

    -- * Impure builders
  , createDPat

    -- * Error and logging messages
  , unsupported
  , dump
  , withColor
  , mutagenLog
  , mutagenError
  )
where

import Control.Exception (bracket_)
import Control.Monad (replicateM)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Function (on)
import Language.Haskell.TH
  ( Name
  , Ppr
  , Q
  , newName
  , pprint
  , runIO
  )
import Language.Haskell.TH.Desugar
  ( DCon (..)
  , DConFields (..)
  , DDec (..)
  , DExp (..)
  , DInfo (..)
  , DPat (..)
  , DTyVarBndr (..)
  , DTyVarBndrVis
  , DType (..)
  , DTypeArg (..)
  , applyDType
  , dsReify
  )
import System.Console.ANSI
  ( Color (..)
  , ColorIntensity (..)
  , ConsoleLayer (..)
  , SGR (..)
  , setSGR
  )

{-------------------------------------------------------------------------------
-- * Reification helpers
-------------------------------------------------------------------------------}

-- | Reify a name or die gracefully.
reifyName :: Name -> Q DInfo
reifyName :: Name -> Q DInfo
reifyName Name
name = do
  Name -> Q (Maybe DInfo)
forall (q :: * -> *). DsMonad q => Name -> q (Maybe DInfo)
dsReify Name
name Q (Maybe DInfo) -> (Maybe DInfo -> Q DInfo) -> Q DInfo
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 DInfo
info -> DInfo -> Q DInfo
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return DInfo
info
    Maybe DInfo
Nothing -> String -> [Name] -> Q DInfo
forall a b. Show a => String -> [a] -> Q b
mutagenError String
"could not reify name" [Name
name]

-- | Reify a type definition or die gracefully.
reifyTypeDef :: Name -> Q ([DTyVarBndrVis], [DCon])
reifyTypeDef :: Name -> Q ([DTyVarBndrVis], [DCon])
reifyTypeDef Name
name = do
  Name -> Q DInfo
reifyName Name
name Q DInfo
-> (DInfo -> Q ([DTyVarBndrVis], [DCon]))
-> Q ([DTyVarBndrVis], [DCon])
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    DTyConI (DDataD DataFlavor
_ DCxt
_ Name
_ [DTyVarBndrVis]
tvbs Maybe DType
_ [DCon]
dcons [DDerivClause]
_) Maybe [DDec]
_ -> ([DTyVarBndrVis], [DCon]) -> Q ([DTyVarBndrVis], [DCon])
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([DTyVarBndrVis]
tvbs, [DCon]
dcons)
    DInfo
info -> String -> [DInfo] -> Q ([DTyVarBndrVis], [DCon])
forall a b. Show a => String -> [a] -> Q b
mutagenError (String
"could not reify type definition " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Name -> String
forall a. Ppr a => a -> String
dump Name
name) [DInfo
info]

{-------------------------------------------------------------------------------
-- * DType helpers
-------------------------------------------------------------------------------}

-- | Compare DTypes for equality.
(.==.) :: DType -> DType -> Bool
.==. :: DType -> DType -> Bool
(.==.) = DType -> DType -> Bool
forall a. Eq a => a -> a -> Bool
(==) (DType -> DType -> Bool)
-> (DType -> DType) -> DType -> DType -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` DType -> DType
simplifyDType
  where
    -- \| Simplify a 'DType' removing foralls and signatures
    simplifyDType :: DType -> DType
simplifyDType = \case
      DForallT DForallTelescope
_ DType
t -> DType -> DType
simplifyDType DType
t
      DSigT DType
t DType
_ -> DType -> DType
simplifyDType DType
t
      DAppT DType
l DType
r -> DType -> DType -> DType
DAppT (DType -> DType
simplifyDType DType
l) (DType -> DType
simplifyDType DType
r)
      DType
t -> DType
t

-- | Is this type a Maybe thing?
isMaybeOf :: DType -> Maybe DType
isMaybeOf :: DType -> Maybe DType
isMaybeOf = \case
  (DConT Name
f `DAppT` DType
a) | Name
f Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== ''Maybe -> DType -> Maybe DType
forall a. a -> Maybe a
Just DType
a
  DType
_ -> Maybe DType
forall a. Maybe a
Nothing

{-------------------------------------------------------------------------------
-- * DTyVarBndrVis helpers
-------------------------------------------------------------------------------}

-- | Get the name of a 'DTyVarBndrVis'.
dTyVarBndrName :: DTyVarBndrVis -> Name
dTyVarBndrName :: DTyVarBndrVis -> Name
dTyVarBndrName = \case
  (DPlainTV Name
tv BndrVis
_) -> Name
tv
  (DKindedTV Name
tv BndrVis
_ DType
_) -> Name
tv

{-------------------------------------------------------------------------------
-- * DCon helpers
-------------------------------------------------------------------------------}

-- | Get the name of a t'DCon'.
dConName :: DCon -> Name
dConName :: DCon -> Name
dConName (DCon [DTyVarBndrSpec]
_ DCxt
_ Name
name DConFields
_ DType
_) = Name
name

-- | Get the fields of a t'DCon'.
dConFields :: DCon -> DConFields
dConFields :: DCon -> DConFields
dConFields (DCon [DTyVarBndrSpec]
_ DCxt
_ Name
_ DConFields
conFields DType
_) = DConFields
conFields

-- | Get the types of the fields of a t'DCon'.
dConFieldsTypes :: DConFields -> [DType]
dConFieldsTypes :: DConFields -> DCxt
dConFieldsTypes (DNormalC Bool
_ [DBangType]
bts) = DBangType -> DType
forall a b. (a, b) -> b
snd (DBangType -> DType) -> [DBangType] -> DCxt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DBangType]
bts
dConFieldsTypes (DRecC [DVarBangType]
bts) = (\(Name
_, Bang
_, DType
t) -> DType
t) (DVarBangType -> DType) -> [DVarBangType] -> DCxt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DVarBangType]
bts

-- | Get the number of fields of a t'DCon'.
dConFieldsNum :: DConFields -> Int
dConFieldsNum :: DConFields -> Int
dConFieldsNum (DNormalC Bool
_ [DBangType]
bts) = [DBangType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DBangType]
bts
dConFieldsNum (DRecC [DVarBangType]
bts) = [DVarBangType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DVarBangType]
bts

{-------------------------------------------------------------------------------
-- * Pure builders
-------------------------------------------------------------------------------}

-- | Apply a list of type variables to a head constructor.
applyTyVars :: Name -> [DTyVarBndrVis] -> DType
applyTyVars :: Name -> [DTyVarBndrVis] -> DType
applyTyVars Name
typeName [DTyVarBndrVis]
vars =
  DType -> [DTypeArg] -> DType
applyDType
    (Name -> DType
DConT Name
typeName)
    (DTyVarBndrVis -> DTypeArg
forall {flag}. DTyVarBndr flag -> DTypeArg
dTyVarBndrToDTypeArg (DTyVarBndrVis -> DTypeArg) -> [DTyVarBndrVis] -> [DTypeArg]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DTyVarBndrVis]
vars)
  where
    dTyVarBndrToDTypeArg :: DTyVarBndr flag -> DTypeArg
dTyVarBndrToDTypeArg = \case
      (DPlainTV Name
tv flag
_) -> DType -> DTypeArg
DTANormal (Name -> DType
DVarT Name
tv)
      (DKindedTV Name
tv flag
_ DType
_) -> DType -> DTypeArg
DTANormal (Name -> DType
DVarT Name
tv)

-- | Apply a constructor name to a list of field expressions.
mkConDExp :: Name -> [DExp] -> DExp
mkConDExp :: Name -> [DExp] -> DExp
mkConDExp Name
conName =
  (DExp -> DExp -> DExp) -> DExp -> [DExp] -> DExp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl DExp -> DExp -> DExp
DAppE (Name -> DExp
DConE Name
conName)

-- | Build an applicative expression by chaining '<*>' after 'pure'.
mkApplicativeDExp :: Name -> [DExp] -> DExp
mkApplicativeDExp :: Name -> [DExp] -> DExp
mkApplicativeDExp Name
headName =
  (DExp -> DExp -> DExp) -> DExp -> [DExp] -> DExp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl DExp -> DExp -> DExp
appExp DExp
pureExp
  where
    pureExp :: DExp
pureExp = Name -> DExp
DVarE 'pure DExp -> DExp -> DExp
`DAppE` Name -> DExp
DConE Name
headName
    appExp :: DExp -> DExp -> DExp
appExp DExp
l DExp
r = Name -> DExp
DVarE '(<*>) DExp -> DExp -> DExp
`DAppE` DExp
l DExp -> DExp -> DExp
`DAppE` DExp
r

-- | Build a list expression by chaining 'Prelude.:'.
mkListDExp :: [DExp] -> DExp
mkListDExp :: [DExp] -> DExp
mkListDExp =
  (DExp -> DExp -> DExp) -> DExp -> [DExp] -> DExp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr DExp -> DExp -> DExp
consExp DExp
nilExp
  where
    nilExp :: DExp
nilExp = Name -> DExp
DConE '[]
    consExp :: DExp -> DExp -> DExp
consExp DExp
l DExp
r = Name -> DExp
DConE '(:) DExp -> DExp -> DExp
`DAppE` DExp
l DExp -> DExp -> DExp
`DAppE` DExp
r

{-------------------------------------------------------------------------------
-- * Impure builders
-------------------------------------------------------------------------------}

-- | Create a pattern from a constructor using fresh variable names.
--
-- Returns the patterns as well as the freshly bound variables.
createDPat :: DCon -> Q ([Name], DPat)
createDPat :: DCon -> Q ([Name], DPat)
createDPat (DCon [DTyVarBndrSpec]
_ DCxt
_ Name
cname DConFields
cfields DType
_) = do
  vars <- Int -> Q Name -> Q [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (DConFields -> Int
dConFieldsNum DConFields
cfields) (String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"_v")
  let pat = Name -> DCxt -> [DPat] -> DPat
DConP Name
cname [] [Name -> DPat
DVarP Name
var | Name
var <- [Name]
vars]
  return (vars, pat)

{-------------------------------------------------------------------------------
-- * Error and logging messages
-------------------------------------------------------------------------------}

-- | Pretty-print a value for debugging purposes.
--
-- NOTE: this is a simple wrapper around 'pprint' in case we want to change
-- the pretty-printing implementation later.
dump :: (Ppr a) => a -> String
dump :: forall a. Ppr a => a -> String
dump = a -> String
forall a. Ppr a => a -> String
pprint

-- | Run an IO action with a given terminal color, resetting afterwards.
withColor :: Color -> IO () -> IO ()
withColor :: Color -> IO () -> IO ()
withColor Color
color IO ()
io = do
  IO () -> IO () -> IO () -> IO ()
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_
    ([SGR] -> IO ()
setSGR [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
Vivid Color
color])
    ([SGR] -> IO ()
setSGR [SGR
Reset])
    IO ()
io

{- FOURMOLU_DISABLE -}
-- | Log a message with the [MUTAGEN] prefix.
mutagenLog :: (MonadIO m) => String -> m ()
mutagenLog :: forall (m :: * -> *). MonadIO m => String -> m ()
mutagenLog String
_str =
#ifdef MUTAGEN_TH_DEBUG
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"[MUTAGEN] " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
_str
#else
  liftIO $ return ()
#endif
{- FOURMOLU_ENABLE -}

-- | Report a derivation error and die gracefully.
mutagenError :: (Show a) => String -> [a] -> Q b
mutagenError :: forall a b. Show a => String -> [a] -> Q b
mutagenError String
msg [a]
inputs = IO b -> Q b
forall a. IO a -> Q a
runIO (IO b -> Q b) -> IO b -> Q b
forall a b. (a -> b) -> a -> b
$ do
  let report :: String
report =
        [String] -> String
unlines
          ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ [ String
"an error occurred during a MUTAGEN derivation:"
            , String
msg
            , String
"inputs were:"
            ]
            [String] -> [String] -> [String]
forall a. Semigroup a => a -> a -> a
<> (a -> [String]) -> [a] -> [String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (String -> [String]
lines (String -> [String]) -> (a -> String) -> a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show) [a]
inputs
  Color -> IO () -> IO ()
withColor Color
Red (String -> IO ()
forall (m :: * -> *). MonadIO m => String -> m ()
mutagenLog String
report)
  -- Finally, die
  String -> IO b
forall a. HasCallStack => String -> a
error String
report

-- | Blow up on unsupported inputs.
unsupported :: (Show a) => Name -> a -> b
unsupported :: forall a b. Show a => Name -> a -> b
unsupported Name
funName a
input =
  String -> b
forall a. HasCallStack => String -> a
error
    (String -> b) -> String -> b
forall a b. (a -> b) -> a -> b
$ String
"[MUTAGEN]"
      String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Name -> String
forall a. Show a => a -> String
show Name
funName
      String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": unsupported input:\n"
      String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
forall a. Show a => a -> String
show a
input