{-# LANGUAGE BangPatterns #-}

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

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

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 (fromMaybe, catMaybes)

import Hhp.Doc (showOneLine, styleUnqualified)
import Hhp.Gap
import Hhp.GHCApi
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 = forall a. Ghc a -> IO a
withGHC' 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 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 <- forall (m :: * -> *). GhcMonad m => m DynFlags
G.getSessionDynFlags forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DynFlags -> Ghc [(ModuleString, ModuleString)]
browseAll
    let !sms :: [(ModuleString, [ModuleString])]
sms = forall a. NFData a => a -> a
force forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b}. [(a, b)] -> (a, [b])
tieup forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [a] -> [a]
sort [(ModuleString, ModuleString)]
sm
        !m :: Map ModuleString [ModuleString]
m = forall a. NFData a => a -> a
force forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(ModuleString, [ModuleString])]
sms
    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 = (forall a. [a] -> a
head (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(a, b)]
x), forall a b. (a -> b) -> [a] -> [b]
map 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) = forall a. ToString a => Options -> a -> ModuleString
convert Options
opt forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe [] (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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ghc UnitState
getUnitState
    [Maybe ModuleInfo]
is <- forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Module -> Ghc (Maybe (Maybe ModuleInfo))
getMaybeModuleInfo [Module]
ms
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (DynFlags
-> (Module, Maybe ModuleInfo) -> [(ModuleString, ModuleString)]
toNameModule DynFlags
dflag) (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 = forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). GhcMonad m => Module -> m (Maybe ModuleInfo)
G.getModuleInfo Module
x forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` (\(SomeException e
_) -> forall (m :: * -> *) a. Monad m => a -> m a
return 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) = 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 (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) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Outputable a => a -> SDoc
ppr

packageModules :: UnitState -> [Module]
packageModules :: UnitState -> [Module]
packageModules UnitState
us = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap UnitInfo -> [Module]
fromUnitInfo 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 = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall compid srcpkgid srcpkgname uid modulename mod.
GenericUnitInfo compid srcpkgid srcpkgname uid modulename mod
-> [(modulename, Maybe mod)]
unitExposedModules UnitInfo
uinfo
    modules :: [Module]
modules = forall a b. (a -> b) -> [a] -> [b]
map (forall u. u -> ModuleName -> GenModule u
G.mkModule Unit
uid) [ModuleName]
moduleNames