-- |
-- Module:      Test.Tasty.Inspection.Plugin
-- Copyright:   (c) 2017 Joachim Breitner, 2021 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  andrew.lelechenko@gmail.com
--
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiWayIf            #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE TupleSections         #-}

module Test.Tasty.Inspection.Plugin (plugin) where

import Control.Monad (foldM)
import qualified Language.Haskell.TH.Syntax as TH
import System.Exit (exitFailure)

#if MIN_VERSION_ghc(9,0,0)
import GHC.Plugins
#else
import GhcPlugins
#endif

#if MIN_VERSION_ghc(9,2,0)
import GHC.Types.TyThing
#endif

import Test.Inspection (Obligation(..))
import qualified Test.Inspection.Plugin as P (checkProperty, CheckResult(..))
import Test.Tasty.Inspection.Internal (CheckResult(..))

#if MIN_VERSION_ghc(9,6,0)
import GHC.Driver.Backend (backendForcesOptimization0)
#elif MIN_VERSION_ghc(9,2,0)
import GHC.Driver.Backend (Backend(Interpreter))
#endif

-- | The plugin for inspection testing.
-- You normally do not need to touch it yourself,
-- 'Test.Tasty.Inspection.inspectTest' will enable it automatically.
plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin
    { installCoreToDos = install
#if MIN_VERSION_ghc(8,6,0)
    , pluginRecompile = \[CommandLineOption]
_args -> PluginRecompile -> IO PluginRecompile
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PluginRecompile
NoForceRecompile
#endif
    }

install :: a -> [CoreToDo] -> CoreM [CoreToDo]
install :: forall a. a -> [CoreToDo] -> CoreM [CoreToDo]
install = ([CoreToDo] -> CoreM [CoreToDo])
-> a -> [CoreToDo] -> CoreM [CoreToDo]
forall a b. a -> b -> a
const (([CoreToDo] -> CoreM [CoreToDo])
 -> a -> [CoreToDo] -> CoreM [CoreToDo])
-> ([CoreToDo] -> CoreM [CoreToDo])
-> a
-> [CoreToDo]
-> CoreM [CoreToDo]
forall a b. (a -> b) -> a -> b
$ \[CoreToDo]
passes -> [CoreToDo] -> CoreM [CoreToDo]
forall a. a -> CoreM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([CoreToDo] -> CoreM [CoreToDo]) -> [CoreToDo] -> CoreM [CoreToDo]
forall a b. (a -> b) -> a -> b
$ [CoreToDo]
passes [CoreToDo] -> [CoreToDo] -> [CoreToDo]
forall a. [a] -> [a] -> [a]
++
    [CommandLineOption -> CorePluginPass -> CoreToDo
CoreDoPluginPass CommandLineOption
"Test.Tasty.Inspection.Plugin" CorePluginPass
proofPass]

extractObligations :: ModGuts -> (ModGuts, [(Name, Obligation)])
extractObligations :: ModGuts -> (ModGuts, [(Name, Obligation)])
extractObligations ModGuts
guts = (ModGuts
guts { mg_anns = anns_clean }, [(Name, Obligation)]
obligations)
  where
    ([Annotation]
anns_clean, [(Name, Obligation)]
obligations) = (Annotation -> Maybe (Name, Obligation))
-> [Annotation] -> ([Annotation], [(Name, Obligation)])
forall a b. (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe Annotation -> Maybe (Name, Obligation)
findObligationAnn (ModGuts -> [Annotation]
mg_anns ModGuts
guts)

findObligationAnn :: Annotation -> Maybe (Name, Obligation)
findObligationAnn :: Annotation -> Maybe (Name, Obligation)
findObligationAnn (Annotation (NamedTarget Name
n) AnnPayload
payload) =
    (Name
n,) (Obligation -> (Name, Obligation))
-> Maybe Obligation -> Maybe (Name, Obligation)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Word8] -> Obligation) -> AnnPayload -> Maybe Obligation
forall a. Typeable a => ([Word8] -> a) -> AnnPayload -> Maybe a
fromSerialized [Word8] -> Obligation
forall a. Data a => [Word8] -> a
deserializeWithData AnnPayload
payload
findObligationAnn Annotation
_ = Maybe (Name, Obligation)
forall a. Maybe a
Nothing

checkObligation :: ModGuts -> (Name, Obligation) -> CoreM ModGuts
checkObligation :: ModGuts -> (Name, Obligation) -> CoreM ModGuts
checkObligation ModGuts
guts (Name
name, Obligation
obl) = do
    CheckResult
res <- ModGuts -> Name -> Property -> CoreM CheckResult
P.checkProperty ModGuts
guts (Obligation -> Name
target Obligation
obl) (Obligation -> Property
property Obligation
obl)
    CoreExpr
e   <- CheckResult -> CoreM CoreExpr
resultToExpr CheckResult
res
    CorePluginPass
forall a. a -> CoreM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CorePluginPass -> CorePluginPass
forall a b. (a -> b) -> a -> b
$ Name -> CoreExpr -> ModGuts -> ModGuts
updateNameInGuts Name
name CoreExpr
e ModGuts
guts

updateNameInGuts :: Name -> CoreExpr -> ModGuts -> ModGuts
updateNameInGuts :: Name -> CoreExpr -> ModGuts -> ModGuts
updateNameInGuts Name
n CoreExpr
expr ModGuts
guts =
    ModGuts
guts {mg_binds = map (updateNameInGut n expr) (mg_binds guts) }

updateNameInGut :: Name -> CoreExpr -> CoreBind -> CoreBind
updateNameInGut :: Name -> CoreExpr -> CoreBind -> CoreBind
updateNameInGut Name
n CoreExpr
e = \case
    NonRec CoreBndr
v CoreExpr
_
        | CoreBndr -> Name
forall a. NamedThing a => a -> Name
getName CoreBndr
v Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n -> CoreBndr -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
v CoreExpr
e
    CoreBind
bind -> CoreBind
bind

fromTHName :: TH.Name -> CoreM Name
fromTHName :: Name -> CoreM Name
fromTHName Name
thn = Name -> CoreM (Maybe Name)
thNameToGhcName Name
thn CoreM (Maybe Name) -> (Maybe Name -> CoreM Name) -> CoreM Name
forall a b. CoreM a -> (a -> CoreM b) -> CoreM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe Name
Nothing -> do
        SDoc -> CoreM ()
errorMsg (SDoc -> CoreM ()) -> SDoc -> CoreM ()
forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
forall doc. IsLine doc => CommandLineOption -> doc
text CommandLineOption
"Could not resolve TH name" SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> CommandLineOption -> SDoc
forall doc. IsLine doc => CommandLineOption -> doc
text (Name -> CommandLineOption
forall a. Show a => a -> CommandLineOption
show Name
thn)
        IO Name -> CoreM Name
forall a. IO a -> CoreM a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Name
forall a. IO a
exitFailure -- kill the compiler. Is there a nicer way?
    Just Name
n -> Name -> CoreM Name
forall a. a -> CoreM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Name
n

dcExpr :: TH.Name -> CoreM CoreExpr
dcExpr :: Name -> CoreM CoreExpr
dcExpr Name
thn = do
    Name
name <- Name -> CoreM Name
fromTHName Name
thn
    DataCon
dc <- Name -> CoreM DataCon
forall (m :: * -> *). MonadThings m => Name -> m DataCon
lookupDataCon Name
name
    CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CoreExpr -> CoreM CoreExpr) -> CoreExpr -> CoreM CoreExpr
forall a b. (a -> b) -> a -> b
$ CoreBndr -> CoreExpr
forall b. CoreBndr -> Expr b
Var (DataCon -> CoreBndr
dataConWrapId DataCon
dc)

resultToExpr :: P.CheckResult -> CoreM CoreExpr
resultToExpr :: CheckResult -> CoreM CoreExpr
resultToExpr CheckResult
P.ResSuccess =
    Name -> CoreM CoreExpr
dcExpr 'ResSuccess
resultToExpr (P.ResSuccessWithMessage SDoc
sdoc) = do
    DynFlags
dflags <- CoreM DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
    CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App (CoreExpr -> CoreExpr -> CoreExpr)
-> CoreM CoreExpr -> CoreM (CoreExpr -> CoreExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> CoreM CoreExpr
dcExpr 'ResSuccessWithMessage CoreM (CoreExpr -> CoreExpr) -> CoreM CoreExpr -> CoreM CoreExpr
forall a b. CoreM (a -> b) -> CoreM a -> CoreM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CommandLineOption -> CoreM CoreExpr
forall (m :: * -> *).
MonadThings m =>
CommandLineOption -> m CoreExpr
mkStringExpr (DynFlags -> SDoc -> CommandLineOption
showSDoc DynFlags
dflags SDoc
sdoc)
resultToExpr (P.ResFailure SDoc
sdoc) = do
    DynFlags
dflags <- CoreM DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
    if DynFlags -> Bool
ghciDetected DynFlags
dflags
        then CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App (CoreExpr -> CoreExpr -> CoreExpr)
-> CoreM CoreExpr -> CoreM (CoreExpr -> CoreExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> CoreM CoreExpr
dcExpr 'ResSuccessWithMessage CoreM (CoreExpr -> CoreExpr) -> CoreM CoreExpr -> CoreM CoreExpr
forall a b. CoreM (a -> b) -> CoreM a -> CoreM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CommandLineOption -> CoreM CoreExpr
forall (m :: * -> *).
MonadThings m =>
CommandLineOption -> m CoreExpr
mkStringExpr CommandLineOption
"Skipped because ghci forces -O0"
        else CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App (CoreExpr -> CoreExpr -> CoreExpr)
-> CoreM CoreExpr -> CoreM (CoreExpr -> CoreExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> CoreM CoreExpr
dcExpr 'ResFailure CoreM (CoreExpr -> CoreExpr) -> CoreM CoreExpr -> CoreM CoreExpr
forall a b. CoreM (a -> b) -> CoreM a -> CoreM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CommandLineOption -> CoreM CoreExpr
forall (m :: * -> *).
MonadThings m =>
CommandLineOption -> m CoreExpr
mkStringExpr (DynFlags -> SDoc -> CommandLineOption
showSDoc DynFlags
dflags SDoc
sdoc)

ghciDetected :: DynFlags -> Bool
#if MIN_VERSION_ghc(9,6,0)
ghciDetected :: DynFlags -> Bool
ghciDetected = Backend -> Bool
backendForcesOptimization0 (Backend -> Bool) -> (DynFlags -> Backend) -> DynFlags -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DynFlags -> Backend
backend
#elif MIN_VERSION_ghc(9,2,0)
ghciDetected = (== Interpreter) . backend
#else
ghciDetected = (< 1) . optLevel
#endif

proofPass :: ModGuts -> CoreM ModGuts
proofPass :: CorePluginPass
proofPass = (ModGuts -> [(Name, Obligation)] -> CoreM ModGuts)
-> (ModGuts, [(Name, Obligation)]) -> CoreM ModGuts
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((ModGuts -> (Name, Obligation) -> CoreM ModGuts)
-> ModGuts -> [(Name, Obligation)] -> CoreM ModGuts
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ModGuts -> (Name, Obligation) -> CoreM ModGuts
checkObligation) ((ModGuts, [(Name, Obligation)]) -> CoreM ModGuts)
-> (ModGuts -> (ModGuts, [(Name, Obligation)])) -> CorePluginPass
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModGuts -> (ModGuts, [(Name, Obligation)])
extractObligations

partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe :: forall a b. (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe a -> Maybe b
f = (a -> ([a], [b]) -> ([a], [b])) -> ([a], [b]) -> [a] -> ([a], [b])
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> ([a], [b]) -> ([a], [b])
go ([], [])
  where
    go :: a -> ([a], [b]) -> ([a], [b])
go a
a ([a]
l, [b]
r) = case a -> Maybe b
f a
a of
        Maybe b
Nothing -> (a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
l, [b]
r)
        Just b
b  -> ([a]
l, b
b b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
r)