{-# LANGUAGE BangPatterns #-}

module Hhp.Find (
    Symbol,
    SymMdlDb,
    findSymbol,
    getSymMdlDb,
    lookupSym,
) where

import GHC (DynFlags, Ghc, Module, ModuleInfo)
import qualified GHC as G
import GHC.Driver.Session (initSDocContext)
import GHC.Unit.Info (UnitInfo, mkUnit, unitExposedModules)
import GHC.Unit.State (UnitState, listUnitInfo)
import GHC.Utils.Outputable (ppr)

import Control.DeepSeq (force)
import Control.Monad.Catch (SomeException (..), catch)
import Data.Function (on)
import Data.List (groupBy, sort)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Data.Maybe (catMaybes, fromMaybe)

import Hhp.Doc (showOneLine, styleUnqualified)
import Hhp.GHCApi
import Hhp.Gap
import Hhp.Types

-- | Type of key for `SymMdlDb`.
type Symbol = String

-- | Database from 'Symbol' to modules.
newtype SymMdlDb = SymMdlDb (Map Symbol [ModuleString])

-- | Finding modules to which the symbol belong.
findSymbol :: Options -> Cradle -> Symbol -> IO String
findSymbol :: Options -> Cradle -> ModuleString -> IO ModuleString
findSymbol Options
opt Cradle
cradle ModuleString
sym = Ghc ModuleString -> IO ModuleString
forall a. Ghc a -> IO a
withGHC' (Ghc ModuleString -> IO ModuleString)
-> Ghc ModuleString -> IO ModuleString
forall a b. (a -> b) -> a -> b
$ do
    Options -> Cradle -> Ghc ()
initializeFlagsWithCradle Options
opt Cradle
cradle
    Options -> ModuleString -> SymMdlDb -> ModuleString
lookupSym Options
opt ModuleString
sym (SymMdlDb -> ModuleString) -> Ghc SymMdlDb -> Ghc ModuleString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ghc SymMdlDb
getSymMdlDb

-- | Creating 'SymMdlDb'.
getSymMdlDb :: Ghc SymMdlDb
getSymMdlDb :: Ghc SymMdlDb
getSymMdlDb = do
    [(ModuleString, ModuleString)]
sm <- Ghc DynFlags
forall (m :: * -> *). GhcMonad m => m DynFlags
G.getSessionDynFlags Ghc DynFlags
-> (DynFlags -> Ghc [(ModuleString, ModuleString)])
-> Ghc [(ModuleString, ModuleString)]
forall a b. Ghc a -> (a -> Ghc b) -> Ghc b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DynFlags -> Ghc [(ModuleString, ModuleString)]
browseAll
    let !sms :: [(ModuleString, [ModuleString])]
sms = [(ModuleString, [ModuleString])]
-> [(ModuleString, [ModuleString])]
forall a. NFData a => a -> a
force ([(ModuleString, [ModuleString])]
 -> [(ModuleString, [ModuleString])])
-> [(ModuleString, [ModuleString])]
-> [(ModuleString, [ModuleString])]
forall a b. (a -> b) -> a -> b
$ ([(ModuleString, ModuleString)] -> (ModuleString, [ModuleString]))
-> [[(ModuleString, ModuleString)]]
-> [(ModuleString, [ModuleString])]
forall a b. (a -> b) -> [a] -> [b]
map [(ModuleString, ModuleString)] -> (ModuleString, [ModuleString])
forall {a} {b}. [(a, b)] -> (a, [b])
tieup ([[(ModuleString, ModuleString)]]
 -> [(ModuleString, [ModuleString])])
-> [[(ModuleString, ModuleString)]]
-> [(ModuleString, [ModuleString])]
forall a b. (a -> b) -> a -> b
$ ((ModuleString, ModuleString)
 -> (ModuleString, ModuleString) -> Bool)
-> [(ModuleString, ModuleString)]
-> [[(ModuleString, ModuleString)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (ModuleString -> ModuleString -> Bool
forall a. Eq a => a -> a -> Bool
(==) (ModuleString -> ModuleString -> Bool)
-> ((ModuleString, ModuleString) -> ModuleString)
-> (ModuleString, ModuleString)
-> (ModuleString, ModuleString)
-> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (ModuleString, ModuleString) -> ModuleString
forall a b. (a, b) -> a
fst) ([(ModuleString, ModuleString)]
 -> [[(ModuleString, ModuleString)]])
-> [(ModuleString, ModuleString)]
-> [[(ModuleString, ModuleString)]]
forall a b. (a -> b) -> a -> b
$ [(ModuleString, ModuleString)] -> [(ModuleString, ModuleString)]
forall a. Ord a => [a] -> [a]
sort [(ModuleString, ModuleString)]
sm
        !m :: Map ModuleString [ModuleString]
m = Map ModuleString [ModuleString] -> Map ModuleString [ModuleString]
forall a. NFData a => a -> a
force (Map ModuleString [ModuleString]
 -> Map ModuleString [ModuleString])
-> Map ModuleString [ModuleString]
-> Map ModuleString [ModuleString]
forall a b. (a -> b) -> a -> b
$ [(ModuleString, [ModuleString])] -> Map ModuleString [ModuleString]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(ModuleString, [ModuleString])]
sms
    SymMdlDb -> Ghc SymMdlDb
forall a. a -> Ghc a
forall (m :: * -> *) a. Monad m => a -> m a
return (Map ModuleString [ModuleString] -> SymMdlDb
SymMdlDb Map ModuleString [ModuleString]
m)
  where
    tieup :: [(a, b)] -> (a, [b])
tieup [(a, b)]
x = ([a] -> a
forall a. [a] -> a
unsafeHead (((a, b) -> a) -> [(a, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> a
forall a b. (a, b) -> a
fst [(a, b)]
x), ((a, b) -> b) -> [(a, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> b
forall a b. (a, b) -> b
snd [(a, b)]
x)

-- | Looking up 'SymMdlDb' with 'Symbol' to find modules.
lookupSym :: Options -> Symbol -> SymMdlDb -> String
lookupSym :: Options -> ModuleString -> SymMdlDb -> ModuleString
lookupSym Options
opt ModuleString
sym (SymMdlDb Map ModuleString [ModuleString]
db) = Options -> [ModuleString] -> ModuleString
forall a. ToString a => Options -> a -> ModuleString
convert Options
opt ([ModuleString] -> ModuleString) -> [ModuleString] -> ModuleString
forall a b. (a -> b) -> a -> b
$ [ModuleString] -> Maybe [ModuleString] -> [ModuleString]
forall a. a -> Maybe a -> a
fromMaybe [] (ModuleString
-> Map ModuleString [ModuleString] -> Maybe [ModuleString]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ModuleString
sym Map ModuleString [ModuleString]
db)

----------------------------------------------------------------

-- | Browsing all functions in all system/user modules.
browseAll :: DynFlags -> Ghc [(String, String)]
browseAll :: DynFlags -> Ghc [(ModuleString, ModuleString)]
browseAll DynFlags
dflag = do
    [Module]
ms <- UnitState -> [Module]
packageModules (UnitState -> [Module]) -> Ghc UnitState -> Ghc [Module]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ghc UnitState
getUnitState
    [Maybe ModuleInfo]
is <- [Maybe (Maybe ModuleInfo)] -> [Maybe ModuleInfo]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Maybe ModuleInfo)] -> [Maybe ModuleInfo])
-> Ghc [Maybe (Maybe ModuleInfo)] -> Ghc [Maybe ModuleInfo]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Module -> Ghc (Maybe (Maybe ModuleInfo)))
-> [Module] -> Ghc [Maybe (Maybe ModuleInfo)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Module -> Ghc (Maybe (Maybe ModuleInfo))
getMaybeModuleInfo [Module]
ms
    [(ModuleString, ModuleString)]
-> Ghc [(ModuleString, ModuleString)]
forall a. a -> Ghc a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(ModuleString, ModuleString)]
 -> Ghc [(ModuleString, ModuleString)])
-> [(ModuleString, ModuleString)]
-> Ghc [(ModuleString, ModuleString)]
forall a b. (a -> b) -> a -> b
$ ((Module, Maybe ModuleInfo) -> [(ModuleString, ModuleString)])
-> [(Module, Maybe ModuleInfo)] -> [(ModuleString, ModuleString)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (DynFlags
-> (Module, Maybe ModuleInfo) -> [(ModuleString, ModuleString)]
toNameModule DynFlags
dflag) ([Module] -> [Maybe ModuleInfo] -> [(Module, Maybe ModuleInfo)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Module]
ms [Maybe ModuleInfo]
is)

-- ghc-bignum causes errors, sigh.
getMaybeModuleInfo :: Module -> Ghc (Maybe (Maybe ModuleInfo))
getMaybeModuleInfo :: Module -> Ghc (Maybe (Maybe ModuleInfo))
getMaybeModuleInfo Module
x = Maybe ModuleInfo -> Maybe (Maybe ModuleInfo)
forall a. a -> Maybe a
Just (Maybe ModuleInfo -> Maybe (Maybe ModuleInfo))
-> Ghc (Maybe ModuleInfo) -> Ghc (Maybe (Maybe ModuleInfo))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Module -> Ghc (Maybe ModuleInfo)
forall (m :: * -> *). GhcMonad m => Module -> m (Maybe ModuleInfo)
G.getModuleInfo Module
x Ghc (Maybe ModuleInfo)
-> (SomeException -> Ghc (Maybe ModuleInfo))
-> Ghc (Maybe ModuleInfo)
forall e a.
(HasCallStack, Exception e) =>
Ghc a -> (e -> Ghc a) -> Ghc a
forall (m :: * -> *) e a.
(MonadCatch m, HasCallStack, Exception e) =>
m a -> (e -> m a) -> m a
`catch` (\(SomeException e
_) -> Maybe ModuleInfo -> Ghc (Maybe ModuleInfo)
forall a. a -> Ghc a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ModuleInfo
forall a. Maybe a
Nothing)

toNameModule :: DynFlags -> (Module, Maybe ModuleInfo) -> [(String, String)]
toNameModule :: DynFlags
-> (Module, Maybe ModuleInfo) -> [(ModuleString, ModuleString)]
toNameModule DynFlags
_ (Module
_, Maybe ModuleInfo
Nothing) = []
toNameModule DynFlags
dflag (Module
m, Just ModuleInfo
inf) = (Name -> (ModuleString, ModuleString))
-> [Name] -> [(ModuleString, ModuleString)]
forall a b. (a -> b) -> [a] -> [b]
map (\Name
name -> (Name -> ModuleString
toStr Name
name, ModuleString
mdl)) [Name]
names
  where
    mdl :: ModuleString
mdl = ModuleName -> ModuleString
G.moduleNameString (Module -> ModuleName
forall unit. GenModule unit -> ModuleName
G.moduleName Module
m)
    names :: [Name]
names = ModuleInfo -> [Name]
G.modInfoExports ModuleInfo
inf
    toStr :: Name -> ModuleString
toStr = SDocContext -> SDoc -> ModuleString
showOneLine (DynFlags -> PprStyle -> SDocContext
initSDocContext DynFlags
dflag PprStyle
styleUnqualified) (SDoc -> ModuleString) -> (Name -> SDoc) -> Name -> ModuleString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> SDoc
forall a. Outputable a => a -> SDoc
ppr

packageModules :: UnitState -> [Module]
packageModules :: UnitState -> [Module]
packageModules UnitState
us = (UnitInfo -> [Module]) -> [UnitInfo] -> [Module]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap UnitInfo -> [Module]
fromUnitInfo ([UnitInfo] -> [Module]) -> [UnitInfo] -> [Module]
forall a b. (a -> b) -> a -> b
$ UnitState -> [UnitInfo]
listUnitInfo UnitState
us

fromUnitInfo :: UnitInfo -> [Module]
fromUnitInfo :: UnitInfo -> [Module]
fromUnitInfo UnitInfo
uinfo = [Module]
modules
  where
    uid :: Unit
uid = UnitInfo -> Unit
mkUnit UnitInfo
uinfo
    moduleNames :: [ModuleName]
moduleNames = ((ModuleName, Maybe Module) -> ModuleName)
-> [(ModuleName, Maybe Module)] -> [ModuleName]
forall a b. (a -> b) -> [a] -> [b]
map (ModuleName, Maybe Module) -> ModuleName
forall a b. (a, b) -> a
fst ([(ModuleName, Maybe Module)] -> [ModuleName])
-> [(ModuleName, Maybe Module)] -> [ModuleName]
forall a b. (a -> b) -> a -> b
$ UnitInfo -> [(ModuleName, Maybe Module)]
forall srcpkgid srcpkgname uid modulename mod.
GenericUnitInfo srcpkgid srcpkgname uid modulename mod
-> [(modulename, Maybe mod)]
unitExposedModules UnitInfo
uinfo
    modules :: [Module]
modules = (ModuleName -> Module) -> [ModuleName] -> [Module]
forall a b. (a -> b) -> [a] -> [b]
map (Unit -> ModuleName -> Module
forall u. u -> ModuleName -> GenModule u
G.mkModule Unit
uid) [ModuleName]
moduleNames