Environment.hs

{-# LANGUAGE RecordWildCards, CPP #-}
{-# LANGUAGE TupleSections #-}
module HIE.Bios.Environment (initSession, initSession', getRuntimeGhcLibDir, getRuntimeGhcVersion, makeDynFlagsAbsolute, makeTargetsAbsolute, getCacheDir, addCmdOpts, extractUnits) where

import GHC (GhcMonad)
import qualified GHC as G

import Control.Applicative
import Control.Monad (void)
import Control.Monad.IO.Class

import System.Directory
import System.FilePath
import System.Environment (lookupEnv)

import qualified Crypto.Hash.SHA1 as H
import qualified Data.ByteString.Char8 as B
import Data.ByteString.Base16
import Data.List
import Data.Char (isSpace)
import Text.ParserCombinators.ReadP hiding (optional)

import HIE.Bios.Types
import qualified HIE.Bios.Ghc.Gap as Gap
import qualified System.OsPath as OsPath
import qualified GHC.Plugins as G
import qualified Data.List.NonEmpty as NE
import qualified GHC.Driver.Monad as G
import Data.IORef (newIORef, readIORef)
import qualified GHC.Unit.Env as G
import GHC.Driver.Env (HscEnv(hsc_unit_env))
import Data.Maybe (fromMaybe)

-- | Start a GHC session and set some sensible options for tooling to use.
-- Creates a folder in the cache directory to cache interface files to make
-- reloading faster.
initSession :: (GhcMonad m)
    => ComponentOptions
    -> m [G.Target]
initSession = initSession' False

initSession' :: (GhcMonad m)
    => Bool
    -> ComponentOptions
    -> m [G.Target]
initSession' workAroundThreadUnsafety ComponentOptions {..} = do
    -- Create a unique folder per set of different GHC options, assuming that each different set of
    -- GHC options will create incompatible interface files.
    let
      -- There seems to be a race condition when writing interface files
      hash_args = (if workAroundThreadUnsafety then (componentRoot :) else id) componentOptions
      opts_hash = B.unpack $ encode $ H.finalize $ H.updates H.init $ map B.pack hash_args

    cache_dir <- liftIO $ makeAbsolute =<< getCacheDir opts_hash

    -- Plan:
    -- - Extract `-unit @resp_file` options if present
    -- - Add the user specified options to a fresh GHC session
    -- - Make all dflags absolute
    -- - Set flags for interactive session
    -- - Reconstruct targets and make them absolute

    let (units,rest) = extractUnits componentOptions

    let liftGhc m = do
          env <- G.getSession
          s <- liftIO (newIORef env)
          x <- liftIO $ G.reflectGhc m (G.Session s)
          G.setSession =<< liftIO (readIORef s)
          pure x
    let
      setFlags df'' =
        void $ G.setSessionDynFlags
          (disableOptimisation -- Compile with -O0 as we are not going to produce object files.
          $ setIgnoreInterfacePragmas            -- Ignore any non-essential information in interface files such as unfoldings changing.
          $ writeInterfaceFiles (Just cache_dir) -- Write interface files to the cache
          $ setVerbosity 0                       -- Set verbosity to zero just in case the user specified `-vx` in the options.
          $ Gap.setWayDynamicIfHostIsDynamic     -- Add dynamic way if GHC is built with dynamic linking
          $ setLinkerOptions df''                -- Set `-fno-code` to avoid generating object files, unless we have to.
          )

    targets <- case NE.nonEmpty units of
      Nothing -> do
        (df',targets) <- addCmdOpts componentOptions =<< G.getSessionDynFlags
        setFlags $ makeDynFlagsAbsolute componentRoot df'
        pure targets
      Just ne_units -> do
        logger <- Gap.getLogger <$> G.getSession
        df <- G.getSessionDynFlags
        -- leftovers won't have targets when units are present.
        (df', _leftovers, _warns) <- Gap.parseDynamicFlags logger df (map G.noLoc rest)
        setFlags $ makeDynFlagsAbsolute componentRoot df'

        hs_srcs <- liftGhc $ Gap.initMulti ne_units (\ _ _ _ _ -> return ())

        G.modifySession $ G.hscUpdateHUG $ fmap (\ ue -> ue {G.homeUnitEnv_dflags = makeDynFlagsAbsolute componentRoot (G.homeUnitEnv_dflags ue)})

        -- This should be an enforced invariant.
        G.modifySession $ \ hsc_env -> G.hscSetActiveUnitId (G.ue_current_unit $ G.hsc_unit_env hsc_env) hsc_env

        mapM (\(src, uid, phase) -> G.guessTarget src uid phase) hs_srcs

    let mkTgtAbs env tgt = tgt {G.targetId = makeTargetIdAbsolute wdir (G.targetId tgt)}
          where
            wdir = fromMaybe componentRoot $
              G.workingDirectory
              . G.homeUnitEnv_dflags . G.ue_findHomeUnitEnv uid
              $ hsc_unit_env env
            uid = G.targetUnitId tgt

    env <- G.getSession
    let targets' = map (mkTgtAbs env) targets

    -- Unset the default log action to avoid output going to stdout.
    Gap.unsetLogAction

    return targets'

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

makeTargetsAbsolute :: FilePath -> [G.Target] -> [G.Target]
makeTargetsAbsolute wdir = map (\target -> target {G.targetId = makeTargetIdAbsolute wdir (G.targetId target)})

makeTargetIdAbsolute :: FilePath -> G.TargetId -> G.TargetId
makeTargetIdAbsolute wdir (G.TargetFile fp phase) = G.TargetFile (wdir </> fp) phase
makeTargetIdAbsolute _ tid = tid

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

-- | @getRuntimeGhcLibDir cradle@ will give you the ghc libDir:
-- __do not__ use 'runGhcCmd' directly.
--
--
-- Obtains libdir by calling 'runCradleGhc' on the provided cradle.
getRuntimeGhcLibDir :: Cradle a
                    -> IO (CradleLoadResult FilePath)
getRuntimeGhcLibDir cradle = fmap (fmap trim) $
      runGhcCmd (cradleOptsProg cradle) ["--print-libdir"]

-- | Gets the version of ghc used when compiling the cradle. It is based off of
-- 'getRuntimeGhcLibDir'. If it can't work out the verison reliably, it will
-- return a 'CradleError'
getRuntimeGhcVersion :: Cradle a
                     -> IO (CradleLoadResult String)
getRuntimeGhcVersion cradle =
  fmap (fmap trim) $ runGhcCmd (cradleOptsProg cradle) ["--numeric-version"]

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

-- | What to call the cache directory in the cache folder.
cacheDir :: String
cacheDir = "hie-bios"

{- |
Back in the day we used to clear the cache at the start of each session,
however, it's not really necessary as
1. There is one cache dir for any change in options.
2. Interface files are resistent to bad option changes anyway.

> clearInterfaceCache :: FilePath -> IO ()
> clearInterfaceCache fp = do
>   cd <- getCacheDir fp
>   res <- doesPathExist cd
>   when res (removeDirectoryRecursive cd)
-}

-- | Prepends the cache directory used by the library to the supplied file path.
-- It tries to use the path under the environment variable `$HIE_BIOS_CACHE_DIR`
-- and falls back to the standard `$XDG_CACHE_HOME/hie-bios` if the former is not set
getCacheDir :: FilePath -> IO FilePath
getCacheDir fp = do
  mbEnvCacheDirectory <- lookupEnv "HIE_BIOS_CACHE_DIR"
  cacheBaseDir <- maybe (getXdgDirectory XdgCache cacheDir) return
                         mbEnvCacheDirectory
  return (cacheBaseDir </> fp)

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

-- we don't want to generate object code so we compile to bytecode
-- (HscInterpreted) which implies LinkInMemory
-- HscInterpreted
setLinkerOptions :: G.DynFlags -> G.DynFlags
setLinkerOptions df = Gap.setNoCode $ df {
    G.ghcLink = G.LinkInMemory
  , G.ghcMode = G.CompManager
  }

setIgnoreInterfacePragmas :: G.DynFlags -> G.DynFlags
setIgnoreInterfacePragmas df = Gap.gopt_set df G.Opt_IgnoreInterfacePragmas

setVerbosity :: Int -> G.DynFlags -> G.DynFlags
setVerbosity n df = df { G.verbosity = n }

writeInterfaceFiles :: Maybe FilePath -> G.DynFlags -> G.DynFlags
writeInterfaceFiles Nothing df = df
writeInterfaceFiles (Just hi_dir) df = setHiDir hi_dir (Gap.gopt_set df G.Opt_WriteInterface)

setHiDir :: FilePath -> G.DynFlags -> G.DynFlags
setHiDir f d = d { G.hiDir      = Just f}

extractUnits :: [String] -> ([String], [String])
extractUnits = go [] []
  where
    -- TODO: we should likely use the 'processCmdLineP' instead
    go units rest ("-unit" : x : xs) = go (x : units) rest xs
    go units rest (x : xs)           = go units (x : rest) xs
    go units rest []                 = (reverse units, reverse rest)


-- | Interpret and set the specific command line options.
-- A lot of this code is just copied from ghc/Main.hs
-- It would be good to move this code into a library module so we can just use it
-- rather than copy it.
addCmdOpts :: (GhcMonad m)
           => [String] -> G.DynFlags -> m (G.DynFlags, [G.Target])
addCmdOpts cmdOpts df1 = do
  logger <- Gap.getLogger <$> G.getSession
  (df2, leftovers', _warns) <- Gap.parseDynamicFlags logger df1 (map G.noLoc cmdOpts)
  -- parse targets from ghci-scripts. Only extract targets that have been ":add"'ed.
  additionalTargets <- concat <$> mapM (liftIO . getTargetsFromGhciScript) (G.ghciScripts df2)

  -- leftovers contains all Targets from the command line
  let leftovers = map G.unLoc leftovers' ++ additionalTargets

  let (df3, srcs, _objs) = Gap.parseTargetFiles df2 leftovers
  ts <- mapM (\(f, phase) -> Gap.guessTarget f (Just $ Gap.homeUnitId_ df3) phase) srcs
  return (df3, ts)

-- | Make filepaths in the given 'DynFlags' absolute.
-- This makes the 'DynFlags' independent of the current working directory.
makeDynFlagsAbsolute :: FilePath -> G.DynFlags -> G.DynFlags
makeDynFlagsAbsolute root df =
  Gap.mapOverIncludePaths makeAbs
  $ df
    { G.importPaths = map makeAbs (G.importPaths df)
    , G.packageDBFlags =
        map (Gap.overPkgDbRef makeAbsOs) (G.packageDBFlags df)
    , G.workingDirectory = (root </>) <$> G.workingDirectory df
    }
  where
    makeAbs =
      case G.workingDirectory df of
        Just fp -> ((root </> fp) </>)
        Nothing ->
          (root </>)

    makeAbsOs p =
      case G.workingDirectory df of
        Just fp -> Gap.unsafeEncodeUtf (root </> fp) OsPath.</> p
        Nothing ->
          Gap.unsafeEncodeUtf root OsPath.</> p

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

disableOptimisation :: G.DynFlags -> G.DynFlags
disableOptimisation df = Gap.updOptLevel 0 df

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

-- | Read a ghci script and extract all targets to load form it.
-- The ghci script is expected to have the following format:
-- @
--  :add Foo Bar Main.hs
-- @
--
-- We strip away ":add" and parse the Targets.
getTargetsFromGhciScript :: FilePath -> IO [String]
getTargetsFromGhciScript script = do
  contents <- lines <$> readFile script
  let parseGhciLine = concatMap fst . filter (null . snd) . readP_to_S parser
  return $ concatMap parseGhciLine contents

-- |This parser aims to parse targets and double-quoted filepaths that are separated by spaces
-- and prefixed with the literal ":add"
--
-- >>> filter (null . snd) $ readP_to_S parser ":add Lib Lib2"
-- [(["Lib","Lib2"],"")]
--
-- >>> filter (null . snd) $ readP_to_S parser ":add Lib Lib2 \"Test Example.hs\""
-- [(["Lib","Lib2","Test Example.hs"],"")]
--
-- >>> filter (null . snd) $ readP_to_S parser ":add Lib Lib2 \"Test Exa\\\"mple.hs\""
-- [(["Lib","Lib2","Test Exa\"mple.hs"],"")]
parser :: ReadP [String]
parser = do
  _ <- string ":add" <* space1
  scriptword `sepBy` space1

space1 :: ReadP [Char]
space1 = many1 (char ' ')

scriptword :: ReadP String
scriptword = quoted <++ value

-- | A balanced double-quoted string
quoted :: ReadP String
quoted = do
    _ <- char '"'
    manyTill (escaped '"' <|> anyToken) $ char '"'

escaped :: Char -> ReadP Char
escaped c = c <$ string ("\\" <> [c])

value :: ReadP String
value = many1 (satisfy (not . isSpace))

anyToken :: ReadP Char
anyToken = satisfy $ const True

-- Used for clipping the trailing newlines on GHC output
-- Also only take the last line of output
-- (Stack's ghc output has a lot of preceding noise from 7zip etc)
trim :: String -> String
trim s = case lines s of
  [] -> s
  ls -> dropWhileEnd isSpace $ last ls