-- |
-- 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 TemplateHaskell #-}
{-# 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

#if !MIN_VERSION_ghc(9,3,0)
import Control.Monad (when)
#endif

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

-- | 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 :: CorePlugin
installCoreToDos = CorePlugin
forall a. a -> [CoreToDo] -> CoreM [CoreToDo]
install
#if MIN_VERSION_ghc(8,6,0)
    , pluginRecompile :: [CommandLineOption] -> IO PluginRecompile
pluginRecompile = \[CommandLineOption]
_args -> PluginRecompile -> IO PluginRecompile
forall (f :: * -> *) a. Applicative f => a -> f a
pure PluginRecompile
NoForceRecompile
#endif
    }

install :: a -> [CoreToDo] -> CoreM [CoreToDo]
install :: 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 (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 :: [Annotation]
mg_anns = [Annotation]
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 (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 :: CoreProgram
mg_binds = (CoreBind -> CoreBind) -> CoreProgram -> CoreProgram
forall a b. (a -> b) -> [a] -> [b]
map (Name -> CoreExpr -> CoreBind -> CoreBind
updateNameInGut Name
n CoreExpr
expr) (ModGuts -> CoreProgram
mg_binds ModGuts
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 (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
text CommandLineOption
"Could not resolve TH name" SDoc -> SDoc -> SDoc
<+> CommandLineOption -> SDoc
text (Name -> CommandLineOption
forall a. Show a => a -> CommandLineOption
show Name
thn)
        IO Name -> CoreM Name
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Name -> CoreM Name) -> IO Name -> CoreM Name
forall a b. (a -> b) -> a -> b
$ IO Name
forall a. IO a
exitFailure -- kill the compiler. Is there a nicer way?
    Just Name
n -> Name -> CoreM Name
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 (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 (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
    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 (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)

proofPass :: ModGuts -> CoreM ModGuts
proofPass :: CorePluginPass
proofPass ModGuts
guts = do
#if !MIN_VERSION_ghc(9,3,0)
    DynFlags
dflags <- CoreM DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
    Bool -> CoreM () -> CoreM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (DynFlags -> Int
optLevel DynFlags
dflags Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1) (CoreM () -> CoreM ()) -> CoreM () -> CoreM ()
forall a b. (a -> b) -> a -> b
$ WarnReason -> SDoc -> CoreM ()
warnMsg
#if MIN_VERSION_ghc(8,9,0)
        WarnReason
NoReason
#endif
        (SDoc -> CoreM ()) -> SDoc -> CoreM ()
forall a b. (a -> b) -> a -> b
$ [SDoc] -> SDoc
fsep ([SDoc] -> SDoc) -> [SDoc] -> SDoc
forall a b. (a -> b) -> a -> b
$ (CommandLineOption -> SDoc) -> [CommandLineOption] -> [SDoc]
forall a b. (a -> b) -> [a] -> [b]
map CommandLineOption -> SDoc
text
        ([CommandLineOption] -> [SDoc]) -> [CommandLineOption] -> [SDoc]
forall a b. (a -> b) -> a -> b
$ CommandLineOption -> [CommandLineOption]
words CommandLineOption
"Test.Inspection: Compilation without -O detected. Expect optimizations to fail."
#endif
    (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 -> (ModGuts, [(Name, Obligation)])
extractObligations ModGuts
guts)

partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe a -> Maybe b
f = (a -> ([a], [b]) -> ([a], [b])) -> ([a], [b]) -> [a] -> ([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)