{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE PatternSynonyms #-}

-- | GHC source plugin that instruments Haskell code to include tracing calls.
--
-- You can enable this plugin in different ways:
--
-- * Globally, by passing the @-fplugin@ flag to GHC:
--
-- @
--  -fplugin=Test.Mutagen.Tracer.Plugin
-- @
--
-- * Per module, by adding the following pragma to the top of the module:
--
-- @
--  {-# OPTIONS_GHC -fplugin=Test.Mutagen.Tracer.Plugin #-}
-- @
--
-- * Per function, by adding the TRACE annotation pragma to the function (in
--  addition to enabling the plugin globally or per-module):
--
-- @
--  {-# ANN myFunction TRACE #-}
-- @
module Test.Mutagen.Tracer.Plugin
  ( -- * GHC Plugin
    plugin
  )
where

import Control.Monad ((>=>))
import Control.Monad.Writer (MonadIO, WriterT, lift, runWriterT, tell)
import Data.Generics (Data, everywhereM, listify, mkM)
import Data.IORef (IORef, atomicModifyIORef', newIORef)
import GHC.Hs
  ( AnnDecl (..)
  , AnnProvenance (..)
  , GRHS (..)
  , GhcPs
  , HsExpr (..)
  , HsLit (..)
  , HsMatchContext (..)
  , HsModule (..)
  , HsParsedModule (..)
  , ImportDecl (..)
  , ImportDeclQualifiedStyle (..)
  , LHsExpr
  , Match (..)
  , gHsPar
  , getLocA
  , noExtField
  , noLocA
  , simpleImportDecl
  )
import qualified GHC.Hs as GHC
import GHC.Plugins hiding ((<>))
import GHC.Types.SourceText (mkIntegralLit)
import System.IO.Unsafe (unsafePerformIO)
import Test.Mutagen.Tracer.Metadata
  ( ModuleMetadata (..)
  , NodeLocation (..)
  , NodeMetadata (..)
  , NodeType (..)
  , mkSingleModuleTracerMetadata
  , saveTracerMetadata
  )
import Test.Mutagen.Tracer.Trace (TraceNode)

{-------------------------------------------------------------------------------
-- * GHC Plugin
-------------------------------------------------------------------------------}

-- | Top-level plugin.
plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin{parsedResultAction = action}
  where
    action :: [String] -> ModSummary -> ParsedResult -> Hsc ParsedResult
action [String]
cli ModSummary
summary ParsedResult
parsed = do
      let ParsedResult HsParsedModule
source PsMessages
msgs = ParsedResult
parsed
      let L SrcSpan
loc HsModule GhcPs
modAST = HsParsedModule -> GenLocated SrcSpan (HsModule GhcPs)
hpm_module HsParsedModule
source
      flags <- Hsc DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
      let modName = ModuleName -> String
moduleNameString (GenModule Unit -> ModuleName
forall unit. GenModule unit -> ModuleName
moduleName (ModSummary -> GenModule Unit
ms_mod ModSummary
summary))
      mutagenLog $ "Plugin started on module " <> modName
      (modAST', modMetadata) <- instrumentModule cli flags modAST
      let tracerMetadata = String -> ModuleMetadata -> TracerMetadata
mkSingleModuleTracerMetadata String
modName ModuleMetadata
modMetadata
      liftIO $ saveTracerMetadata tracerMetadata
      mutagenLog "Done"
      return (ParsedResult (source{hpm_module = L loc modAST'}) msgs)

{-------------------------------------------------------------------------------
-- * Instrumentation
-------------------------------------------------------------------------------}

-- | Instrument a parsed module.
instrumentModule
  :: [CommandLineOption]
  -> DynFlags
  -> HsModule GhcPs
  -> Hsc (HsModule GhcPs, ModuleMetadata)
instrumentModule :: [String]
-> DynFlags
-> HsModule GhcPs
-> Hsc (HsModule GhcPs, ModuleMetadata)
instrumentModule [String]
_cli DynFlags
flags HsModule GhcPs
modAST = do
  (modAST', nodes) <- WriterT [NodeMetadata] Hsc (HsModule GhcPs)
-> Hsc (HsModule GhcPs, [NodeMetadata])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT
    (WriterT [NodeMetadata] Hsc (HsModule GhcPs)
 -> Hsc (HsModule GhcPs, [NodeMetadata]))
-> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
-> Hsc (HsModule GhcPs, [NodeMetadata])
forall a b. (a -> b) -> a -> b
$ case [RdrName]
traceAnnotations of
      [] -> do
        String -> WriterT [NodeMetadata] Hsc ()
forall (m :: * -> *). MonadIO m => String -> m ()
mutagenLog String
"Run mode: full module"
        HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
instrumentEntireModule HsModule GhcPs
modAST
      [RdrName]
bindings -> do
        String -> WriterT [NodeMetadata] Hsc ()
forall (m :: * -> *). MonadIO m => String -> m ()
mutagenLog (String -> WriterT [NodeMetadata] Hsc ())
-> String -> WriterT [NodeMetadata] Hsc ()
forall a b. (a -> b) -> a -> b
$ String
"Run mode: trace only " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> [RdrName] -> String
forall a. Outputable a => DynFlags -> a -> String
showPpr DynFlags
flags [RdrName]
bindings
        [RdrName]
-> HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
instrumentTopLevelBindings [RdrName]
bindings HsModule GhcPs
modAST
  let metadata = [NodeMetadata] -> ModuleMetadata
ModuleMetadata [NodeMetadata]
nodes
  return (modAST', metadata)
  where
    -- Extract all the TRACE annotations from this module.
    traceAnnotations :: [RdrName]
traceAnnotations = AnnDecl GhcPs -> RdrName
extractAnn (AnnDecl GhcPs -> RdrName) -> [AnnDecl GhcPs] -> [RdrName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (AnnDecl GhcPs -> Bool) -> GenericQ [AnnDecl GhcPs]
forall r. Typeable r => (r -> Bool) -> GenericQ [r]
listify (DynFlags -> AnnDecl GhcPs -> Bool
isAnn DynFlags
flags) HsModule GhcPs
modAST

    -- Add an import to the module exporting @trace@, so it's always in scope
    addTracerModuleImport
      :: HsModule GhcPs
      -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
    addTracerModuleImport :: HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
addTracerModuleImport HsModule GhcPs
m = do
      String -> WriterT [NodeMetadata] Hsc ()
forall (m :: * -> *). MonadIO m => String -> m ()
mutagenLog (String -> WriterT [NodeMetadata] Hsc ())
-> String -> WriterT [NodeMetadata] Hsc ()
forall a b. (a -> b) -> a -> b
$ String
"Adding module import for" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> ModuleName -> String
forall a. Outputable a => DynFlags -> a -> String
showPpr DynFlags
flags ModuleName
tracerModuleName
      let tracerModuleImport :: GenLocated SrcSpanAnnA (ImportDecl GhcPs)
tracerModuleImport =
            ImportDecl GhcPs -> GenLocated SrcSpanAnnA (ImportDecl GhcPs)
forall e a. HasAnnotation e => a -> GenLocated e a
noLocA
              ( (ModuleName -> ImportDecl GhcPs
simpleImportDecl ModuleName
tracerModuleName)
                  { ideclQualified = QualifiedPre
                  }
              )
      HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
forall a. a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a. Monad m => a -> m a
return HsModule GhcPs
m{hsmodImports = tracerModuleImport : hsmodImports m}

    -- Instrument and entire module
    instrumentEntireModule
      :: HsModule GhcPs
      -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
    instrumentEntireModule :: HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
instrumentEntireModule =
      HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
addTracerModuleImport
        (HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs))
-> (HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs))
-> HsModule GhcPs
-> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
forall a. Data a => a -> WriterT [NodeMetadata] Hsc a
instrumentEverywhere

    -- Instrument only specific bindings
    instrumentTopLevelBindings
      :: [RdrName]
      -> HsModule GhcPs
      -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
    instrumentTopLevelBindings :: [RdrName]
-> HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
instrumentTopLevelBindings [RdrName]
bindings =
      HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
addTracerModuleImport
        (HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs))
-> (HsModule GhcPs -> WriterT [NodeMetadata] Hsc (HsModule GhcPs))
-> HsModule GhcPs
-> WriterT [NodeMetadata] Hsc (HsModule GhcPs)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (forall a. Data a => a -> WriterT [NodeMetadata] Hsc a)
-> forall a. Data a => a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *). Monad m => GenericM m -> GenericM m
everywhereM ((Match GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))
 -> WriterT
      [NodeMetadata]
      Hsc
      (Match GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))))
-> a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(b -> m b) -> a -> m a
mkM ([RdrName]
-> Match GhcPs (LHsExpr GhcPs)
-> WriterT [NodeMetadata] Hsc (Match GhcPs (LHsExpr GhcPs))
instrumentTopLevelBinding [RdrName]
bindings))

    -- Instrument top-level functions having TRACE annotation pragmas
    instrumentTopLevelBinding
      :: [RdrName]
      -> Match GhcPs (LHsExpr GhcPs)
      -> WriterT [NodeMetadata] Hsc (Match GhcPs (LHsExpr GhcPs))
    instrumentTopLevelBinding :: [RdrName]
-> Match GhcPs (LHsExpr GhcPs)
-> WriterT [NodeMetadata] Hsc (Match GhcPs (LHsExpr GhcPs))
instrumentTopLevelBinding [RdrName]
annotations Match GhcPs (LHsExpr GhcPs)
match =
      case Match GhcPs (LHsExpr GhcPs)
match of
        Match XCMatch GhcPs (LHsExpr GhcPs)
m_x HsMatchContext (LIdP (NoGhcTc GhcPs))
m_ctx [LPat GhcPs]
m_ps GRHSs GhcPs (LHsExpr GhcPs)
m_bodies
          | HsMatchContext (GenLocated SrcSpanAnnN RdrName) -> Bool
forall id. HsMatchContext id -> Bool
isFunRhs HsMatchContext (LIdP (NoGhcTc GhcPs))
HsMatchContext (GenLocated SrcSpanAnnN RdrName)
m_ctx Bool -> Bool -> Bool
&& GenLocated SrcSpanAnnN RdrName -> RdrName
forall l e. GenLocated l e -> e
unLoc (HsMatchContext (GenLocated SrcSpanAnnN RdrName)
-> GenLocated SrcSpanAnnN RdrName
forall fn. HsMatchContext fn -> fn
mc_fun HsMatchContext (LIdP (NoGhcTc GhcPs))
HsMatchContext (GenLocated SrcSpanAnnN RdrName)
m_ctx) RdrName -> [RdrName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [RdrName]
annotations -> do
              m_bodies' <- GRHSs GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))
-> WriterT
     [NodeMetadata]
     Hsc
     (GRHSs GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs)))
forall a. Data a => a -> WriterT [NodeMetadata] Hsc a
instrumentEverywhere GRHSs GhcPs (LHsExpr GhcPs)
GRHSs GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))
m_bodies
              return (Match m_x m_ctx m_ps m_bodies')
        Match GhcPs (LHsExpr GhcPs)
x -> Match GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))
-> WriterT
     [NodeMetadata]
     Hsc
     (Match GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs)))
forall a. a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a. Monad m => a -> m a
return Match GhcPs (LHsExpr GhcPs)
Match GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))
x

    -- Recursively instrument every sub-expression
    instrumentEverywhere
      :: (Data a)
      => a
      -> WriterT [NodeMetadata] Hsc a
    instrumentEverywhere :: forall a. Data a => a -> WriterT [NodeMetadata] Hsc a
instrumentEverywhere =
      (forall a. Data a => a -> WriterT [NodeMetadata] Hsc a)
-> forall a. Data a => a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *). Monad m => GenericM m -> GenericM m
everywhereM ((GRHS GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))
 -> WriterT
      [NodeMetadata]
      Hsc
      (GRHS GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))))
-> a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(b -> m b) -> a -> m a
mkM GRHS GhcPs (LHsExpr GhcPs)
-> WriterT [NodeMetadata] Hsc (GRHS GhcPs (LHsExpr GhcPs))
GRHS GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs))
-> WriterT
     [NodeMetadata]
     Hsc
     (GRHS GhcPs (GenLocated SrcSpanAnnA (HsExpr GhcPs)))
instrumentGRHSs)
        (a -> WriterT [NodeMetadata] Hsc a)
-> (a -> WriterT [NodeMetadata] Hsc a)
-> a
-> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (forall a. Data a => a -> WriterT [NodeMetadata] Hsc a)
-> forall a. Data a => a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *). Monad m => GenericM m -> GenericM m
everywhereM ((HsExpr GhcPs -> WriterT [NodeMetadata] Hsc (HsExpr GhcPs))
-> a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(b -> m b) -> a -> m a
mkM HsExpr GhcPs -> WriterT [NodeMetadata] Hsc (HsExpr GhcPs)
instrumentIFs)

    -- Instrument every RHS with a tracer node.
    -- These come after function clauses, case statements, multi-way ifs, etc.
    instrumentGRHSs
      :: GRHS GhcPs (LHsExpr GhcPs)
      -> WriterT [NodeMetadata] Hsc (GRHS GhcPs (LHsExpr GhcPs))
    instrumentGRHSs :: GRHS GhcPs (LHsExpr GhcPs)
-> WriterT [NodeMetadata] Hsc (GRHS GhcPs (LHsExpr GhcPs))
instrumentGRHSs (GRHS XCGRHS GhcPs (LHsExpr GhcPs)
ext [GuardLStmt GhcPs]
guards LHsExpr GhcPs
rhs) = do
      rhsNode <- IO Int -> WriterT [NodeMetadata] Hsc Int
forall a. IO a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Int
freshTraceNode
      let rhsLoc = GenLocated SrcSpanAnnA (HsExpr GhcPs) -> SrcSpan
forall a e. HasLoc a => GenLocated a e -> SrcSpan
getLocA LHsExpr GhcPs
GenLocated SrcSpanAnnA (HsExpr GhcPs)
rhs
      tell [NodeMetadata rhsNode GRHSNode (srcSpanToNodeLocation rhsLoc)]
      logInstrumentedNode "RHS" rhsNode rhsLoc
      let rhs' = Int -> LHsExpr GhcPs -> LHsExpr GhcPs
wrapTracer Int
rhsNode LHsExpr GhcPs
rhs
      return (GRHS ext guards rhs')

    -- Instrument each branch of an if-then-else expression with a tracer
    instrumentIFs
      :: HsExpr GhcPs
      -> WriterT [NodeMetadata] Hsc (HsExpr GhcPs)
    instrumentIFs :: HsExpr GhcPs -> WriterT [NodeMetadata] Hsc (HsExpr GhcPs)
instrumentIFs HsExpr GhcPs
expr =
      case HsExpr GhcPs
expr of
        HsIf XIf GhcPs
ext LHsExpr GhcPs
cond LHsExpr GhcPs
th LHsExpr GhcPs
el -> do
          -- then branch
          thNode <- IO Int -> WriterT [NodeMetadata] Hsc Int
forall a. IO a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Int
freshTraceNode
          let thLoc = GenLocated SrcSpanAnnA (HsExpr GhcPs) -> SrcSpan
forall a e. HasLoc a => GenLocated a e -> SrcSpan
getLocA LHsExpr GhcPs
GenLocated SrcSpanAnnA (HsExpr GhcPs)
th
          tell [NodeMetadata thNode ThenNode (srcSpanToNodeLocation thLoc)]
          logInstrumentedNode "then branch" thNode thLoc
          let th' = Int -> LHsExpr GhcPs -> LHsExpr GhcPs
wrapTracer Int
thNode LHsExpr GhcPs
th
          -- else branch
          elNode <- liftIO freshTraceNode
          let elLoc = GenLocated SrcSpanAnnA (HsExpr GhcPs) -> SrcSpan
forall a e. HasLoc a => GenLocated a e -> SrcSpan
getLocA LHsExpr GhcPs
GenLocated SrcSpanAnnA (HsExpr GhcPs)
el
          tell [NodeMetadata elNode ElseNode (srcSpanToNodeLocation elLoc)]
          logInstrumentedNode "else branch" elNode elLoc
          let el' = Int -> LHsExpr GhcPs -> LHsExpr GhcPs
wrapTracer Int
elNode LHsExpr GhcPs
el
          -- wrap it up again
          return (HsIf ext cond th' el')
        HsExpr GhcPs
x -> HsExpr GhcPs -> WriterT [NodeMetadata] Hsc (HsExpr GhcPs)
forall a. a -> WriterT [NodeMetadata] Hsc a
forall (m :: * -> *) a. Monad m => a -> m a
return HsExpr GhcPs
x

    -- Log an instrumentation message
    logInstrumentedNode
      :: String
      -> TraceNode
      -> SrcSpan
      -> WriterT [NodeMetadata] Hsc ()
    logInstrumentedNode :: String -> Int -> SrcSpan -> WriterT [NodeMetadata] Hsc ()
logInstrumentedNode String
reason Int
node SrcSpan
loc = do
      Hsc () -> WriterT [NodeMetadata] Hsc ()
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [NodeMetadata] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift
        (Hsc () -> WriterT [NodeMetadata] Hsc ())
-> Hsc () -> WriterT [NodeMetadata] Hsc ()
forall a b. (a -> b) -> a -> b
$ String -> Hsc ()
forall (m :: * -> *). MonadIO m => String -> m ()
mutagenLog
        (String -> Hsc ()) -> String -> Hsc ()
forall a b. (a -> b) -> a -> b
$ String
"Inoculating tracer #"
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
node
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" on "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
reason
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" at "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SrcSpan -> String
forall a. Outputable a => DynFlags -> a -> String
showPpr DynFlags
flags SrcSpan
loc

{-------------------------------------------------------------------------------
-- * Helpers
-------------------------------------------------------------------------------}

-- ** Trace node generation

-- | Global counter for trace nodes.
traceNodeCounter :: IORef TraceNode
traceNodeCounter :: IORef Int
traceNodeCounter = IO (IORef Int) -> IORef Int
forall a. IO a -> a
unsafePerformIO (Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0)
{-# NOINLINE traceNodeCounter #-}

-- | Generate a fresh trace node.
freshTraceNode :: IO TraceNode
freshTraceNode :: IO Int
freshTraceNode = IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
traceNodeCounter ((Int -> (Int, Int)) -> IO Int) -> (Int -> (Int, Int)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Int
n ->
  (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE freshTraceNode #-}

-- ** Logging

{- FOURMOLU_DISABLE -}
-- | Print a message if debugging is enabled.
mutagenLog :: MonadIO m => String -> m ()
mutagenLog :: forall (m :: * -> *). MonadIO m => String -> m ()
mutagenLog String
_str =
#ifdef MUTAGEN_PLUGIN_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 -}

-- ** Annotation pragmas

-- | Pattern for matching against annotation pragmas.
pattern HsAnn :: RdrName -> RdrName -> AnnDecl GhcPs
pattern $mHsAnn :: forall {r}.
AnnDecl GhcPs -> (RdrName -> RdrName -> r) -> ((# #) -> r) -> r
HsAnn lhs rhs <-
  HsAnnotation
    _
    (ValueAnnProvenance (L _ lhs))
    (L _ (HsVar _ (L _ rhs)))

-- | Check whether an annotation pragma is of the shape:
-- {-# ANN ident TRACE #-}
isAnn :: DynFlags -> AnnDecl GhcPs -> Bool
isAnn :: DynFlags -> AnnDecl GhcPs -> Bool
isAnn DynFlags
flags (HsAnn RdrName
_ RdrName
rhs) = DynFlags -> RdrName -> String
forall a. Outputable a => DynFlags -> a -> String
showPpr DynFlags
flags RdrName
rhs String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== DynFlags -> OccName -> String
forall a. Outputable a => DynFlags -> a -> String
showPpr DynFlags
flags OccName
tracerAnnName
isAnn DynFlags
_ AnnDecl GhcPs
_ = Bool
False

-- | Extract the target of an annotation pragma.
extractAnn :: AnnDecl GhcPs -> RdrName
extractAnn :: AnnDecl GhcPs -> RdrName
extractAnn (HsAnn RdrName
target RdrName
_) = RdrName
target
extractAnn AnnDecl GhcPs
_ = String -> RdrName
forall a. HasCallStack => String -> a
error String
"this should not happen"

-- ** Source locations

-- | Turn a generic 'SrcSpan' into something more amenable to serialization.
srcSpanToNodeLocation :: SrcSpan -> Maybe NodeLocation
srcSpanToNodeLocation :: SrcSpan -> Maybe NodeLocation
srcSpanToNodeLocation SrcSpan
loc =
  case SrcSpan
loc of
    RealSrcSpan RealSrcSpan
realLoc Maybe BufSpan
_ ->
      NodeLocation -> Maybe NodeLocation
forall a. a -> Maybe a
Just
        (NodeLocation -> Maybe NodeLocation)
-> NodeLocation -> Maybe NodeLocation
forall a b. (a -> b) -> a -> b
$ NodeLocation
          { filePath :: String
filePath = FastString -> String
unpackFS (RealSrcSpan -> FastString
srcSpanFile RealSrcSpan
realLoc)
          , startLine :: Int
startLine = RealSrcSpan -> Int
srcSpanStartLine RealSrcSpan
realLoc
          , startCol :: Int
startCol = RealSrcSpan -> Int
srcSpanStartCol RealSrcSpan
realLoc
          , endLine :: Int
endLine = RealSrcSpan -> Int
srcSpanEndLine RealSrcSpan
realLoc
          , endCol :: Int
endCol = RealSrcSpan -> Int
srcSpanEndCol RealSrcSpan
realLoc
          }
    SrcSpan
_ -> Maybe NodeLocation
forall a. Maybe a
Nothing

-- ** Predicates

-- | Is this a pattern matching an argument of a function binding?
isFunRhs :: HsMatchContext id -> Bool
isFunRhs :: forall id. HsMatchContext id -> Bool
isFunRhs (FunRhs{}) = Bool
True
isFunRhs HsMatchContext id
_ = Bool
False

-- ** Constants

-- | Module name of the tracing module.
tracerModuleName :: ModuleName
tracerModuleName :: ModuleName
tracerModuleName = String -> ModuleName
mkModuleName String
"Test.Mutagen.Tracer"

-- | Name of the tracing function.
tracerFunName :: OccName
tracerFunName :: OccName
tracerFunName = String -> OccName
mkVarOcc String
"trace"

-- | Name of the tracing annotation.
tracerAnnName :: OccName
tracerAnnName :: OccName
tracerAnnName = String -> OccName
mkVarOcc String
"TRACE"

-- ** AST Builders

-- | Build a variable expression from a 'RdrName'.
var :: RdrName -> LHsExpr GhcPs
var :: RdrName -> LHsExpr GhcPs
var RdrName
v = HsExpr GhcPs -> GenLocated SrcSpanAnnA (HsExpr GhcPs)
forall e a. HasAnnotation e => a -> GenLocated e a
noLocA (XVar GhcPs -> LIdP GhcPs -> HsExpr GhcPs
forall p. XVar p -> LIdP p -> HsExpr p
HsVar XVar GhcPs
NoExtField
noExtField (RdrName -> GenLocated SrcSpanAnnN RdrName
forall e a. HasAnnotation e => a -> GenLocated e a
noLocA RdrName
v))

-- | Wrap an expression in parentheses.
paren :: LHsExpr GhcPs -> LHsExpr GhcPs
paren :: LHsExpr GhcPs -> LHsExpr GhcPs
paren LHsExpr GhcPs
x = HsExpr GhcPs -> GenLocated SrcSpanAnnA (HsExpr GhcPs)
forall e a. HasAnnotation e => a -> GenLocated e a
noLocA (LHsExpr GhcPs -> HsExpr GhcPs
forall (p :: Pass).
IsPass p =>
LHsExpr (GhcPass p) -> HsExpr (GhcPass p)
gHsPar LHsExpr GhcPs
x)

-- | Apply one expression to another.
app :: LHsExpr GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs
#if MIN_VERSION_ghc(9,10,1)
app :: LHsExpr GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs
app LHsExpr GhcPs
x LHsExpr GhcPs
y = HsExpr GhcPs -> GenLocated SrcSpanAnnA (HsExpr GhcPs)
forall e a. HasAnnotation e => a -> GenLocated e a
noLocA (XApp GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs -> HsExpr GhcPs
forall p. XApp p -> LHsExpr p -> LHsExpr p -> HsExpr p
HsApp XApp GhcPs
NoExtField
GHC.noExtField LHsExpr GhcPs
x LHsExpr GhcPs
y)
#else
app x y = noLocA (HsApp GHC.noComments x y)
#endif

infixl 5 `app`

-- | Build a numeric literal expression.
numLit :: Int -> LHsExpr GhcPs
#if MIN_VERSION_ghc(9,10,1)
numLit :: Int -> LHsExpr GhcPs
numLit Int
n = HsExpr GhcPs -> GenLocated SrcSpanAnnA (HsExpr GhcPs)
forall e a. HasAnnotation e => a -> GenLocated e a
noLocA (XLitE GhcPs -> HsLit GhcPs -> HsExpr GhcPs
forall p. XLitE p -> HsLit p -> HsExpr p
HsLit XLitE GhcPs
NoExtField
GHC.noExtField (XHsInt GhcPs -> IntegralLit -> HsLit GhcPs
forall x. XHsInt x -> IntegralLit -> HsLit x
HsInt XHsInt GhcPs
NoExtField
noExtField (Int -> IntegralLit
forall a. Integral a => a -> IntegralLit
mkIntegralLit Int
n)))
#else
numLit n = noLocA (HsLit GHC.noComments (HsInt noExtField (mkIntegralLit n)))
#endif

-- | Wrap an expression with the tracer function.
wrapTracer :: TraceNode -> LHsExpr GhcPs -> LHsExpr GhcPs
wrapTracer :: Int -> LHsExpr GhcPs -> LHsExpr GhcPs
wrapTracer Int
node LHsExpr GhcPs
expr =
  RdrName -> LHsExpr GhcPs
var (ModuleName -> OccName -> RdrName
Qual ModuleName
tracerModuleName OccName
tracerFunName)
    LHsExpr GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs
`app` Int -> LHsExpr GhcPs
numLit Int
node
    LHsExpr GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs
`app` LHsExpr GhcPs -> LHsExpr GhcPs
paren LHsExpr GhcPs
expr