TidyImportsPlugin.hs

{-# LANGUAGE TypeApplications #-}

module TidyImportsPlugin (plugin) where

import Data.Function ((&))
import Data.Generics (Data, cast, gmapQr)
import GHC.Hs
import GhcPlugins

plugin :: Plugin
plugin =
  defaultPlugin
    { pluginRecompile = const $ return NoForceRecompile,
      parsedResultAction = \_ _ parsed ->
        return $ parsed {hpm_module = modifyHsMod <$> hpm_module parsed}
    }

modifyHsMod :: HsModule GhcPs -> HsModule GhcPs
modifyHsMod m =
  let newImports =
        hsmodDecls m
          & map referencedModules
          & unionManyUniqSets
          & nonDetEltsUniqSet
          & map (\n -> noLoc (simpleImportDecl n) {ideclQualified = QualifiedPre})
   in m {hsmodImports = hsmodImports m ++ newImports}

referencedModules :: Data a => a -> UniqSet ModuleName
referencedModules =
  gmapQr
    unionUniqSets
    emptyUniqSet
    ( \d ->
        cast @_ @ModuleName d
          & maybe (referencedModules d) unitUniqSet
    )