{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE LambdaCase #-}
module AutoInstrument.Internal.Plugin.Parser
( parsedResultAction
) where
import Control.Monad.IO.Class (liftIO)
import qualified Data.ByteString.Char8 as BS8
import Data.Maybe (mapMaybe)
import qualified Data.Set as S
import qualified AutoInstrument.Internal.GhcFacade as Ghc
import qualified AutoInstrument.Internal.Config as Cfg
parsedResultAction
:: [Ghc.CommandLineOption]
-> Ghc.ModSummary
-> Ghc.ParsedResult
-> Ghc.Hsc Ghc.ParsedResult
parsedResultAction opts modSummary
parsedResult@Ghc.ParsedResult
{Ghc.parsedResultModule = prm@Ghc.HsParsedModule
{Ghc.hpm_module = Ghc.L modLoc mo@Ghc.HsModule{Ghc.hsmodDecls}}} = do
let modName = Ghc.moduleName $ Ghc.ms_mod modSummary
unitId = Ghc.toUnitId . Ghc.moduleUnit $ Ghc.ms_mod modSummary
hscEnv <- Ghc.getHscEnv
result <- liftIO $
Ghc.findImportedModule hscEnv (Ghc.mkModuleName "AutoInstrument.Internal.Types") Ghc.NoPkgQual
otelMod <-
case result of
Ghc.Found _ m -> pure m
_ -> error "AutoInstrument.Internal.Types module not found"
let occ = Ghc.mkVarOcc "autoInstrument"
autoInstrumentName <- liftIO $ Ghc.lookupNameCache (Ghc.hsc_NC hscEnv) otelMod occ
mConfig <- liftIO $ fmap Cfg.getConfig <$> Cfg.getConfigCache opts
case mConfig of
Nothing -> pure parsedResult
{ Ghc.parsedResultMessages = (Ghc.parsedResultMessages parsedResult)
{ Ghc.psErrors =
let msg = Ghc.mkParseError "Failed to load auto instrumentation config"
in Ghc.addMessage msg . Ghc.psErrors $ Ghc.parsedResultMessages parsedResult
}
}
Just config -> do
let matches = S.fromList $ getMatches config hsmodDecls
newDecls = instrumentDecl modName unitId autoInstrumentName matches <$> hsmodDecls
pure parsedResult
{ Ghc.parsedResultModule = prm
{ Ghc.hpm_module = Ghc.L modLoc mo
{ Ghc.hsmodDecls = newDecls
}
}
}
getMatches
:: Cfg.Config
-> [Ghc.LHsDecl Ghc.GhcPs]
-> [Ghc.OccName]
getMatches cfg = concat . mapMaybe go where
go (Ghc.L _ (Ghc.SigD _ (Ghc.TypeSig _ lhs (Ghc.HsWC _ (Ghc.L _ (Ghc.HsSig _ _ (Ghc.L _ ty)))))))
| isTargetTy [] ty = Just (Ghc.rdrNameOcc . Ghc.unLoc <$> lhs)
go _ = Nothing
isTargetTy preds = \case
Ghc.HsForAllTy _ _ (Ghc.L _ body) -> isTargetTy preds body
Ghc.HsQualTy _ (Ghc.L _ ctx) (Ghc.L _ body) ->
isTargetTy (preds ++ fmap Ghc.unLoc ctx) body
app@Ghc.HsAppTy{} -> check preds app
var@Ghc.HsTyVar{} -> check preds var
Ghc.HsFunTy _ _ _ (Ghc.L _ nxt) -> isTargetTy preds nxt
Ghc.HsParTy _ (Ghc.L _ nxt) -> isTargetTy preds nxt
Ghc.HsDocTy _ (Ghc.L _ nxt) _ -> isTargetTy preds nxt
_ -> False
check
:: [Ghc.HsType Ghc.GhcPs]
-> Ghc.HsType Ghc.GhcPs
-> Bool
check preds expr =
any (matchTarget preds expr) (Cfg.targets cfg)
&& not (any (matchTarget preds expr) (Cfg.exclusions cfg))
matchTarget preds expr = \case
Cfg.Constructor conTarget -> checkTy True conTarget expr
Cfg.Constraints predTarget -> checkPred preds predTarget
checkTy
:: Bool
-> Cfg.TargetCon
-> Ghc.HsType Ghc.GhcPs
-> Bool
checkTy top t (Ghc.HsParTy _ (Ghc.L _ x)) = checkTy top t x
checkTy top t (Ghc.HsDocTy _ (Ghc.L _ x) _) = checkTy top t x
checkTy _ (Cfg.TyVar name) (Ghc.HsTyVar _ _ (Ghc.L _ rdrName)) =
BS8.pack name == Ghc.bytesFS (Ghc.occNameFS $ Ghc.rdrNameOcc rdrName)
checkTy top target@(Cfg.App x y) (Ghc.HsAppTy _ (Ghc.L _ con) (Ghc.L _ arg)) =
(checkTy False y arg && checkTy False x con )
|| (top && checkTy True target con)
checkTy True target@(Cfg.TyVar _) (Ghc.HsAppTy _ (Ghc.L _ con) _) =
checkTy True target con
checkTy _ Cfg.Unit (Ghc.HsTupleTy _ Ghc.HsBoxedOrConstraintTuple []) = True
checkTy _ (Cfg.Tuple targets) (Ghc.HsTupleTy _ Ghc.HsBoxedOrConstraintTuple exprs) =
and $ zipWith (checkTy False) targets (Ghc.unLoc <$> exprs)
checkTy _ Cfg.WC _ = True
checkTy _ _ _ = False
checkPred
:: [Ghc.HsType Ghc.GhcPs]
-> Cfg.ConstraintSet
-> Bool
checkPred preds predSet =
all (\p -> any (checkTy True p) preds)
(S.toList predSet)
instrumentDecl
:: Ghc.ModuleName
-> Ghc.UnitId
-> Ghc.Name
-> S.Set Ghc.OccName
-> Ghc.LHsDecl Ghc.GhcPs
-> Ghc.LHsDecl Ghc.GhcPs
instrumentDecl modName unitId instrName targets
(Ghc.L loc (Ghc.ValD vX fb@Ghc.FunBind
{ Ghc.fun_matches = mg@Ghc.MG
{ Ghc.mg_alts = Ghc.L altsLoc alts }, Ghc.fun_id}))
| Ghc.rdrNameOcc (Ghc.unLoc fun_id) `S.member` targets
= let newAlts = (fmap . fmap)
(instrumentMatch modName unitId (Ghc.unLoc fun_id) instrName)
alts
in Ghc.L loc (Ghc.ValD vX (fb
{ Ghc.fun_matches = mg
{ Ghc.mg_alts = Ghc.L altsLoc newAlts }}))
instrumentDecl _ _ _ _ x = x
instrumentMatch
:: Ghc.ModuleName
-> Ghc.UnitId
-> Ghc.RdrName
-> Ghc.Name
-> Ghc.Match Ghc.GhcPs (Ghc.GenLocated Ghc.SrcSpanAnnA (Ghc.HsExpr Ghc.GhcPs))
-> Ghc.Match Ghc.GhcPs (Ghc.GenLocated Ghc.SrcSpanAnnA (Ghc.HsExpr Ghc.GhcPs))
instrumentMatch modName unitId bindName instrName match =
match
{ Ghc.m_grhss = (Ghc.m_grhss match)
{ Ghc.grhssGRHSs = (fmap . fmap) modifyGRH (Ghc.grhssGRHSs (Ghc.m_grhss match)) }
}
where
modifyGRH :: Ghc.GRHS Ghc.GhcPs (Ghc.LHsExpr Ghc.GhcPs)
-> Ghc.GRHS Ghc.GhcPs (Ghc.LHsExpr Ghc.GhcPs)
modifyGRH (Ghc.GRHS x guards body) =
Ghc.GRHS x guards (go body)
go :: Ghc.LHsExpr Ghc.GhcPs -> Ghc.LHsExpr Ghc.GhcPs
go (Ghc.L loc x) =
let instrVar = Ghc.HsVar Ghc.noExtField (Ghc.L Ghc.noSrcSpanA (Ghc.Exact instrName))
mkStringExpr = Ghc.L Ghc.noSrcSpanA . Ghc.HsLit Ghc.noAnn
. Ghc.HsString Ghc.NoSourceText
app :: Ghc.LHsExpr Ghc.GhcPs -> Ghc.LHsExpr Ghc.GhcPs -> Ghc.LHsExpr Ghc.GhcPs
app l r = Ghc.L Ghc.noSrcSpanA $ Ghc.HsApp Ghc.noAnn l r
srcSpan = Ghc.realSrcSpan . Ghc.locA $ loc :: Ghc.RealSrcSpan
instr =
Ghc.L Ghc.noSrcSpanA instrVar
`app`
(mkStringExpr . Ghc.occNameFS $ Ghc.rdrNameOcc bindName)
`app`
mkStringExpr (Ghc.moduleNameFS modName)
`app`
mkStringExpr (Ghc.srcSpanFile srcSpan)
`app`
(mkStringExpr . Ghc.fsLit . show $ Ghc.srcSpanStartLine srcSpan)
`app`
mkStringExpr (Ghc.unitIdFS unitId)
in Ghc.L loc $ Ghc.HsApp Ghc.noAnn instr (Ghc.L loc x)