{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

module QualifiedImportsPlugin where

import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Generics
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (mapMaybe)

#if __GLASGOW_HASKELL__ >= 900
import GHC.Driver.Main (getHscEnv)
import GHC.Hs
import GHC.Plugins hiding (getHscEnv, (<>))
import GHC.Utils.Error
import GHC.Data.Bag

unhelpfulOther :: FastString -> UnhelpfulSpanReason
unhelpfulOther = UnhelpfulOther
#else
import GHC.Hs hiding (HsModule)
import qualified GHC.Hs as GHC
import Bag
import GhcPlugins hiding (getHscEnv, (<>))
import ErrUtils
import HscMain

type HsModule = GHC.HsModule GhcPs

unitState :: a -> a
unitState :: a -> a
unitState = a -> a
forall a. a -> a
id

unhelpfulOther :: String -> FastString
unhelpfulOther :: String -> FastString
unhelpfulOther = String -> FastString
mkFastString
#endif

defaultImports :: [(String, String)]
defaultImports :: [(String, String)]
defaultImports =
  [ (String
"Data.Text", String
"Text"),
    (String
"Data.Text.IO", String
"Text"),
    (String
"Data.Text.Lazy", String
"LText"),
    (String
"Data.Text.Lazy.IO", String
"LText"),
    (String
"Data.ByteString", String
"ByteString"),
    (String
"Data.ByteString.Lazy", String
"LByteString"),
    (String
"Data.Map.Strict", String
"Map"),
    (String
"Data.Map.Lazy", String
"LMap"),
    (String
"Data.IntMap.Strict", String
"IntMap"),
    (String
"Data.IntMap.Lazy", String
"LIntMap"),
    (String
"Data.HashMap.Strict", String
"HashMap"),
    (String
"Data.HashMap.Lazy", String
"LHashMap"),
    (String
"Data.HashSet", String
"HashSet"),
    (String
"Data.Set", String
"Set"),
    (String
"Data.Aeson", String
"Aeson"),
    (String
"Data.Vector", String
"Vector"),
    (String
"Data.Vector.Mutable", String
"MVector")
  ]

data Opts = Opts
  { Opts -> Bool
optsNoDefaults :: Bool,
    Opts -> [(String, String)]
optsCustomImports :: [(String, String)]
  }

instance Semigroup Opts where
  Opts Bool
a1 [(String, String)]
a2 <> :: Opts -> Opts -> Opts
<> Opts Bool
b1 [(String, String)]
b2 = Bool -> [(String, String)] -> Opts
Opts (Bool
a1 Bool -> Bool -> Bool
|| Bool
b1) ([(String, String)]
a2 [(String, String)] -> [(String, String)] -> [(String, String)]
forall a. Semigroup a => a -> a -> a
<> [(String, String)]
b2)

instance Monoid Opts where
  mempty :: Opts
mempty = Bool -> [(String, String)] -> Opts
Opts Bool
False []

parseOpts :: [CommandLineOption] -> Hsc Opts
parseOpts :: [String] -> Hsc Opts
parseOpts [] = Opts -> Hsc Opts
forall (m :: * -> *) a. Monad m => a -> m a
return Opts
forall a. Monoid a => a
mempty
parseOpts (String
x : [String]
xs) = case String -> Maybe Opts
parseOpt String
x of
  Maybe Opts
Nothing -> do
    () <- (HscEnv -> WarningMessages -> IO ((), WarningMessages)) -> Hsc ()
forall a.
(HscEnv -> WarningMessages -> IO (a, WarningMessages)) -> Hsc a
Hsc ((HscEnv -> WarningMessages -> IO ((), WarningMessages)) -> Hsc ())
-> (HscEnv -> WarningMessages -> IO ((), WarningMessages))
-> Hsc ()
forall a b. (a -> b) -> a -> b
$ \HscEnv
env WarningMessages
wm ->
      -- There must a better way than delving into Hsc manually, but I couldn't find the function
      -- 'WarnMsg -> Hsc ()'.
      let msg :: ErrMsg
msg =
            DynFlags -> SrcSpan -> MsgDoc -> ErrMsg
mkPlainWarnMsg
              (HscEnv -> DynFlags
hsc_dflags HscEnv
env)
              (FastString -> SrcSpan
UnhelpfulSpan (FastString -> SrcSpan) -> FastString -> SrcSpan
forall a b. (a -> b) -> a -> b
$ String -> FastString
unhelpfulOther String
"QualifiedImportsPlugin")
              (MsgDoc
"Unknown argument:" MsgDoc -> MsgDoc -> MsgDoc
<+> String -> MsgDoc
text String
x)
       in ((), WarningMessages) -> IO ((), WarningMessages)
forall (m :: * -> *) a. Monad m => a -> m a
return ((), ErrMsg -> WarningMessages -> WarningMessages
forall a. a -> Bag a -> Bag a
consBag ErrMsg
msg WarningMessages
wm)
    [String] -> Hsc Opts
parseOpts [String]
xs
  Just Opts
opts -> (Opts
opts Opts -> Opts -> Opts
forall a. Semigroup a => a -> a -> a
<>) (Opts -> Opts) -> Hsc Opts -> Hsc Opts
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String] -> Hsc Opts
parseOpts [String]
xs

{-
Valid options:
- no-defaults
- <actual_module_name>:<qualified_name>
-}
parseOpt :: CommandLineOption -> Maybe Opts
parseOpt :: String -> Maybe Opts
parseOpt String
"no-defaults" = Opts -> Maybe Opts
forall a. a -> Maybe a
Just (Opts -> Maybe Opts) -> Opts -> Maybe Opts
forall a b. (a -> b) -> a -> b
$ Opts
forall a. Monoid a => a
mempty {optsNoDefaults :: Bool
optsNoDefaults = Bool
True}
parseOpt String
xs =
  -- Check if something looks like two module names separated by a colon. This still allows some
  -- invalid module names, but I think it is good enough.
  case (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
':') String
xs of
    (String
from, Char
':' : String
to)
      | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
allowed String
from Bool -> Bool -> Bool
&& (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
allowed String
to -> Opts -> Maybe Opts
forall a. a -> Maybe a
Just (Opts -> Maybe Opts) -> Opts -> Maybe Opts
forall a b. (a -> b) -> a -> b
$ Opts
forall a. Monoid a => a
mempty {optsCustomImports :: [(String, String)]
optsCustomImports = [(String
from, String
to)]}
    (String, String)
_ -> Maybe Opts
forall a. Maybe a
Nothing
  where
    allowed :: Char -> Bool
allowed Char
c =
      [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or
        [ Char
'a' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'z',
          Char
'A' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'Z',
          Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'.'
        ]

plugin :: Plugin
plugin :: Plugin
plugin =
  Plugin
defaultPlugin
    { pluginRecompile :: [String] -> IO PluginRecompile
pluginRecompile = [String] -> IO PluginRecompile
purePlugin,
      parsedResultAction :: [String] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
parsedResultAction = \[String]
args ModSummary
_ HsParsedModule
parsed -> do
        Opts
opts <- [String] -> Hsc Opts
parseOpts [String]
args
        GenLocated SrcSpan HsModule
nm <-
          HsParsedModule -> GenLocated SrcSpan HsModule
hpm_module HsParsedModule
parsed
            GenLocated SrcSpan HsModule
-> (HsModule -> Hsc HsModule) -> GenLocated SrcSpan (Hsc HsModule)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Opts -> HsModule -> Hsc HsModule
modifyHsMod Opts
opts
            GenLocated SrcSpan (Hsc HsModule)
-> (GenLocated SrcSpan (Hsc HsModule)
    -> Hsc (GenLocated SrcSpan HsModule))
-> Hsc (GenLocated SrcSpan HsModule)
forall a b. a -> (a -> b) -> b
& GenLocated SrcSpan (Hsc HsModule)
-> Hsc (GenLocated SrcSpan HsModule)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA
        HsParsedModule -> Hsc HsParsedModule
forall (m :: * -> *) a. Monad m => a -> m a
return (HsParsedModule -> Hsc HsParsedModule)
-> HsParsedModule -> Hsc HsParsedModule
forall a b. (a -> b) -> a -> b
$ HsParsedModule
parsed {hpm_module :: GenLocated SrcSpan HsModule
hpm_module = GenLocated SrcSpan HsModule
nm}
    }

{-
Insert qualified import statements to a module.

It's not very straightforward which imports to insert. We can not insert all of them,
because then GHC will complain about missing modules/modules in hidden packages. We can
only insert the ones accessed from the code, but then if the user makes a typo in the name,
GHC won't suggest the correct names.

So:

1. We start with a combination of 'defaultImports' and user specified imports.
2. For every candidate import, we insert if it's either:
   * Referred from the source code
   * Available to import
3. When we insert an import which if referred from the source code, we use the source
   location of the referent as the source code of the import, in order for the error
   message to be useful.
-}
modifyHsMod :: Opts -> HsModule -> Hsc HsModule
modifyHsMod :: Opts -> HsModule -> Hsc HsModule
modifyHsMod Opts
opts HsModule
m = do
  HscEnv
env <- Hsc HscEnv
getHscEnv

  let imports :: [(String, String)]
imports =
        (if Opts -> Bool
optsNoDefaults Opts
opts then [] else [(String, String)]
defaultImports)
          [(String, String)] -> [(String, String)] -> [(String, String)]
forall a. [a] -> [a] -> [a]
++ Opts -> [(String, String)]
optsCustomImports Opts
opts

      refs :: Map ModuleName SrcSpan
refs = HsModule -> Map ModuleName SrcSpan
referencedModules HsModule
m
      newImports :: [GenLocated SrcSpan (ImportDecl (GhcPass p))]
newImports =
        [(String, String)]
imports
          [(String, String)]
-> ([(String, String)] -> [(ModuleName, ModuleName)])
-> [(ModuleName, ModuleName)]
forall a b. a -> (a -> b) -> b
& ((String, String) -> (ModuleName, ModuleName))
-> [(String, String)] -> [(ModuleName, ModuleName)]
forall a b. (a -> b) -> [a] -> [b]
map (\(String
n, String
qn) -> (String -> ModuleName
mkModuleName String
n, String -> ModuleName
mkModuleName String
qn))
          [(ModuleName, ModuleName)]
-> ([(ModuleName, ModuleName)]
    -> [(Maybe SrcSpan, ModuleName, ModuleName)])
-> [(Maybe SrcSpan, ModuleName, ModuleName)]
forall a b. a -> (a -> b) -> b
& ((ModuleName, ModuleName)
 -> Maybe (Maybe SrcSpan, ModuleName, ModuleName))
-> [(ModuleName, ModuleName)]
-> [(Maybe SrcSpan, ModuleName, ModuleName)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
            ( \(ModuleName
n, ModuleName
qn) ->
                case ModuleName -> Map ModuleName SrcSpan -> Maybe SrcSpan
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ModuleName
qn Map ModuleName SrcSpan
refs of
                  Maybe SrcSpan
Nothing
                    | HscEnv -> ModuleName -> Bool
isModuleAvailable HscEnv
env ModuleName
n -> (Maybe SrcSpan, ModuleName, ModuleName)
-> Maybe (Maybe SrcSpan, ModuleName, ModuleName)
forall a. a -> Maybe a
Just (Maybe SrcSpan
forall a. Maybe a
Nothing, ModuleName
n, ModuleName
qn)
                    | Bool
otherwise -> Maybe (Maybe SrcSpan, ModuleName, ModuleName)
forall a. Maybe a
Nothing
                  Just SrcSpan
loc ->
                    (Maybe SrcSpan, ModuleName, ModuleName)
-> Maybe (Maybe SrcSpan, ModuleName, ModuleName)
forall a. a -> Maybe a
Just (SrcSpan -> Maybe SrcSpan
forall a. a -> Maybe a
Just SrcSpan
loc, ModuleName
n, ModuleName
qn)
            )
          [(Maybe SrcSpan, ModuleName, ModuleName)]
-> ([(Maybe SrcSpan, ModuleName, ModuleName)]
    -> [GenLocated SrcSpan (ImportDecl (GhcPass p))])
-> [GenLocated SrcSpan (ImportDecl (GhcPass p))]
forall a b. a -> (a -> b) -> b
& ((Maybe SrcSpan, ModuleName, ModuleName)
 -> GenLocated SrcSpan (ImportDecl (GhcPass p)))
-> [(Maybe SrcSpan, ModuleName, ModuleName)]
-> [GenLocated SrcSpan (ImportDecl (GhcPass p))]
forall a b. (a -> b) -> [a] -> [b]
map
            ( \(Maybe SrcSpan
loc, ModuleName
n, ModuleName
qn) ->
                ((ImportDecl (GhcPass p)
 -> GenLocated SrcSpan (ImportDecl (GhcPass p)))
-> (SrcSpan
    -> ImportDecl (GhcPass p)
    -> GenLocated SrcSpan (ImportDecl (GhcPass p)))
-> Maybe SrcSpan
-> ImportDecl (GhcPass p)
-> GenLocated SrcSpan (ImportDecl (GhcPass p))
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ImportDecl (GhcPass p)
-> GenLocated SrcSpan (ImportDecl (GhcPass p))
forall a. HasSrcSpan a => SrcSpanLess a -> a
noLoc SrcSpan
-> ImportDecl (GhcPass p)
-> GenLocated SrcSpan (ImportDecl (GhcPass p))
forall l e. l -> e -> GenLocated l e
L Maybe SrcSpan
loc)
                  (ModuleName -> ImportDecl (GhcPass p)
forall (p :: Pass). ModuleName -> ImportDecl (GhcPass p)
simpleImportDecl ModuleName
n)
                    { ideclQualified :: ImportDeclQualifiedStyle
ideclQualified = ImportDeclQualifiedStyle
QualifiedPre,
                      ideclAs :: Maybe (Located ModuleName)
ideclAs = Located ModuleName -> Maybe (Located ModuleName)
forall a. a -> Maybe a
Just (SrcSpanLess (Located ModuleName) -> Located ModuleName
forall a. HasSrcSpan a => SrcSpanLess a -> a
noLoc SrcSpanLess (Located ModuleName)
ModuleName
qn),
                      -- This makes GHC not complain about unused imports.
                      ideclImplicit :: Bool
ideclImplicit = Bool
True
                    }
            )
  HsModule -> Hsc HsModule
forall (m :: * -> *) a. Monad m => a -> m a
return (HsModule -> Hsc HsModule) -> HsModule -> Hsc HsModule
forall a b. (a -> b) -> a -> b
$ HsModule
m {hsmodImports :: [LImportDecl GhcPs]
hsmodImports = HsModule -> [LImportDecl GhcPs]
forall pass. HsModule pass -> [LImportDecl pass]
hsmodImports HsModule
m [LImportDecl GhcPs] -> [LImportDecl GhcPs] -> [LImportDecl GhcPs]
forall a. [a] -> [a] -> [a]
++ [LImportDecl GhcPs]
forall (p :: Pass). [GenLocated SrcSpan (ImportDecl (GhcPass p))]
newImports}

-- Figure out if the module is available to import.
isModuleAvailable :: HscEnv -> ModuleName -> Bool
isModuleAvailable :: HscEnv -> ModuleName -> Bool
isModuleAvailable HscEnv
env ModuleName
n =
  let us :: DynFlags
us = DynFlags -> DynFlags
forall a. a -> a
unitState (DynFlags -> DynFlags) -> DynFlags -> DynFlags
forall a b. (a -> b) -> a -> b
$ HscEnv -> DynFlags
hsc_dflags HscEnv
env
   in case DynFlags -> ModuleName -> Maybe FastString -> LookupResult
lookupModuleWithSuggestions DynFlags
us ModuleName
n Maybe FastString
forall a. Maybe a
Nothing of
        LookupFound Module
_ PackageConfig
_ -> Bool
True -- "found"
        LookupMultiple [(Module, ModuleOrigin)]
_ -> Bool
False -- "multiple"
        LookupHidden [(Module, ModuleOrigin)]
_ [(Module, ModuleOrigin)]
_ -> Bool
False -- "hidden " ++ show (length l) ++ " " ++ show (length r)
        LookupUnusable [(Module, ModuleOrigin)]
_ -> Bool
False -- "unusable"
        LookupNotFound [ModuleSuggestion]
_ -> Bool
False -- "not found"

-- We also carry the references to modules, so we can the error messages can point to
-- the use site (since there is no visible import statement).
referencedModules :: HsModule -> Map ModuleName SrcSpan
referencedModules :: HsModule -> Map ModuleName SrcSpan
referencedModules HsModule
m =
  HsModule -> [LHsDecl GhcPs]
forall pass. HsModule pass -> [LHsDecl pass]
hsmodDecls HsModule
m
    [LHsDecl GhcPs]
-> ([LHsDecl GhcPs] -> [Map ModuleName SrcSpan])
-> [Map ModuleName SrcSpan]
forall a b. a -> (a -> b) -> b
& (LHsDecl GhcPs -> Map ModuleName SrcSpan)
-> [LHsDecl GhcPs] -> [Map ModuleName SrcSpan]
forall a b. (a -> b) -> [a] -> [b]
map LHsDecl GhcPs -> Map ModuleName SrcSpan
forall a. Data a => a -> Map ModuleName SrcSpan
go
    [Map ModuleName SrcSpan]
-> ([Map ModuleName SrcSpan] -> [Map ModuleName SrcSpan])
-> [Map ModuleName SrcSpan]
forall a b. a -> (a -> b) -> b
& [Map ModuleName SrcSpan] -> [Map ModuleName SrcSpan]
forall a. [a] -> [a]
reverse
    [Map ModuleName SrcSpan]
-> ([Map ModuleName SrcSpan] -> Map ModuleName SrcSpan)
-> Map ModuleName SrcSpan
forall a b. a -> (a -> b) -> b
& [Map ModuleName SrcSpan] -> Map ModuleName SrcSpan
forall a. Monoid a => [a] -> a
mconcat
  where
    go :: Data a => a -> Map ModuleName SrcSpan
    go :: a -> Map ModuleName SrcSpan
go =
      (Map ModuleName SrcSpan
 -> Map ModuleName SrcSpan -> Map ModuleName SrcSpan)
-> Map ModuleName SrcSpan
-> (forall a. Data a => a -> Map ModuleName SrcSpan)
-> a
-> Map ModuleName SrcSpan
forall a r r'.
Data a =>
(r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> a -> r
gmapQr
        ((Map ModuleName SrcSpan
 -> Map ModuleName SrcSpan -> Map ModuleName SrcSpan)
-> Map ModuleName SrcSpan
-> Map ModuleName SrcSpan
-> Map ModuleName SrcSpan
forall a b c. (a -> b -> c) -> b -> a -> c
flip Map ModuleName SrcSpan
-> Map ModuleName SrcSpan -> Map ModuleName SrcSpan
forall a. Monoid a => a -> a -> a
mappend)
        Map ModuleName SrcSpan
forall a. Monoid a => a
mempty
        ( \d
d ->
            case d -> Maybe (Located RdrName)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast @_ @(Located RdrName) d
d of
              Maybe (Located RdrName)
Nothing -> d -> Map ModuleName SrcSpan
forall a. Data a => a -> Map ModuleName SrcSpan
go d
d
              Just (L SrcSpan
loc (Qual ModuleName
m OccName
_)) -> ModuleName -> SrcSpan -> Map ModuleName SrcSpan
forall k a. k -> a -> Map k a
Map.singleton ModuleName
m SrcSpan
loc
              Just Located RdrName
_ -> Map ModuleName SrcSpan
forall a. Monoid a => a
mempty
        )