Ergonomic improvements

- Avoiding recompile (was always compiling twice)
 - Only importing `GHC.Stack` when we've modified the AST
 - Qualified import of `GHC.Stack` as `AutoImported.GHC.Stack`,
   to avoid any name shadowing
 - Cleaned up unused imports and binds
 - More liberal version constraints on ghc in cabal file

haskell-stack-trace-plugin.cabal

@@ -1,5 +1,5 @@
 name:                haskell-stack-trace-plugin
-version:             0.1.0.0
+version:             0.1.0.1
 synopsis:            haskell-stack-trace-plugin
 description:         This plugin allow implicitly add HasCallStack class to every top-level function for all module. Hence, we can to get completely continuous call stack.
 homepage:            https://github.com/waddlaw/haskell-stack-trace-plugin
@@ -23,7 +23,7 @@ library
   hs-source-dirs: src
   build-depends:
     base >=4.12 && <4.13,
-    ghc==8.6.2
+    ghc>=8.6.2 && <8.6.6
 
   exposed-modules:
     StackTrace.Plugin
@@ -38,4 +38,4 @@ executable example
     -fplugin=StackTrace.Plugin
   build-depends:
     base >=4.12 && <4.13,
-    haskell-stack-trace-plugin
+    haskell-stack-trace-plugin

src/StackTrace/Plugin.hs

@@ -1,74 +1,109 @@
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE OverloadedStrings #-}
 module StackTrace.Plugin (plugin) where
-  
+
+import Control.Arrow (first)
+import Data.Monoid (Any(Any, getAny))
 import GhcPlugins
-import TcRnTypes (IfM, TcM, TcGblEnv, tcg_binds, tcg_rn_decls)
-import HsExtension (GhcTc, GhcRn)
-import HsDecls (HsGroup)
-import HsExpr (LHsExpr)
 import HsSyn
 
-import Data.Maybe
+type Traversal s t a b
+   = forall f. Applicative f =>
+                 (a -> f b) -> s -> f t
+
+type Traversal' s a = Traversal s s a a
 
 plugin :: Plugin
-plugin = defaultPlugin
+plugin = defaultPlugin {parsedResultAction = parsedPlugin, pluginRecompile = purePlugin}
-  { parsedResultAction = parsedPlugin
-  }
 
-parsedPlugin :: [CommandLineOption] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
+parsedPlugin ::
+     [CommandLineOption] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
 parsedPlugin _ _ pm = do
-  dflags <- getDynFlags
+  let m = updateHsModule <$> hpm_module pm
-
+      pm' = pm {hpm_module = m}
-  let extract = hsmodImports . unLoc
-      addGHCStackMod = noLoc $ simpleImportDecl $ mkModuleName "GHC.Stack"
-
-      m = updateHsModule addGHCStackMod updateHsmodDecl <$> hpm_module pm
-      pm' = pm { hpm_module = m }
-
   return pm'
 
-updateHsModule :: LImportDecl GhcPs -> (LHsDecl GhcPs -> LHsDecl GhcPs) -> HsModule GhcPs -> HsModule GhcPs
+-- Use qualified import for GHC.Stack as "AutoImported.GHC.Stack"
-updateHsModule importDecl update hsm = hsm
+-- ...this should not interfere with other imports...
-    { hsmodImports = importDecl:decls
+ghcStackModuleName :: ModuleName
-    , hsmodDecls = map update lhss
+ghcStackModuleName = mkModuleName "AutoImported.GHC.Stack"
-    }
+
+ghcStackImport :: Located (ImportDecl (GhcPass p))
+ghcStackImport =
+  noLoc $
+  (simpleImportDecl $ mkModuleName "GHC.Stack")
+    {ideclQualified = True, ideclAs = Just $ noLoc ghcStackModuleName}
+
+updateHsModule :: HsModule GhcPs -> HsModule GhcPs
+updateHsModule hsm =
+  hsm {hsmodImports = hsmodImports', hsmodDecls = hsmodDecls'}
   where
-    decls = hsmodImports hsm
+    -- Traverse the haskell AST; if we have to add some HasStack
-    lhss = hsmodDecls hsm
+    -- constraint we set a flag in a (Any,) functor.
+    -- ...it'd be simpler to check if before == after, but Haskell AST
+    -- doesn't have Eq instances.
+    (updatedP, hsmodDecls') =
+      first getAny $
+      (traverse . astTraversal) updateHsType (hsmodDecls hsm)
+
+    -- Only import GHC.Stack if needed for a constraint we introduced
+    hsmodImports' =
+      (if updatedP
+         then [ghcStackImport]
+         else []) ++
+      hsmodImports hsm
+
+    astTraversal :: Traversal' (LHsDecl GhcPs) (HsType GhcPs)
+    astTraversal = updateHsmodDecl
+                 . updateHsDecl
+                 . updateSig
+                 . updateLHsSigWsType
+                 . updateLHsSigType
+                 . updateLHsType
 
 --------------
+updateHsmodDecl :: Traversal' (LHsDecl GhcPs) (HsDecl GhcPs)
+updateHsmodDecl = traverse
 
-updateHsmodDecl :: LHsDecl GhcPs -> LHsDecl GhcPs
+updateHsDecl :: Traversal' (HsDecl GhcPs) (Sig GhcPs)
-updateHsmodDecl = fmap updateHsDecl
+updateHsDecl f (SigD xSig s) = SigD xSig <$> f s
+updateHsDecl _ sig = pure sig
 
-updateHsDecl :: HsDecl GhcPs -> HsDecl GhcPs
+updateSig :: Traversal' (Sig GhcPs) (LHsSigWcType GhcPs)
-updateHsDecl (SigD xSig s) = SigD xSig (updateSig s)
+updateSig f (TypeSig xSig ls t) = TypeSig xSig ls <$> f t
-updateHsDecl decl = decl
+updateSig _ sig = pure sig
 
-updateSig :: Sig GhcPs -> Sig GhcPs
+updateLHsSigWsType :: Traversal' (LHsSigWcType GhcPs) (LHsSigType GhcPs)
-updateSig (TypeSig xSig ls t) = TypeSig xSig ls (updateLHsSigWcType t)
+updateLHsSigWsType f lhs@HsWC {} =
-updateSig sig = sig
+  (\x -> lhs {hswc_body = x}) <$> f (hswc_body lhs)
+updateLHsSigWsType _ lhs = pure lhs
 
-updateLHsSigWcType :: LHsSigWcType GhcPs -> LHsSigWcType GhcPs
+updateLHsSigType :: Traversal' (LHsSigType GhcPs) (LHsType GhcPs)
-updateLHsSigWcType lhs@HsWC{} = lhs { hswc_body = updateLHsSigType (hswc_body lhs) }
+updateLHsSigType f lhs@HsIB {} =
-updateLHsSigWcType lhs@XHsWildCardBndrs{} = lhs
+  (\x -> lhs {hsib_body = x}) <$> f (hsib_body lhs)
+updateLHsSigType _ lhs = pure lhs
 
-updateLHsSigType :: LHsSigType GhcPs -> LHsSigType GhcPs
+updateLHsType :: Traversal' (LHsType GhcPs) (HsType GhcPs)
-updateLHsSigType lhs@HsIB{} = lhs { hsib_body = updateLHsType (hsib_body lhs )}
+updateLHsType = traverse
-updateLHsSigType lhs@XHsImplicitBndrs{} = lhs
-
-updateLHsType :: LHsType GhcPs -> LHsType GhcPs
-updateLHsType = fmap updateHsType
 
 -- Main process
-updateHsType :: HsType GhcPs -> HsType GhcPs
+updateHsType :: HsType GhcPs -> (Any, HsType GhcPs)
-updateHsType ty@(HsQualTy xty ctxt body) = HsQualTy xty (fmap appendHSC ctxt) body
+updateHsType (HsQualTy xty ctxt body) =
-updateHsType ty@HsTyVar{}   = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+  flagASTModified $ HsQualTy xty (fmap appendHSC ctxt) body
-updateHsType ty@HsAppTy{}   = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+updateHsType ty@HsTyVar {} =
-updateHsType ty@HsFunTy{}   = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+  flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
-updateHsType ty@HsListTy{}  = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+updateHsType ty@HsAppTy {} =
-updateHsType ty@HsTupleTy{} = HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+  flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
-updateHsType ty = ty
+updateHsType ty@HsFunTy {} =
+  flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+updateHsType ty@HsListTy {} =
+  flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+updateHsType ty@HsTupleTy {} =
+  flagASTModified $ HsQualTy noExt (noLoc $ appendHSC []) (noLoc ty)
+updateHsType ty = pure ty
+
+flagASTModified :: a -> (Any, a)
+flagASTModified a = (Any True, a)
 
 appendHSC :: HsContext GhcPs -> HsContext GhcPs
 appendHSC cs = mkHSC : cs
@@ -78,4 +113,4 @@ mkHSC :: LHsType GhcPs
 mkHSC = noLoc $ HsTyVar noExt NotPromoted lId
 
 lId :: Located (IdP GhcPs)
-lId = noLoc $ mkRdrUnqual $ mkClsOcc "HasCallStack"
+lId = noLoc $ mkRdrQual ghcStackModuleName $ mkClsOcc "HasCallStack"