{-# 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 (DynFlags(dumpDir), ModSummary, ModuleName, Name, moduleName,
  moduleNameString)
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, occEnvElts, showSDoc)
import GHC.Tc.Utils.Monad (ImportAvails(imp_mods), TcGblEnv(tcg_imports,
  tcg_mod, tcg_used_gres), MonadIO, TcM)
import GHC.Types.Avail (greNamePrintableName)
import GHC.Unit.Module.Imported (ImportedBy(ImportedByUser),
  ImportedModsVal(imv_all_exports))
import qualified Data.Map as Map
import qualified Data.Set as Set


plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin
  { typeCheckResultAction :: [String] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typeCheckResultAction = [String] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typeCheckResultActionImpl
  , pluginRecompile :: [String] -> IO PluginRecompile
pluginRecompile = \[String]
_ -> 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
_ TcGblEnv
env = do
  Map ModuleImport (Map Name (Set Name))
used <- forall (m :: * -> *).
MonadIO m =>
TcGblEnv -> m (Map ModuleImport (Map Name (Set Name)))
getUsedImports TcGblEnv
env
  DynFlags
flags <- forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
  forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadIO m =>
TcGblEnv
-> DynFlags
-> Map ModuleImport (Map Name (Set Name))
-> m (Maybe String)
writeToDumpFile TcGblEnv
env DynFlags
flags Map ModuleImport (Map Name (Set Name))
used
  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 String)
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 String
dumpDir DynFlags
flags of
    Maybe String
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    Just String
dir ->
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let 
          modName :: FilePath
          modName :: String
modName = ModuleName -> String
moduleNameString forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall unit. GenModule unit -> ModuleName
moduleName forall b c a. (b -> c) -> (a -> b) -> a -> c
. TcGblEnv -> Module
tcg_mod forall a b. (a -> b) -> a -> b
$ TcGblEnv
env

          filename :: FilePath
          filename :: String
filename = String
dir forall a. Semigroup a => a -> a -> a
<> String
"/" forall a. Semigroup a => a -> a -> a
<> String
modName 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)
        forall (f :: * -> *) a. Applicative f => a -> f a
pure (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 <- (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IORef a -> IO a
readIORef) (TcGblEnv -> TcRef [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 =
      forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith
        forall a. Ord a => Set a -> Set a -> Set a
Set.union
        [ forall k a. k -> a -> Map k a
Map.singleton
            (forall unit. GenModule unit -> ModuleName
moduleName Module
m)
            (forall a. a -> Set a
Set.singleton (GreName -> Name
greNamePrintableName GreName
name))
        | (Module
m, [ImportedBy]
ibs)
            <- forall a. ModuleEnv a -> [(Module, a)]
moduleEnvToList forall b c a. (b -> c) -> (a -> b) -> a -> c
. ImportAvails -> ModuleEnv [ImportedBy]
imp_mods forall b c a. (b -> c) -> (a -> b) -> a -> c
. TcGblEnv -> ImportAvails
tcg_imports forall a b. (a -> b) -> a -> b
$ TcGblEnv
env
        , ImportedByUser ImportedModsVal
imv <- [ImportedBy]
ibs
        , GRE { gre_name :: GlobalRdrElt -> GreName
gre_name = GreName
name } <- forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. OccEnv a -> [a]
occEnvElts forall b c a. (b -> c) -> (a -> b) -> a -> c
. ImportedModsVal -> GlobalRdrEnv
imv_all_exports forall a b. (a -> b) -> a -> b
$ ImportedModsVal
imv
        ]

    used :: Map ModuleImport (Map Name (Set Name))
    used :: Map ModuleImport (Map Name (Set Name))
used =
      forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith
        (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith 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 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
            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
                      forall a. Ord a => a -> Set a -> Bool
Set.member Name
parentName forall a b. (a -> b) -> a -> b
$
                        forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault
                          forall a. Monoid a => a
mempty
                          ModuleName
modName
                          Map ModuleName (Set Name)
availableParents
                    then
                      forall k a. k -> a -> Map k a
Map.singleton Name
parentName (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 = forall k a. k -> a -> Map k a
Map.singleton Name
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 :: GreName
gre_name :: GlobalRdrElt -> GreName
gre_name
            , gre_par :: GlobalRdrElt -> Parent
gre_par = Parent
parent
            , gre_imp :: GlobalRdrElt -> [ImportSpec]
gre_imp = [ImportSpec]
imps
            } <- [GlobalRdrElt]
rawUsed
        , let
            name :: Name
            name :: Name
name = GreName -> Name
greNamePrintableName GreName
gre_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
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
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)) -> String
renderNewImports DynFlags
flags Map ModuleImport (Map Name (Set Name))
used =
    [String] -> String
unlines
      [
        case ModuleImport
modImport of
          Unqualified ModuleName
modName ->
            String
"import " forall a. Semigroup a => a -> a -> a
<> forall o. Outputable o => o -> String
shown ModuleName
modName forall a. Semigroup a => a -> a -> a
<> String
" (" forall a. Semigroup a => a -> a -> a
<> Map Name (Set Name) -> String
showParents Map Name (Set Name)
parents forall a. Semigroup a => a -> a -> a
<> String
")"
          Qualified ModuleName
modName ->
            String
"import qualified " forall a. Semigroup a => a -> a -> a
<> forall o. Outputable o => o -> String
shown ModuleName
modName
          QualifiedAs ModuleName
modName ModuleName
asName ->
            String
"import qualified " forall a. Semigroup a => a -> a -> a
<> forall o. Outputable o => o -> String
shown ModuleName
modName forall a. Semigroup a => a -> a -> a
<> String
" as " forall a. Semigroup a => a -> a -> a
<> forall o. Outputable o => o -> String
shown ModuleName
asName
      | (ModuleImport
modImport, Map Name (Set Name)
parents) <- 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 =
      forall a. [a] -> [[a]] -> [a]
intercalate String
", "
        [ forall o. Outputable o => o -> String
shown Name
parent forall a. Semigroup a => a -> a -> a
<> Set Name -> String
showChildren Set Name
children
        | (Name
parent, Set Name
children) <- 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 forall a. Set a -> Bool
Set.null Set Name
children then
        String
""
      else
        String
"(" forall a. Semigroup a => a -> a -> a
<> forall a. [a] -> [[a]] -> [a]
intercalate String
", " (forall o. Outputable o => o -> String
shown forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Set a -> [a]
Set.toAscList Set Name
children) forall a. Semigroup a => a -> a -> a
<> String
")"

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