{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | Description: Plugin for helping close open imports. -}
module OM.Plugin.Imports (
  plugin,
) where


import Control.Monad (void)
import Data.IORef (readIORef)
import Data.List (intercalate)
import Data.Map (Map)
import Data.Set (Set)
import GHC (ModSummary(ms_hspp_file), DynFlags, ModuleName, Name,
  moduleName)
import GHC.Data.Bag (bagToList)
import GHC.Plugins (GlobalRdrElt(GRE, gre_imp, gre_name, gre_par),
  HasDynFlags(getDynFlags), ImpDeclSpec(ImpDeclSpec, is_as, is_mod,
  is_qual), ImportSpec(is_decl), Outputable(ppr), Parent(NoParent,
  ParentIs), Plugin(pluginRecompile, typeCheckResultAction),
  PluginRecompile(NoForceRecompile), CommandLineOption, bestImport,
  defaultPlugin, liftIO, moduleEnvToList, nonDetOccEnvElts, showSDoc)
import GHC.Tc.Utils.Monad (ImportAvails(imp_mods), TcGblEnv(tcg_imports,
  tcg_used_gres), MonadIO, TcM)
import GHC.Types.Avail (greNamePrintableName)
import GHC.Unit.Module.Imported (ImportedBy(ImportedByUser),
  ImportedModsVal(imv_all_exports))
import Prelude (Applicative(pure), Bool(False, True), Eq((==)),
  Maybe(Just, Nothing), Monoid(mempty), Semigroup((<>)), ($), (.),
  (<$>), (||), FilePath, Ord, String, concat, otherwise, putStrLn,
  unlines, writeFile)
import Safe (headMay)
import qualified Data.Char as Char
import qualified Data.Map as Map
import qualified Data.Set as Set


plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin
  { typeCheckResultAction = typeCheckResultActionImpl
  , pluginRecompile = \[String]
_ -> PluginRecompile -> IO PluginRecompile
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PluginRecompile
NoForceRecompile
  }


typeCheckResultActionImpl
  :: [CommandLineOption]
  -> ModSummary
  -> TcGblEnv
  -> TcM TcGblEnv
typeCheckResultActionImpl :: [String] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typeCheckResultActionImpl [String]
_ ModSummary
modSummary TcGblEnv
env = do
  IO () -> IOEnv (Env TcGblEnv TcLclEnv) ()
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO ()
putStrLn (ModSummary -> String
ms_hspp_file ModSummary
modSummary))
  Map ModuleImport (Map Name (Set Name))
used <- TcGblEnv
-> IOEnv
     (Env TcGblEnv TcLclEnv) (Map ModuleImport (Map Name (Set Name)))
forall (m :: * -> *).
MonadIO m =>
TcGblEnv -> m (Map ModuleImport (Map Name (Set Name)))
getUsedImports TcGblEnv
env
  DynFlags
flags <- IOEnv (Env TcGblEnv TcLclEnv) DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
  IOEnv (Env TcGblEnv TcLclEnv) (Maybe String)
-> IOEnv (Env TcGblEnv TcLclEnv) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IOEnv (Env TcGblEnv TcLclEnv) (Maybe String)
 -> IOEnv (Env TcGblEnv TcLclEnv) ())
-> IOEnv (Env TcGblEnv TcLclEnv) (Maybe String)
-> IOEnv (Env TcGblEnv TcLclEnv) ()
forall a b. (a -> b) -> a -> b
$ String
-> DynFlags
-> Map ModuleImport (Map Name (Set Name))
-> IOEnv (Env TcGblEnv TcLclEnv) (Maybe String)
forall (m :: * -> *).
MonadIO m =>
String
-> DynFlags
-> Map ModuleImport (Map Name (Set Name))
-> m (Maybe String)
writeToDumpFile (ModSummary -> String
ms_hspp_file ModSummary
modSummary) DynFlags
flags Map ModuleImport (Map Name (Set Name))
used
  TcGblEnv -> TcM TcGblEnv
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TcGblEnv
env


writeToDumpFile
  :: (MonadIO m)
  => FilePath
  -> DynFlags
  -> Map ModuleImport (Map Name (Set Name))
  -> m (Maybe FilePath)
writeToDumpFile :: forall (m :: * -> *).
MonadIO m =>
String
-> DynFlags
-> Map ModuleImport (Map Name (Set Name))
-> m (Maybe String)
writeToDumpFile String
srcFile DynFlags
flags Map ModuleImport (Map Name (Set Name))
used =
  IO (Maybe String) -> m (Maybe String)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe String) -> m (Maybe String))
-> IO (Maybe String) -> m (Maybe String)
forall a b. (a -> b) -> a -> b
$ do
    let
      filename :: FilePath
      filename :: String
filename = String
srcFile String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
".full-imports"
    String -> String -> IO ()
writeFile String
filename (DynFlags -> Map ModuleImport (Map Name (Set Name)) -> String
renderNewImports DynFlags
flags Map ModuleImport (Map Name (Set Name))
used)
    Maybe String -> IO (Maybe String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Maybe String
forall a. a -> Maybe a
Just String
filename)


getUsedImports
  :: forall m.
     (MonadIO m)
  => TcGblEnv
  -> m (Map ModuleImport (Map Name (Set Name)))
getUsedImports :: forall (m :: * -> *).
MonadIO m =>
TcGblEnv -> m (Map ModuleImport (Map Name (Set Name)))
getUsedImports TcGblEnv
env = do
  [GlobalRdrElt]
rawUsed <- (IO [GlobalRdrElt] -> m [GlobalRdrElt]
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [GlobalRdrElt] -> m [GlobalRdrElt])
-> (IORef [GlobalRdrElt] -> IO [GlobalRdrElt])
-> IORef [GlobalRdrElt]
-> m [GlobalRdrElt]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef [GlobalRdrElt] -> IO [GlobalRdrElt]
forall a. IORef a -> IO a
readIORef) (TcGblEnv -> IORef [GlobalRdrElt]
tcg_used_gres TcGblEnv
env) :: m [GlobalRdrElt]
  let
    {-
      Sometimes, the module from which the name is imported may not
      export the Parent of the name. E.g. Data.List exports 'foldl',
      but not 'Foldable'. So we check to see if the parent is available
      from the module. If it isn't then we just omit the parent. If it
      is, we include the parent with the justification that it provides
      more explicit information to the reader.
    -}
    availableParents :: Map ModuleName (Set Name)
    availableParents :: Map ModuleName (Set Name)
availableParents =
      (Set Name -> Set Name -> Set Name)
-> [Map ModuleName (Set Name)] -> Map ModuleName (Set Name)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith
        Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union
        [ ModuleName -> Set Name -> Map ModuleName (Set Name)
forall k a. k -> a -> Map k a
Map.singleton
            (GenModule Unit -> ModuleName
forall unit. GenModule unit -> ModuleName
moduleName GenModule Unit
m)
            (Name -> Set Name
forall a. a -> Set a
Set.singleton (GreName -> Name
greNamePrintableName GreName
name))
        | (GenModule Unit
m, [ImportedBy]
ibs)
            <- ModuleEnv [ImportedBy] -> [(GenModule Unit, [ImportedBy])]
forall a. ModuleEnv a -> [(GenModule Unit, a)]
moduleEnvToList (ModuleEnv [ImportedBy] -> [(GenModule Unit, [ImportedBy])])
-> (TcGblEnv -> ModuleEnv [ImportedBy])
-> TcGblEnv
-> [(GenModule Unit, [ImportedBy])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ImportAvails -> ModuleEnv [ImportedBy]
imp_mods (ImportAvails -> ModuleEnv [ImportedBy])
-> (TcGblEnv -> ImportAvails) -> TcGblEnv -> ModuleEnv [ImportedBy]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TcGblEnv -> ImportAvails
tcg_imports (TcGblEnv -> [(GenModule Unit, [ImportedBy])])
-> TcGblEnv -> [(GenModule Unit, [ImportedBy])]
forall a b. (a -> b) -> a -> b
$ TcGblEnv
env
        , ImportedByUser ImportedModsVal
imv <- [ImportedBy]
ibs
        , GRE { gre_name :: GlobalRdrElt -> GreName
gre_name = GreName
name } <- [[GlobalRdrElt]] -> [GlobalRdrElt]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[GlobalRdrElt]] -> [GlobalRdrElt])
-> (ImportedModsVal -> [[GlobalRdrElt]])
-> ImportedModsVal
-> [GlobalRdrElt]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccEnv [GlobalRdrElt] -> [[GlobalRdrElt]]
forall a. OccEnv a -> [a]
nonDetOccEnvElts (OccEnv [GlobalRdrElt] -> [[GlobalRdrElt]])
-> (ImportedModsVal -> OccEnv [GlobalRdrElt])
-> ImportedModsVal
-> [[GlobalRdrElt]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ImportedModsVal -> OccEnv [GlobalRdrElt]
imv_all_exports (ImportedModsVal -> [GlobalRdrElt])
-> ImportedModsVal -> [GlobalRdrElt]
forall a b. (a -> b) -> a -> b
$ ImportedModsVal
imv
        ]

    used :: Map ModuleImport (Map Name (Set Name))
    used :: Map ModuleImport (Map Name (Set Name))
used =
      (Map Name (Set Name) -> Map Name (Set Name) -> Map Name (Set Name))
-> [Map ModuleImport (Map Name (Set Name))]
-> Map ModuleImport (Map Name (Set Name))
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith
        ((Set Name -> Set Name -> Set Name)
-> Map Name (Set Name)
-> Map Name (Set Name)
-> Map Name (Set Name)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union)
        [ let
            imp :: ImportSpec
            imp :: ImportSpec
imp = [ImportSpec] -> ImportSpec
bestImport (Bag ImportSpec -> [ImportSpec]
forall a. Bag a -> [a]
bagToList Bag ImportSpec
imps)

            modName :: ModuleName
            modImport :: ModuleImport
            (ModuleImport
modImport, ModuleName
modName) =
              let
                ImpDeclSpec { ModuleName
is_mod :: ImpDeclSpec -> ModuleName
is_mod :: ModuleName
is_mod , ModuleName
is_as :: ImpDeclSpec -> ModuleName
is_as :: ModuleName
is_as , Bool
is_qual :: ImpDeclSpec -> Bool
is_qual :: Bool
is_qual } = ImportSpec -> ImpDeclSpec
is_decl ImportSpec
imp
              in
                ( case (Bool
is_qual, ModuleName
is_as ModuleName -> ModuleName -> Bool
forall a. Eq a => a -> a -> Bool
== ModuleName
is_mod) of
                    (Bool
True, Bool
True) -> ModuleName -> ModuleImport
Qualified ModuleName
is_mod
                    (Bool
True, Bool
False) -> ModuleName -> ModuleName -> ModuleImport
QualifiedAs ModuleName
is_mod ModuleName
is_as
                    (Bool
False, Bool
_) -> ModuleName -> ModuleImport
Unqualified ModuleName
is_mod
                , ModuleName
is_mod
                )
          in
            ModuleImport
-> Map Name (Set Name) -> Map ModuleImport (Map Name (Set Name))
forall k a. k -> a -> Map k a
Map.singleton
              ModuleImport
modImport
              (
                let
                  {-
                    Figure out if we need to omit the parent name because
                    it isn't exported by the module from which the name
                    itself is imported.
                  -}
                  withPossibleParent :: Name -> Map Name (Set Name)
                  withPossibleParent :: Name -> Map Name (Set Name)
withPossibleParent Name
parentName =
                    if
                      Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Name
parentName (Set Name -> Bool) -> Set Name -> Bool
forall a b. (a -> b) -> a -> b
$
                        Set Name -> ModuleName -> Map ModuleName (Set Name) -> Set Name
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault
                          Set Name
forall a. Monoid a => a
mempty
                          ModuleName
modName
                          Map ModuleName (Set Name)
availableParents
                    then
                      Name -> Set Name -> Map Name (Set Name)
forall k a. k -> a -> Map k a
Map.singleton Name
parentName (Name -> Set Name
forall a. a -> Set a
Set.singleton Name
name)
                    else
                      Map Name (Set Name)
noParent

                  noParent :: Map Name (Set Name)
                  noParent :: Map Name (Set Name)
noParent = Name -> Set Name -> Map Name (Set Name)
forall k a. k -> a -> Map k a
Map.singleton Name
name Set Name
forall a. Monoid a => a
mempty
                in
                  case Parent
parent of
                    Parent
NoParent -> Map Name (Set Name)
noParent
                    ParentIs Name
parentName ->
                      Name -> Map Name (Set Name)
withPossibleParent Name
parentName
              )
        | GRE
            { GreName
gre_name :: GlobalRdrElt -> GreName
gre_name :: GreName
gre_name
            , gre_par :: GlobalRdrElt -> Parent
gre_par = Parent
parent
            , gre_imp :: GlobalRdrElt -> Bag ImportSpec
gre_imp = Bag ImportSpec
imps
            } <- [GlobalRdrElt]
rawUsed
        , let
            name :: Name
            name :: Name
name = GreName -> Name
greNamePrintableName GreName
gre_name
        ]
  Map ModuleImport (Map Name (Set Name))
-> m (Map ModuleImport (Map Name (Set Name)))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map ModuleImport (Map Name (Set Name))
used


data ModuleImport
  = Unqualified ModuleName
  | Qualified ModuleName
  | QualifiedAs ModuleName ModuleName
  deriving stock (ModuleImport -> ModuleImport -> Bool
(ModuleImport -> ModuleImport -> Bool)
-> (ModuleImport -> ModuleImport -> Bool) -> Eq ModuleImport
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ModuleImport -> ModuleImport -> Bool
== :: ModuleImport -> ModuleImport -> Bool
$c/= :: ModuleImport -> ModuleImport -> Bool
/= :: ModuleImport -> ModuleImport -> Bool
Eq, Eq ModuleImport
Eq ModuleImport =>
(ModuleImport -> ModuleImport -> Ordering)
-> (ModuleImport -> ModuleImport -> Bool)
-> (ModuleImport -> ModuleImport -> Bool)
-> (ModuleImport -> ModuleImport -> Bool)
-> (ModuleImport -> ModuleImport -> Bool)
-> (ModuleImport -> ModuleImport -> ModuleImport)
-> (ModuleImport -> ModuleImport -> ModuleImport)
-> Ord ModuleImport
ModuleImport -> ModuleImport -> Bool
ModuleImport -> ModuleImport -> Ordering
ModuleImport -> ModuleImport -> ModuleImport
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: ModuleImport -> ModuleImport -> Ordering
compare :: ModuleImport -> ModuleImport -> Ordering
$c< :: ModuleImport -> ModuleImport -> Bool
< :: ModuleImport -> ModuleImport -> Bool
$c<= :: ModuleImport -> ModuleImport -> Bool
<= :: ModuleImport -> ModuleImport -> Bool
$c> :: ModuleImport -> ModuleImport -> Bool
> :: ModuleImport -> ModuleImport -> Bool
$c>= :: ModuleImport -> ModuleImport -> Bool
>= :: ModuleImport -> ModuleImport -> Bool
$cmax :: ModuleImport -> ModuleImport -> ModuleImport
max :: ModuleImport -> ModuleImport -> ModuleImport
$cmin :: ModuleImport -> ModuleImport -> ModuleImport
min :: ModuleImport -> ModuleImport -> ModuleImport
Ord)

renderNewImports
  :: DynFlags
  -> Map ModuleImport (Map Name (Set Name))
  -> String
renderNewImports :: DynFlags -> Map ModuleImport (Map Name (Set Name)) -> String
renderNewImports DynFlags
flags Map ModuleImport (Map Name (Set Name))
used =
    [String] -> String
unlines
      [
        case ModuleImport
modImport of
          Unqualified ModuleName
modName ->
            String
"import " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ModuleName -> String
forall o. Outputable o => o -> String
shown ModuleName
modName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" (" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Map Name (Set Name) -> String
showParents Map Name (Set Name)
parents String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
")"
          Qualified ModuleName
modName ->
            String
"import qualified " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ModuleName -> String
forall o. Outputable o => o -> String
shown ModuleName
modName
          QualifiedAs ModuleName
modName ModuleName
asName ->
            String
"import qualified " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ModuleName -> String
forall o. Outputable o => o -> String
shown ModuleName
modName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" as " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ModuleName -> String
forall o. Outputable o => o -> String
shown ModuleName
asName
      | (ModuleImport
modImport, Map Name (Set Name)
parents) <- Map ModuleImport (Map Name (Set Name))
-> [(ModuleImport, Map Name (Set Name))]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map ModuleImport (Map Name (Set Name))
used
      ]
  where
    showParents :: Map Name (Set Name) -> String
    showParents :: Map Name (Set Name) -> String
showParents Map Name (Set Name)
parents =
      String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", "
        [ Name -> String
forall o. Outputable o => o -> String
shown Name
parent String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Set Name -> String
showChildren Set Name
children
        | (Name
parent, Set Name
children) <- Map Name (Set Name) -> [(Name, Set Name)]
forall k a. Map k a -> [(k, a)]
Map.toList Map Name (Set Name)
parents
        ]

    showChildren :: Set Name -> String
    showChildren :: Set Name -> String
showChildren Set Name
children =
      if Set Name -> Bool
forall a. Set a -> Bool
Set.null Set Name
children then
        String
""
      else
        String
"(" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " (Name -> String
forall o. Outputable o => o -> String
shown (Name -> String) -> [Name] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Name -> [Name]
forall a. Set a -> [a]
Set.toAscList Set Name
children) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
")"

    shown :: Outputable o => o -> String
    shown :: forall o. Outputable o => o -> String
shown = String -> String
fixInlineName (String -> String) -> (o -> String) -> o -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DynFlags -> SDoc -> String
showSDoc DynFlags
flags (SDoc -> String) -> (o -> SDoc) -> o -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. o -> SDoc
forall a. Outputable a => a -> SDoc
ppr

    fixInlineName :: String -> String
    fixInlineName :: String -> String
fixInlineName String
name =
      case String -> Maybe Char
forall a. [a] -> Maybe a
headMay String
name of
        Maybe Char
Nothing -> String
name
        Just Char
c
          | Char -> Bool
Char.isAlphaNum Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'_' -> String
name
          | Bool
otherwise -> String
"(" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
name String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
")"