{-# 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 Control.Monad.IO.Class (MonadIO)
import Data.IORef (readIORef)
import Data.List (intercalate)
import Data.Map (Map)
import Data.Set (Set)
import GHC.Plugins (DynFlags(dumpDir), GenModule(moduleName),
  GlobalRdrElt(GRE, gre_imp, gre_name, gre_par),
  HasDynFlags(getDynFlags), ImpDeclSpec(ImpDeclSpec, is_as,
  is_mod, is_qual), ImportSpec(is_decl), ImportedBy(ImportedByUser),
  ImportedModsVal(imv_all_exports), Outputable(ppr), Parent(FldParent,
  NoParent, ParentIs), Plugin(pluginRecompile, typeCheckResultAction),
  PluginRecompile(NoForceRecompile), CommandLineOption, ModSummary,
  ModuleName, Name, bestImport, defaultPlugin, liftIO, moduleEnvToList,
  moduleNameString, occEnvElts, showSDoc)
import GHC.Tc.Types (ImportAvails(imp_mods), TcGblEnv(tcg_imports,
  tcg_mod, tcg_used_gres), TcM)
import qualified Data.Map as Map
import qualified Data.Set as Set


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


typeCheckResultActionImpl
  :: [CommandLineOption]
  -> ModSummary
  -> TcGblEnv
  -> TcM TcGblEnv
typeCheckResultActionImpl :: [[Char]] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typeCheckResultActionImpl [[Char]]
_ ModSummary
_ TcGblEnv
env = do
  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 [Char])
-> IOEnv (Env TcGblEnv TcLclEnv) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IOEnv (Env TcGblEnv TcLclEnv) (Maybe [Char])
 -> IOEnv (Env TcGblEnv TcLclEnv) ())
-> IOEnv (Env TcGblEnv TcLclEnv) (Maybe [Char])
-> IOEnv (Env TcGblEnv TcLclEnv) ()
forall a b. (a -> b) -> a -> b
$ TcGblEnv
-> DynFlags
-> Map ModuleImport (Map Name (Set Name))
-> IOEnv (Env TcGblEnv TcLclEnv) (Maybe [Char])
forall (m :: * -> *).
MonadIO m =>
TcGblEnv
-> DynFlags
-> Map ModuleImport (Map Name (Set Name))
-> m (Maybe [Char])
writeToDumpFile TcGblEnv
env DynFlags
flags Map ModuleImport (Map Name (Set Name))
used
  TcGblEnv -> TcM TcGblEnv
forall (f :: * -> *) a. Applicative f => a -> f a
pure TcGblEnv
env


writeToDumpFile
  :: (MonadIO m)
  => TcGblEnv
  -> DynFlags
  -> Map ModuleImport (Map Name (Set Name))
  -> m (Maybe FilePath)
writeToDumpFile :: forall (m :: * -> *).
MonadIO m =>
TcGblEnv
-> DynFlags
-> Map ModuleImport (Map Name (Set Name))
-> m (Maybe [Char])
writeToDumpFile TcGblEnv
env DynFlags
flags Map ModuleImport (Map Name (Set Name))
used =
  {-
    If `-dumpdir` has been specified, then write the output into
    the dumpdir.  Mainly this  is because I can't figure out how to
    programmatically find the default dump dir.
  -}
  case DynFlags -> Maybe [Char]
dumpDir DynFlags
flags of
    Maybe [Char]
Nothing -> Maybe [Char] -> m (Maybe [Char])
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Char]
forall a. Maybe a
Nothing
    Just [Char]
dir ->
      IO (Maybe [Char]) -> m (Maybe [Char])
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe [Char]) -> m (Maybe [Char]))
-> IO (Maybe [Char]) -> m (Maybe [Char])
forall a b. (a -> b) -> a -> b
$ do
        let 
          modName :: FilePath
          modName :: [Char]
modName = ModuleName -> [Char]
moduleNameString (ModuleName -> [Char])
-> (TcGblEnv -> ModuleName) -> TcGblEnv -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenModule Unit -> ModuleName
forall unit. GenModule unit -> ModuleName
moduleName (GenModule Unit -> ModuleName)
-> (TcGblEnv -> GenModule Unit) -> TcGblEnv -> ModuleName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TcGblEnv -> GenModule Unit
tcg_mod (TcGblEnv -> [Char]) -> TcGblEnv -> [Char]
forall a b. (a -> b) -> a -> b
$ TcGblEnv
env

          filename :: FilePath
          filename :: [Char]
filename = [Char]
dir [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"/" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
modName [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
".full-imports"
        [Char] -> [Char] -> IO ()
writeFile [Char]
filename (DynFlags -> Map ModuleImport (Map Name (Set Name)) -> [Char]
renderNewImports DynFlags
flags Map ModuleImport (Map Name (Set Name))
used)
        Maybe [Char] -> IO (Maybe [Char])
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
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 (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 Name
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 -> Name
gre_name = Name
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]
occEnvElts (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 [ImportSpec]
imps

            modName :: ModuleName
            modImport :: ModuleImport
            (ModuleImport
modImport, ModuleName
modName) =
              let
                ImpDeclSpec { ModuleName
is_mod :: ModuleName
is_mod :: ImpDeclSpec -> ModuleName
is_mod , ModuleName
is_as :: ModuleName
is_as :: ImpDeclSpec -> ModuleName
is_as , Bool
is_qual :: Bool
is_qual :: ImpDeclSpec -> 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
                    FldParent Name
parentName Maybe FieldLabelString
_ -> Name -> Map Name (Set Name)
withPossibleParent Name
parentName
              )
        | GRE
            { gre_name :: GlobalRdrElt -> Name
gre_name = Name
name
            , gre_par :: GlobalRdrElt -> Parent
gre_par = Parent
parent
            , gre_imp :: GlobalRdrElt -> [ImportSpec]
gre_imp = [ImportSpec]
imps
            } <- [GlobalRdrElt]
rawUsed
        ]
  Map ModuleImport (Map Name (Set Name))
-> m (Map ModuleImport (Map Name (Set Name)))
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
/= :: ModuleImport -> ModuleImport -> Bool
$c/= :: ModuleImport -> ModuleImport -> Bool
== :: ModuleImport -> ModuleImport -> Bool
$c== :: 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
min :: ModuleImport -> ModuleImport -> ModuleImport
$cmin :: ModuleImport -> ModuleImport -> ModuleImport
max :: ModuleImport -> ModuleImport -> ModuleImport
$cmax :: ModuleImport -> ModuleImport -> ModuleImport
>= :: ModuleImport -> ModuleImport -> Bool
$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
compare :: ModuleImport -> ModuleImport -> Ordering
$ccompare :: ModuleImport -> ModuleImport -> Ordering
Ord)

renderNewImports
  :: DynFlags
  -> Map ModuleImport (Map Name (Set Name))
  -> String
renderNewImports :: DynFlags -> Map ModuleImport (Map Name (Set Name)) -> [Char]
renderNewImports DynFlags
flags Map ModuleImport (Map Name (Set Name))
used =
    [[Char]] -> [Char]
unlines
      [
        case ModuleImport
modImport of
          Unqualified ModuleName
modName ->
            [Char]
"import " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> ModuleName -> [Char]
forall o. Outputable o => o -> [Char]
shown ModuleName
modName [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" (" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Map Name (Set Name) -> [Char]
showParents Map Name (Set Name)
parents [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
")"
          Qualified ModuleName
modName ->
            [Char]
"import qualified " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> ModuleName -> [Char]
forall o. Outputable o => o -> [Char]
shown ModuleName
modName
          QualifiedAs ModuleName
modName ModuleName
asName ->
            [Char]
"import qualified " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> ModuleName -> [Char]
forall o. Outputable o => o -> [Char]
shown ModuleName
modName [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" as " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> ModuleName -> [Char]
forall o. Outputable o => o -> [Char]
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) -> [Char]
showParents Map Name (Set Name)
parents =
      [Char] -> [[Char]] -> [Char]
forall a. [a] -> [[a]] -> [a]
intercalate [Char]
", "
        [ Name -> [Char]
forall o. Outputable o => o -> [Char]
shown Name
parent [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Set Name -> [Char]
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 -> [Char]
showChildren Set Name
children =
      if Set Name -> Bool
forall a. Set a -> Bool
Set.null Set Name
children then
        [Char]
""
      else
        [Char]
"(" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char] -> [[Char]] -> [Char]
forall a. [a] -> [[a]] -> [a]
intercalate [Char]
", " (Name -> [Char]
forall o. Outputable o => o -> [Char]
shown (Name -> [Char]) -> [Name] -> [[Char]]
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) [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
")"

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