@@ -1,74 +1,109 @@
1+{-# LANGUAGE RankNTypes #-}
1{-# LANGUAGE OverloadedStrings #-}
2{-# LANGUAGE OverloadedStrings #-}
2module StackTrace.Plugin (plugin) where
3module StackTrace.Plugin (plugin) where
5+import Control.Arrow (first)
6+import Data.Monoid (Any(Any, getAny))
4import GhcPlugins
7import GhcPlugins
5−import TcRnTypes (IfM, TcM, TcGblEnv, tcg_binds, tcg_rn_decls)
6−import HsExtension (GhcTc, GhcRn)
7−import HsDecls (HsGroup)
8−import HsExpr (LHsExpr)
9import HsSyn
8import HsSyn
11−import Data.Maybe
10+type Traversal s t a b
11+ = forall f. Applicative f =>
12+ (a -> f b) -> s -> f t
14+type Traversal' s a = Traversal s s a a
13plugin :: Plugin
16plugin :: Plugin
14−plugin = defaultPlugin
15− { parsedResultAction = parsedPlugin
17+plugin = defaultPlugin {parsedResultAction = parsedPlugin, pluginRecompile = purePlugin}
18−parsedPlugin :: [CommandLineOption] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
20+ [CommandLineOption] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
19parsedPlugin _ _ pm = do
21parsedPlugin _ _ pm = do
20− dflags <- getDynFlags
22+ let m = updateHsModule <$> hpm_module pm
21−
23+ pm' = pm {hpm_module = m}
22− let extract = hsmodImports . unLoc
24+ return pm'
23− addGHCStackMod = noLoc $ simpleImportDecl $ mkModuleName "GHC.Stack"
25− m = updateHsModule addGHCStackMod updateHsmodDecl <$> hpm_module pm
26+-- Use qualified import for GHC.Stack as "AutoImported.GHC.Stack"
26− pm' = pm { hpm_module = m }
27+-- ...this should not interfere with other imports...
28+ghcStackModuleName :: ModuleName
29+ghcStackModuleName = mkModuleName "AutoImported.GHC.Stack"
28− return pm'
31+ghcStackImport :: Located (ImportDecl (GhcPass p))
34+ (simpleImportDecl $ mkModuleName "GHC.Stack")
35+ {ideclQualified = True, ideclAs = Just $ noLoc ghcStackModuleName}
30−updateHsModule :: LImportDecl GhcPs -> (LHsDecl GhcPs -> LHsDecl GhcPs) -> HsModule GhcPs -> HsModule GhcPs
37+updateHsModule :: HsModule GhcPs -> HsModule GhcPs
31−updateHsModule importDecl update hsm = hsm
38+updateHsModule hsm =
32− { hsmodImports = importDecl:decls
39+ hsm {hsmodImports = hsmodImports', hsmodDecls = hsmodDecls'}
33− , hsmodDecls = map update lhss
36− decls = hsmodImports hsm
41+ -- Traverse the haskell AST; if we have to add some HasStack
37− lhss = hsmodDecls hsm
42+ -- constraint we set a flag in a (Any,) functor.
43+ -- ...it'd be simpler to check if before == after, but Haskell AST
44+ -- doesn't have Eq instances.
45+ (updatedP, hsmodDecls') =
47+ (traverse . astTraversal) updateHsType (hsmodDecls hsm)
49+ -- Only import GHC.Stack if needed for a constraint we introduced
52+ then [ghcStackImport]
56+ astTraversal :: Traversal' (LHsDecl GhcPs) (HsType GhcPs)
57+ astTraversal = updateHsmodDecl
39--------------
64--------------
65+updateHsmodDecl :: Traversal' (LHsDecl GhcPs) (HsDecl GhcPs)
66+updateHsmodDecl = traverse
41−updateHsmodDecl :: LHsDecl GhcPs -> LHsDecl GhcPs
42−updateHsmodDecl = fmap updateHsDecl
44−updateHsDecl :: HsDecl GhcPs -> HsDecl GhcPs
68+updateHsDecl :: Traversal' (HsDecl GhcPs) (Sig GhcPs)
45−updateHsDecl (SigD xSig s) = SigD xSig (updateSig s)
69+updateHsDecl f (SigD xSig s) = SigD xSig <$> f s
46−updateHsDecl decl = decl
70+updateHsDecl _ sig = pure sig
48−updateSig :: Sig GhcPs -> Sig GhcPs
72+updateSig :: Traversal' (Sig GhcPs) (LHsSigWcType GhcPs)
49−updateSig (TypeSig xSig ls t) = TypeSig xSig ls (updateLHsSigWcType t)
73+updateSig f (TypeSig xSig ls t) = TypeSig xSig ls <$> f t
50−updateSig sig = sig
74+updateSig _ sig = pure sig
52−updateLHsSigWcType :: LHsSigWcType GhcPs -> LHsSigWcType GhcPs
76+updateLHsSigWsType :: Traversal' (LHsSigWcType GhcPs) (LHsSigType GhcPs)
77+updateLHsSigWsType f lhs@HsWC {} =
53−updateLHsSigWcType lhs@HsWC{} = lhs { hswc_body = updateLHsSigType (hswc_body lhs) }
78+ (\x -> lhs {hswc_body = x}) <$> f (hswc_body lhs)
54−updateLHsSigWcType lhs@XHsWildCardBndrs{} = lhs
79+updateLHsSigWsType _ lhs = pure lhs
56−updateLHsSigType :: LHsSigType GhcPs -> LHsSigType GhcPs
81+updateLHsSigType :: Traversal' (LHsSigType GhcPs) (LHsType GhcPs)
82+updateLHsSigType f lhs@HsIB {} =
57−updateLHsSigType lhs@HsIB{} = lhs { hsib_body = updateLHsType (hsib_body lhs )}
83+ (\x -> lhs {hsib_body = x}) <$> f (hsib_body lhs)
58−updateLHsSigType lhs@XHsImplicitBndrs{} = lhs
84+updateLHsSigType _ lhs = pure lhs
60−updateLHsType :: LHsType GhcPs -> LHsType GhcPs
86+updateLHsType :: Traversal' (LHsType GhcPs) (HsType GhcPs)
61−updateLHsType = fmap updateHsType
87+updateLHsType = traverse
63-- Main process
89-- Main process
64−updateHsType :: HsType GhcPs -> HsType GhcPs
90+updateHsType :: HsType GhcPs -> (Any, HsType GhcPs)
65−updateHsType ty@(HsQualTy xty ctxt body) = HsQualTy xty (fmap appendHSC ctxt) body
91+updateHsType (HsQualTy xty ctxt body) =
92+ flagASTModified $ HsQualTy xty (fmap appendHSC ctxt) body
93+updateHsType ty@HsTyVar {} =
66−updateHsType ty@HsTyVar{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
94+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
95+updateHsType ty@HsAppTy {} =
67−updateHsType ty@HsAppTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
96+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
97+updateHsType ty@HsFunTy {} =
68−updateHsType ty@HsFunTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
98+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
99+updateHsType ty@HsListTy {} =
69−updateHsType ty@HsListTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
100+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
101+updateHsType ty@HsTupleTy {} =
70−updateHsType ty@HsTupleTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
102+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
71−updateHsType ty = ty
103+updateHsType ty = pure ty
105+flagASTModified :: a -> (Any, a)
106+flagASTModified a = (Any True, a)
73appendHSC :: HsContext GhcPs -> HsContext GhcPs
108appendHSC :: HsContext GhcPs -> HsContext GhcPs
74appendHSC cs = mkHSC : cs
109appendHSC cs = mkHSC : cs
@@ -78,4 +113,4 @@ mkHSC :: LHsType GhcPs
78mkHSC = noLoc $ HsTyVar noExt NotPromoted lId
113mkHSC = noLoc $ HsTyVar noExt NotPromoted lId
80lId :: Located (IdP GhcPs)
115lId :: Located (IdP GhcPs)
81−lId = noLoc $ mkRdrUnqual $ mkClsOcc "HasCallStack"
116+lId = noLoc $ mkRdrQual ghcStackModuleName $ mkClsOcc "HasCallStack"