auto-split-0.1.0.4: src/AutoSplit.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE CPP #-}
module AutoSplit
  ( plugin
  ) where

import           Control.Exception
import           Data.Functor
import           Data.IORef
import           Data.String (IsString, fromString)
import qualified Data.Typeable as Typeable
import qualified GHC.Paths as Paths
import qualified Language.Haskell.GHC.ExactPrint.Parsers as EP
import qualified Language.Haskell.GHC.ExactPrint.Utils as EP

import qualified AutoSplit.Fields as Fields
import qualified AutoSplit.GhcFacade as Ghc
import qualified AutoSplit.Split as Split

plugin :: Ghc.Plugin
plugin = Ghc.defaultPlugin
  { Ghc.driverPlugin = \_ hscEnv -> pure $ addDsHook hscEnv
  , Ghc.parsedResultAction = \_ _ result -> pure $ addImport result
  , Ghc.pluginRecompile = Ghc.purePlugin
  , Ghc.typeCheckResultAction = \_ _ env ->
      removeUnusedImportWarn >> pure env
  }

-- | Incomplete patterns warning is emitted by the desugarer. Some shenanigans
-- are needed to intercept these warnings: a custom error is added to the
-- frontend result which causes the desugarer to throw a 'SourceError'
-- exception containing all diagnostics after running. This exception is then
-- caught (and later rethrown) by the plugin.
addDsHook :: Ghc.HscEnv -> Ghc.HscEnv
addDsHook hscEnv = hscEnv
  { Ghc.hsc_hooks =
      let hooks = Ghc.hsc_hooks hscEnv
       in hooks
          { Ghc.runPhaseHook = Just $ phaseHook (Ghc.runPhaseHook hooks) }
  }
  where
    phaseHook mExistingHook = Ghc.PhaseHook $ \tPhase -> do
      let isMissingPatWarn msgEnv =
            case Ghc.errMsgDiagnostic msgEnv of
              Ghc.GhcDsMessage (Ghc.DsNonExhaustivePatterns _ _ _ patIds nablas)
                | Just srcSpan <- Ghc.srcSpanToRealSrcSpan (Ghc.errMsgSpan msgEnv) ->
                    Left Split.NonExhaustivePatterns
                      { Split.patternIds = patIds
                      , Split.patternNablas = nablas
                      , Split.srcCodeLoc = srcSpan
                      }
              _ -> Right msgEnv
          isMissingFieldsErr msgEnv =
            case Ghc.errMsgDiagnostic msgEnv of
              Ghc.GhcTcRnMessage (Ghc.TcRnMessageWithInfo _ (Ghc.TcRnMessageDetailed _ (Ghc.TcRnMissingFields conLike _fields)))
                | Just srcSpan <- Ghc.srcSpanToRealSrcSpan (Ghc.errMsgSpan msgEnv) ->
                    Left Fields.MissingFields
                      { Fields.conLike = conLike
                      , Fields.srcCodeLoc = srcSpan
                      }
              Ghc.GhcTcRnMessage (Ghc.TcRnMessageWithInfo _ (Ghc.TcRnMessageDetailed _ (Ghc.TcRnMissingStrictFields conLike _fields)))
                | Just srcSpan <- Ghc.srcSpanToRealSrcSpan (Ghc.errMsgSpan msgEnv) ->
                    Left Fields.MissingFields
                      { Fields.conLike = conLike
                      , Fields.srcCodeLoc = srcSpan
                      }
              _ -> Right msgEnv
          isAutoSplitError msgEnv =
            case Ghc.errMsgDiagnostic msgEnv of
              Ghc.GhcUnknownMessage (Ghc.UnknownDiagnostic' a)
                | Just PatternSplitDiag <- Typeable.cast a -> True
              _ -> False
          parseFailedErr filePath =
            Ghc.mkPlainErrorMsgEnvelope
              (Ghc.mkGeneralSrcSpan $ fromString filePath)
              (Ghc.ghcUnknownMessage ParseFailed)
          patSplitErr filePath =
            Ghc.mkPlainErrorMsgEnvelope
              (Ghc.mkGeneralSrcSpan $ fromString filePath)
              (Ghc.ghcUnknownMessage PatternSplitDiag)
          improperUseErr filePath =
            Ghc.mkPlainErrorMsgEnvelope
              (Ghc.mkGeneralSrcSpan $ fromString filePath)
              (Ghc.ghcUnknownMessage ImproperUse)
          runPhaseOrExistingHook :: Ghc.TPhase res -> IO res
          runPhaseOrExistingHook = maybe Ghc.runPhase (\(Ghc.PhaseHook h) -> h) mExistingHook
      case tPhase of
        Ghc.T_HscPostTc env modSum tcResult@(Ghc.FrontendTypecheck gblEnv) warns mOldHash -> do
          usedGres <- readIORef $ Ghc.tcg_used_gres gblEnv
          let usesSplit = any ((== Split.splitName) . Ghc.occNameFS . Ghc.occName . Ghc.gre_name) usedGres
              mFilePath = Ghc.ml_hs_file (Ghc.ms_location modSum)
          case mFilePath of
            Just filePath | usesSplit -> do
              let noMissingPatErr =
                    Ghc.mkPlainErrorMsgEnvelope
                      (Ghc.mkGeneralSrcSpan $ fromString filePath)
                      (Ghc.ghcUnknownMessage NoMissingPat)
                  warnsWithError = Ghc.addMessage (patSplitErr filePath) warns
                  updEnv = env
                    -- Plugin only functions if incomplete patterns warning is enabled, so we force it on.
                    -- If this is instead done as driver plugin, ghci sessions won't pick it up.
                    { Ghc.hsc_dflags = (Ghc.hsc_dflags env `Ghc.wopt_set` Ghc.Opt_WarnIncompletePatterns)
                      -- Need to override the number of uncovered patterns reported.
                      { Ghc.maxUncoveredPatterns = maxBound - 1 }
                    }
              catch
                (runPhaseOrExistingHook $ Ghc.T_HscPostTc updEnv modSum tcResult warnsWithError mOldHash)
                (\(Ghc.SourceError msgs) -> do
                  let (missingPatWarns, otherDiags) = Ghc.partitionBagWith isMissingPatWarn (Ghc.getMessages msgs)
                      -- Need to parse the module because GHC removes
                      -- ParsedResult from ModSummary after the frontend finishes.
                      -- Use the options from compilation to parse the module, otherwise certain
                      -- syntax extensions won't parse correctly.
                      dynFlags = Ghc.ms_hspp_opts modSum `Ghc.gopt_set` Ghc.Opt_KeepRawTokenStream
                  eResult <- parseModule env dynFlags filePath
                  case eResult of
                    (Right parsedMod, usesCpp)
                      | not (null missingPatWarns) -> do
                          let gblRdrEnv = Ghc.tcg_rdr_env gblEnv
                          updated <- Split.modifyModule gblRdrEnv parsedMod usesCpp missingPatWarns filePath
                          if updated
                             then throw . Ghc.SourceError $ Ghc.mkMessages otherDiags
                             else throw . Ghc.SourceError . Ghc.mkMessages
                                   . Ghc.consBag (improperUseErr filePath)
                                   $ Ghc.filterBag (not . isAutoSplitError) otherDiags
                      | otherwise -> throw . Ghc.SourceError . Ghc.mkMessages
                                   . Ghc.consBag noMissingPatErr
                                   $ Ghc.filterBag (not . isAutoSplitError) otherDiags
                    (Left _, _) ->
                      throw . Ghc.SourceError . Ghc.mkMessages
                            . Ghc.consBag (parseFailedErr filePath)
                            $ Ghc.filterBag (not . isAutoSplitError) otherDiags
                )
            _ -> runPhaseOrExistingHook tPhase -- no SPLIT or no file path

        -- Missing fields
        Ghc.T_Hsc env modSum -> do
          gblEnvRef <- newIORef Nothing :: IO (IORef (Maybe Ghc.TcGblEnv))
          let rnPlugin =
                Ghc.defaultPlugin
                { Ghc.renamedResultAction = \_ gblEnv group -> do
                    Ghc.liftIO $ writeIORef gblEnvRef (Just gblEnv)
                    pure (gblEnv, group)
                }
              staticPlugin = Ghc.StaticPlugin
                { Ghc.spPlugin = Ghc.PluginWithArgs rnPlugin []
#if MIN_VERSION_ghc(9,12,0)
                , Ghc.spInitialised = True
#endif
                }
              envWithPlugin = env
                { Ghc.hsc_plugins = let ps = Ghc.hsc_plugins env in ps
                  { Ghc.staticPlugins = staticPlugin : Ghc.staticPlugins ps }
                }
          eTcRes <- try . runPhaseOrExistingHook $ Ghc.T_Hsc envWithPlugin modSum
          let msgs = case eTcRes of
                Right (_res, m) -> m
                Left (Ghc.SourceError m) -> m
              result = case eTcRes of
                Right res -> pure res
                Left err -> throw err
          mGblEnv <- readIORef gblEnvRef
          case mGblEnv of
            Nothing -> result -- renaming failed
            Just gblEnv -> do
              usedGres <- readIORef $ Ghc.tcg_used_gres gblEnv
              let usesFields = any ((== Fields.fieldsName) . Ghc.occNameFS . Ghc.occName . Ghc.gre_name) usedGres
                  mFilePath = Ghc.ml_hs_file (Ghc.ms_location modSum)
                  (missingFields, otherDiags) = Ghc.partitionBagWith isMissingFieldsErr (Ghc.getMessages msgs)
              case mFilePath of
                Just filePath | usesFields, not (null missingFields) -> do
                  -- Parse module
                  let dynFlags = Ghc.ms_hspp_opts modSum `Ghc.gopt_set` Ghc.Opt_KeepRawTokenStream
                  eParseRes <- parseModule env dynFlags filePath
                  case eParseRes of
                    (Right parsedMod, usesCpp) -> do
                      let gblRdrEnv = Ghc.tcg_rdr_env gblEnv
                      updated <- Fields.modifyModule gblRdrEnv parsedMod usesCpp missingFields filePath
                      if updated
                         then throw . Ghc.SourceError . Ghc.mkMessages
                                $ Ghc.consBag (patSplitErr filePath) otherDiags
                         else throw . Ghc.SourceError . Ghc.mkMessages
                                $ Ghc.consBag (improperUseErr filePath) otherDiags
                    (Left _, _) ->
                      throw . Ghc.SourceError . Ghc.mkMessages
                            $ Ghc.consBag (parseFailedErr filePath) otherDiags
                _ -> result -- no FIELDS or no file path

        _ -> runPhaseOrExistingHook tPhase

-- | Parse the given module file. Accounts for CPP comments
parseModule
  :: Ghc.HscEnv
  -> Ghc.DynFlags
  -> String
  -> IO (EP.ParseResult Ghc.ParsedSource, Bool)
parseModule env dynFlags filePath = EP.ghcWrapper Paths.libdir $ do
  Ghc.setSession env { Ghc.hsc_dflags = dynFlags }
  res <- EP.parseModuleEpAnnsWithCppInternal EP.defaultCppOptions dynFlags filePath
  let eCppComments = fmap (\(c, _, _) -> c) res
      hasCpp = case eCppComments of
                 Right cs -> not $ null cs
                 _ -> False
  pure
    ( liftA2 EP.insertCppComments
        (EP.postParseTransform res)
        eCppComments
    , hasCpp
    )

-- | Diagnostic thrown when case splitting should be attempted.
data PatternSplitDiag = PatternSplitDiag

instance Ghc.Diagnostic PatternSplitDiag where
  type DiagnosticOpts PatternSplitDiag = Ghc.NoDiagnosticOpts
  diagnosticMessage _ _ = Ghc.mkSimpleDecorated $
    Ghc.text "Module updated by auto-split, compilation aborted"
  diagnosticReason _ = Ghc.ErrorWithoutFlag
  diagnosticHints _ = []
  diagnosticCode _ = Nothing
#if !MIN_VERSION_ghc(9,8,0)
  defaultDiagnosticOpts = Ghc.NoDiagnosticOpts
#endif

-- | Diagnostic thrown when SPLIT is used but there are no resulting warnings
data NoMissingPat = NoMissingPat

instance Ghc.Diagnostic NoMissingPat where
  type DiagnosticOpts NoMissingPat = Ghc.NoDiagnosticOpts
  diagnosticMessage _ _ = Ghc.mkSimpleDecorated $
    Ghc.text "Module was not updated because all cases are already covered where SPLIT occurs"
  diagnosticReason _ = Ghc.ErrorWithoutFlag
  diagnosticHints _ = []
  diagnosticCode _ = Nothing
#if !MIN_VERSION_ghc(9,8,0)
  defaultDiagnosticOpts = Ghc.NoDiagnosticOpts
#endif

-- | Diagnostic thrown when SPLIT is used but there are no resulting warnings
data ImproperUse = ImproperUse

instance Ghc.Diagnostic ImproperUse where
  type DiagnosticOpts ImproperUse = Ghc.NoDiagnosticOpts
  diagnosticMessage _ _ = Ghc.mkSimpleDecorated $
    Ghc.text "It appears that an auto-split pattern was used improperly. The module has not been updated."
  diagnosticReason _ = Ghc.ErrorWithoutFlag
  diagnosticHints _ = []
  diagnosticCode _ = Nothing
#if !MIN_VERSION_ghc(9,8,0)
  defaultDiagnosticOpts = Ghc.NoDiagnosticOpts
#endif

-- | Diagnostic thrown if the module fails to parse
data ParseFailed = ParseFailed

instance Ghc.Diagnostic ParseFailed where
  type DiagnosticOpts ParseFailed = Ghc.NoDiagnosticOpts
  diagnosticMessage _ _ = Ghc.mkSimpleDecorated $
    Ghc.text "auto-split failed to parse the module"
  diagnosticReason _ = Ghc.ErrorWithoutFlag
  diagnosticHints _ = []
  diagnosticCode _ = Nothing
#if !MIN_VERSION_ghc(9,8,0)
  defaultDiagnosticOpts = Ghc.NoDiagnosticOpts
#endif

-- | Adds the import for SPLIT to the module being compiled. Otherwise users
-- would have to manually add this import everytime they want to do pattern splitting.
addImport :: Ghc.ParsedResult -> Ghc.ParsedResult
addImport result = result
  { Ghc.parsedResultModule =
    let resMod = Ghc.parsedResultModule result
     in resMod
        { Ghc.hpm_module = Ghc.hpm_module resMod <&> \pMod ->
            pMod
              { Ghc.hsmodImports =
                  patternImport : Ghc.hsmodImports pMod
              }
        }
  }
  where
    patternImport = Ghc.noLocA . Ghc.simpleImportDecl $ Ghc.mkModuleName patternModName

-- | The automatically added import gets flagged as unused even if it is used.
-- The solution here is to simply suppress the warning.
removeUnusedImportWarn :: Ghc.TcM ()
removeUnusedImportWarn = do
  errsVar <- Ghc.getErrsVar
#if MIN_VERSION_ghc(9,8,0)
  let isAutoSplitImportWarn msgEnv =
        case Ghc.errMsgDiagnostic msgEnv of
          Ghc.TcRnMessageWithInfo _ (Ghc.TcRnMessageDetailed _ (Ghc.TcRnUnusedImport decl _)) ->
            Ghc.unLoc (Ghc.ideclName decl) == Ghc.mkModuleName patternModName
          _ -> False
  Ghc.liftIO . modifyIORef errsVar $
    Ghc.mkMessages . Ghc.filterBag (not . isAutoSplitImportWarn) . Ghc.getMessages
#else
  -- 9.6 lacks the specific diagnostic
  let isAutoSplitImportWarn msgEnv =
        case Ghc.errMsgDiagnostic msgEnv of
          Ghc.TcRnMessageWithInfo _ (Ghc.TcRnMessageDetailed _ (Ghc.TcRnUnknownMessage (Ghc.UnknownDiagnostic diag)))
            | Ghc.WarningWithFlag Ghc.Opt_WarnUnusedImports <- Ghc.diagnosticReason diag
            -> True
          _ -> False
  Ghc.liftIO . modifyIORef errsVar $ \msgs ->
    -- the target import warning always shows up as the last occurrence
    case break isAutoSplitImportWarn . reverse . Ghc.bagToList $ Ghc.getMessages msgs of
      (before, _ : after) -> Ghc.mkMessages . Ghc.listToBag . reverse $ before ++ after
      _ -> msgs
#endif

patternModName :: IsString a => a
patternModName = "AutoSplit.Pattern"