{-# LANGUAGE GADTs #-}

module Plugin.HieDb (plugin) where

import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Data.IORef
import GHC
import GHC.Driver.Hooks
import GHC.Driver.Pipeline
import GHC.Driver.Pipeline.Phases
import GHC.Plugins as Plugins
import GHC.Types.Name.Cache
import HieDb.Create
import HieDb.Types
import System.Directory (doesPathExist, makeAbsolute)
import System.FilePath
import qualified System.IO.Unsafe as Unsafe

plugin :: Plugin
plugin :: Plugin
plugin =
  Plugin
defaultPlugin
    { pluginRecompile = Plugins.purePlugin
    , driverPlugin = driver
    }

driver :: [CommandLineOption] -> HscEnv -> IO HscEnv
driver :: [FilePath] -> HscEnv -> IO HscEnv
driver [FilePath]
_ HscEnv
hscEnv = do
  IO ()
initializeHiedb
  HscEnv -> IO HscEnv
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    HscEnv
hscEnv
      { hsc_hooks =
          (hsc_hooks hscEnv)
            { runPhaseHook = Just phaseHook
            }
      }
 where
  initializeHiedb :: IO ()
initializeHiedb = IO () -> IO ()
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> (HieDb -> IO ()) -> IO ()
forall a. FilePath -> (HieDb -> IO a) -> IO a
withHieDb FilePath
defaultHiedbFile HieDb -> IO ()
initConn

  -- We index using a phase hook instead of typeCheckResultAction since
  -- the hie file can be written after that plugin phase
  phaseHook :: PhaseHook
phaseHook =
    (forall a. TPhase a -> IO a) -> PhaseHook
PhaseHook ((forall a. TPhase a -> IO a) -> PhaseHook)
-> (forall a. TPhase a -> IO a) -> PhaseHook
forall a b. (a -> b) -> a -> b
$ \TPhase a
phase -> do
      case TPhase a
phase of
        T_HscPostTc HscEnv
_ ModSummary
modSummary FrontendResult
_ Messages GhcMessage
_ Maybe Fingerprint
_ -> do
          let dynFlags :: DynFlags
dynFlags = HscEnv -> DynFlags
hsc_dflags HscEnv
hscEnv
              hieDirectory :: Maybe FilePath
hieDirectory = DynFlags -> Maybe FilePath
hieDir DynFlags
dynFlags
          Maybe ()
_ <-
            IO (Maybe ()) -> IO (Maybe ())
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ()) -> IO (Maybe ())) -> IO (Maybe ()) -> IO (Maybe ())
forall a b. (a -> b) -> a -> b
$
              (FilePath -> IO ()) -> Maybe FilePath -> IO (Maybe ())
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) -> Maybe a -> m (Maybe b)
mapM
                (FilePath -> Module -> FilePath -> IO ()
addModuleToDb FilePath
defaultHiedbFile (Module -> FilePath -> IO ()) -> Module -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ ModSummary -> Module
ms_mod ModSummary
modSummary)
                Maybe FilePath
hieDirectory
          TPhase a -> IO a
forall a. TPhase a -> IO a
runPhase TPhase a
phase
        TPhase a
_ -> TPhase a -> IO a
forall a. TPhase a -> IO a
runPhase TPhase a
phase

addModuleToDb :: FilePath -> Module -> FilePath -> IO ()
addModuleToDb :: FilePath -> Module -> FilePath -> IO ()
addModuleToDb FilePath
hiedbFile Module
mod' FilePath
hieBaseDir = do
  let
    -- Note: For performance reasons we intentionally skip the type
    -- indexing phase
    -- TODO: pass this in as a user defined option
    skipOptions :: SkipOptions
skipOptions = SkipOptions
defaultSkipOptions{skipTypes = True}
    modToPath :: GenModule unit -> FilePath
modToPath = ModuleName -> FilePath
moduleNameSlashes (ModuleName -> FilePath)
-> (GenModule unit -> ModuleName) -> GenModule unit -> FilePath
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenModule unit -> ModuleName
forall unit. GenModule unit -> ModuleName
moduleName

  let hieFile :: FilePath
hieFile = FilePath
hieBaseDir FilePath -> FilePath -> FilePath
</> Module -> FilePath
forall {unit}. GenModule unit -> FilePath
modToPath Module
mod' FilePath -> FilePath -> FilePath
-<.> FilePath
".hie"

  FilePath
absoluteHieFile <- FilePath -> IO FilePath
makeAbsolute FilePath
hieFile
  Bool
hieExists <- FilePath -> IO Bool
doesPathExist FilePath
absoluteHieFile
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
hieExists (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    ()
_ <- IO () -> IO ()
withDbLock (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      IORef NameCache
nc <- NameCache -> IO (IORef NameCache)
forall a. a -> IO (IORef a)
newIORef (NameCache -> IO (IORef NameCache))
-> IO NameCache -> IO (IORef NameCache)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Char -> [Name] -> IO NameCache
initNameCache Char
'a' []
      Bool
_ <-
        FilePath -> (HieDb -> IO Bool) -> IO Bool
forall a. FilePath -> (HieDb -> IO a) -> IO a
withHieDb
          FilePath
hiedbFile
          (\HieDb
conn -> IORef NameCache -> DbMonad Bool -> IO Bool
forall a. IORef NameCache -> DbMonad a -> IO a
runDbM IORef NameCache
nc (DbMonad Bool -> IO Bool) -> DbMonad Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ HieDb -> Maybe FilePath -> SkipOptions -> FilePath -> DbMonad Bool
forall (m :: * -> *).
(MonadIO m, NameCacheMonad m) =>
HieDb -> Maybe FilePath -> SkipOptions -> FilePath -> m Bool
addRefsFrom HieDb
conn (FilePath -> Maybe FilePath
forall a. a -> Maybe a
Just FilePath
".") SkipOptions
skipOptions FilePath
absoluteHieFile)
          -- TODO: report this and maybe make configurable in future versions
          IO Bool -> (SomeException -> IO Bool) -> IO Bool
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(SomeException
_ :: SomeException) -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False)
      () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
 where
  acquireDbLock :: IO ()
acquireDbLock =
    IO () -> IO ()
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar () -> STM ()
forall a. TMVar a -> STM a
takeTMVar TMVar ()
dbLock
  releaseDbLock :: IO ()
releaseDbLock =
    IO () -> IO ()
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar () -> () -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar ()
dbLock ()
  -- Safely use a db lock - ensure the lock is released if an exception occurs
  withDbLock :: IO () -> IO ()
  withDbLock :: IO () -> IO ()
withDbLock IO ()
fn = do
    IO () -> IO () -> IO () -> IO ()
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_
      IO ()
acquireDbLock
      IO ()
releaseDbLock
      IO ()
fn

defaultHiedbFile :: String
defaultHiedbFile :: FilePath
defaultHiedbFile = FilePath
".hiedb"

-----------------------------------------------------
-- Since we cant pass state along through the phases we use unsafePerformIO
-- to define global mutable state
-----------------------------------------------------

-- | We need to ensure only one thread writes to the db at once since sqlite
-- only maintains one WAL file and will throw an error on concurrent writes
dbLock :: TMVar ()
dbLock :: TMVar ()
dbLock = IO (TMVar ()) -> TMVar ()
forall a. IO a -> a
Unsafe.unsafePerformIO (IO (TMVar ()) -> TMVar ()) -> IO (TMVar ()) -> TMVar ()
forall a b. (a -> b) -> a -> b
$ () -> IO (TMVar ())
forall a. a -> IO (TMVar a)
newTMVarIO ()
{-# NOINLINE dbLock #-}