@@ -1,74 +1,109 @@
1+{-# LANGUAGE RankNTypes #-}
12{-# LANGUAGE OverloadedStrings #-}
23module StackTrace.Plugin (plugin) where
3−
4+
5+import Control.Arrow (first)
6+import Data.Monoid (Any(Any, getAny))
47import 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)
98import HsSyn
109
11−import Data.Maybe
10+type Traversal s t a b
11+ = forall f. Applicative f =>
12+ (a -> f b) -> s -> f t
13+
14+type Traversal' s a = Traversal s s a a
1215
1316plugin :: Plugin
14−plugin = defaultPlugin
15− { parsedResultAction = parsedPlugin
16− }
17+plugin = defaultPlugin {parsedResultAction = parsedPlugin, pluginRecompile = purePlugin}
1718
18−parsedPlugin :: [CommandLineOption] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
19+parsedPlugin ::
20+ [CommandLineOption] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
1921parsedPlugin _ _ pm = do
20− dflags <- getDynFlags
21−
22− let extract = hsmodImports . unLoc
23− addGHCStackMod = noLoc $ simpleImportDecl $ mkModuleName "GHC.Stack"
22+ let m = updateHsModule <$> hpm_module pm
23+ pm' = pm {hpm_module = m}
24+ return pm'
2425
25− m = updateHsModule addGHCStackMod updateHsmodDecl <$> hpm_module pm
26− pm' = pm { hpm_module = m }
26+-- Use qualified import for GHC.Stack as "AutoImported.GHC.Stack"
27+-- ...this should not interfere with other imports...
28+ghcStackModuleName :: ModuleName
29+ghcStackModuleName = mkModuleName "AutoImported.GHC.Stack"
2730
28− return pm'
31+ghcStackImport :: Located (ImportDecl (GhcPass p))
32+ghcStackImport =
33+ noLoc $
34+ (simpleImportDecl $ mkModuleName "GHC.Stack")
35+ {ideclQualified = True, ideclAs = Just $ noLoc ghcStackModuleName}
2936
30−updateHsModule :: LImportDecl GhcPs -> (LHsDecl GhcPs -> LHsDecl GhcPs) -> HsModule GhcPs -> HsModule GhcPs
31−updateHsModule importDecl update hsm = hsm
32− { hsmodImports = importDecl:decls
33− , hsmodDecls = map update lhss
34− }
37+updateHsModule :: HsModule GhcPs -> HsModule GhcPs
38+updateHsModule hsm =
39+ hsm {hsmodImports = hsmodImports', hsmodDecls = hsmodDecls'}
3540 where
36− decls = hsmodImports hsm
37− lhss = hsmodDecls hsm
41+ -- Traverse the haskell AST; if we have to add some HasStack
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') =
46+ first getAny $
47+ (traverse . astTraversal) updateHsType (hsmodDecls hsm)
48+
49+ -- Only import GHC.Stack if needed for a constraint we introduced
50+ hsmodImports' =
51+ (if updatedP
52+ then [ghcStackImport]
53+ else []) ++
54+ hsmodImports hsm
55+
56+ astTraversal :: Traversal' (LHsDecl GhcPs) (HsType GhcPs)
57+ astTraversal = updateHsmodDecl
58+ . updateHsDecl
59+ . updateSig
60+ . updateLHsSigWsType
61+ . updateLHsSigType
62+ . updateLHsType
3863
3964--------------
65+updateHsmodDecl :: Traversal' (LHsDecl GhcPs) (HsDecl GhcPs)
66+updateHsmodDecl = traverse
4067
41−updateHsmodDecl :: LHsDecl GhcPs -> LHsDecl GhcPs
42−updateHsmodDecl = fmap updateHsDecl
43−
44−updateHsDecl :: HsDecl GhcPs -> HsDecl GhcPs
45−updateHsDecl (SigD xSig s) = SigD xSig (updateSig s)
46−updateHsDecl decl = decl
68+updateHsDecl :: Traversal' (HsDecl GhcPs) (Sig GhcPs)
69+updateHsDecl f (SigD xSig s) = SigD xSig <$> f s
70+updateHsDecl _ sig = pure sig
4771
48−updateSig :: Sig GhcPs -> Sig GhcPs
49−updateSig (TypeSig xSig ls t) = TypeSig xSig ls (updateLHsSigWcType t)
50−updateSig sig = sig
72+updateSig :: Traversal' (Sig GhcPs) (LHsSigWcType GhcPs)
73+updateSig f (TypeSig xSig ls t) = TypeSig xSig ls <$> f t
74+updateSig _ sig = pure sig
5175
52−updateLHsSigWcType :: LHsSigWcType GhcPs -> LHsSigWcType GhcPs
53−updateLHsSigWcType lhs@HsWC{} = lhs { hswc_body = updateLHsSigType (hswc_body lhs) }
54−updateLHsSigWcType lhs@XHsWildCardBndrs{} = lhs
76+updateLHsSigWsType :: Traversal' (LHsSigWcType GhcPs) (LHsSigType GhcPs)
77+updateLHsSigWsType f lhs@HsWC {} =
78+ (\x -> lhs {hswc_body = x}) <$> f (hswc_body lhs)
79+updateLHsSigWsType _ lhs = pure lhs
5580
56−updateLHsSigType :: LHsSigType GhcPs -> LHsSigType GhcPs
57−updateLHsSigType lhs@HsIB{} = lhs { hsib_body = updateLHsType (hsib_body lhs )}
58−updateLHsSigType lhs@XHsImplicitBndrs{} = lhs
81+updateLHsSigType :: Traversal' (LHsSigType GhcPs) (LHsType GhcPs)
82+updateLHsSigType f lhs@HsIB {} =
83+ (\x -> lhs {hsib_body = x}) <$> f (hsib_body lhs)
84+updateLHsSigType _ lhs = pure lhs
5985
60−updateLHsType :: LHsType GhcPs -> LHsType GhcPs
61−updateLHsType = fmap updateHsType
86+updateLHsType :: Traversal' (LHsType GhcPs) (HsType GhcPs)
87+updateLHsType = traverse
6288
6389-- Main process
64−updateHsType :: HsType GhcPs -> HsType GhcPs
65−updateHsType ty@(HsQualTy xty ctxt body) = HsQualTy xty (fmap appendHSC ctxt) body
66−updateHsType ty@HsTyVar{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
67−updateHsType ty@HsAppTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
68−updateHsType ty@HsFunTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
69−updateHsType ty@HsListTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
70−updateHsType ty@HsTupleTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
71−updateHsType ty = ty
90+updateHsType :: HsType GhcPs -> (Any, HsType GhcPs)
91+updateHsType (HsQualTy xty ctxt body) =
92+ flagASTModified $ HsQualTy xty (fmap appendHSC ctxt) body
93+updateHsType ty@HsTyVar {} =
94+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
95+updateHsType ty@HsAppTy {} =
96+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
97+updateHsType ty@HsFunTy {} =
98+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
99+updateHsType ty@HsListTy {} =
100+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
101+updateHsType ty@HsTupleTy {} =
102+ flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
103+updateHsType ty = pure ty
104+
105+flagASTModified :: a -> (Any, a)
106+flagASTModified a = (Any True, a)
72107
73108appendHSC :: HsContext GhcPs -> HsContext GhcPs
74109appendHSC cs = mkHSC : cs
@@ -78,4 +113,4 @@ mkHSC :: LHsType GhcPs
78113mkHSC = noLoc $ HsTyVar noExt NotPromoted lId
79114
80115lId :: Located (IdP GhcPs)
81−lId = noLoc $ mkRdrUnqual $ mkClsOcc "HasCallStack"
116+lId = noLoc $ mkRdrQual ghcStackModuleName $ mkClsOcc "HasCallStack"