Util.hs

module IHaskell.Eval.Util (
  -- * Initialization
  initGhci,

  -- * Flags and extensions
  -- ** Set and unset flags.
  extensionFlag, setExtension,
  ExtFlag(..),
  setFlags,

  -- * Code Evaluation
  evalImport,
  evalDeclarations,
  getType,

  -- * Pretty printing
  doc,
  ) where

-- GHC imports.
import DynFlags
import FastString
import GHC
import GhcMonad
import HsImpExp
import HscTypes
import InteractiveEval
import Module
import Outputable
import Packages
import RdrName
import qualified Pretty

import Control.Monad (void)
import Data.Function (on)
import Data.List (find)
import Data.String.Utils (replace)

-- | A extension flag that can be set or unset.
data ExtFlag
     = SetFlag   ExtensionFlag
     | UnsetFlag ExtensionFlag

-- | Find the extension that corresponds to a given flag. Create the
-- corresponding 'ExtFlag' via @SetFlag@ or @UnsetFlag@.
-- If no such extension exist, yield @Nothing@.
extensionFlag :: String         -- Extension name, such as @"DataKinds"@
              -> Maybe ExtFlag
extensionFlag ext = 
  case find (flagMatches ext) xFlags of
    Just (_, flag, _) -> Just $ SetFlag flag
    -- If it doesn't match an extension name, try matching against
    -- disabling an extension.
    Nothing ->
      case find (flagMatchesNo ext) xFlags of
        Just (_, flag, _) -> Just $ UnsetFlag flag
        Nothing -> Nothing

  where
    -- Check if a FlagSpec matches an extension name.
    flagMatches ext (name, _, _) = ext == name

    -- Check if a FlagSpec matches "No<ExtensionName>".
    -- In that case, we disable the extension.
    flagMatchesNo ext (name, _, _) = ext == "No"  ++ name

-- | Set an extension and update flags.
-- Return @Nothing@ on success. On failure, return an error message.
setExtension :: GhcMonad m => String -> m (Maybe String)
setExtension ext = do
  flags <- getSessionDynFlags
  case extensionFlag ext of
    Nothing -> return $ Just $ "Could not parse extension name: " ++ ext
    Just flag -> do
      setSessionDynFlags $ 
        case flag of
          SetFlag ghcFlag -> xopt_set flags ghcFlag
          UnsetFlag ghcFlag -> xopt_unset flags ghcFlag
      return Nothing

-- | Set a list of flags, as per GHCi's `:set`.
-- This was adapted from GHC's InteractiveUI.hs (newDynFlags).
-- It returns a list of error messages.
setFlags :: GhcMonad m => [String] -> m [String]
setFlags ext = do
    -- Try to parse flags.
    flags <- getSessionDynFlags
    (flags', unrecognized, warnings) <- parseDynamicFlags flags (map noLoc ext)

    -- First, try to check if this flag matches any extension name.
    let restorePkg x = x { packageFlags = packageFlags flags }
    let restoredPkgs = flags' { packageFlags = packageFlags flags}
    GHC.setProgramDynFlags restoredPkgs
    GHC.setInteractiveDynFlags restoredPkgs

    -- Create the parse errors.
    let noParseErrs = map (("Could not parse: " ++) . unLoc) unrecognized
        allWarns = map unLoc warnings ++ 
                     ["-package not supported yet" | packageFlags flags /= packageFlags flags']
        warnErrs    = map ("Warning: " ++) allWarns
    return $ noParseErrs ++ warnErrs

-- | Convert an 'SDoc' into a string. This is similar to the family of
-- 'showSDoc' functions, but does not impose an arbitrary width limit on
-- the output (in terms of number of columns). Instead, it respsects the
-- 'pprCols' field in the structure returned by 'getSessionDynFlags', and
-- thus gives a configurable width of output.
doc :: GhcMonad m => SDoc -> m String
doc sdoc = do
  flags <- getSessionDynFlags
  unqual <- getPrintUnqual
  let style = mkUserStyle unqual AllTheWay
  let cols = pprCols flags
      d = runSDoc sdoc (initSDocContext flags style)
  return $ Pretty.fullRender Pretty.PageMode cols 1.5 string_txt "" d
  where
    string_txt :: Pretty.TextDetails -> String -> String
    string_txt (Pretty.Chr c)   s  = c:s
    string_txt (Pretty.Str s1)  s2 = s1 ++ s2
    string_txt (Pretty.PStr s1) s2 = unpackFS s1 ++ s2
    string_txt (Pretty.LStr s1 _) s2 = unpackLitString s1 ++ s2

-- | Initialize the GHC API. Run this as the first thing in the `runGhc`.
-- This initializes some dyn flags (@ExtendedDefaultRules@,
-- @NoMonomorphismRestriction@), sets the target to interpreted, link in
-- memory, sets a reasonable output width, and potentially a few other
-- things. It should be invoked before other functions from this module.
initGhci :: GhcMonad m => m ()
initGhci = do
  -- Initialize dyn flags.
  -- Start with -XExtendedDefaultRules and -XNoMonomorphismRestriction.
  originalFlags <- getSessionDynFlags
  let flag = flip xopt_set
      unflag = flip xopt_unset
      dflags = flag Opt_ExtendedDefaultRules . unflag Opt_MonomorphismRestriction $ originalFlags

  void $ setSessionDynFlags $ dflags { hscTarget = HscInterpreted,
                                       ghcLink = LinkInMemory,
                                       pprCols = 300 }

-- | Evaluate a single import statement.
-- If this import statement is importing a module which was previously
-- imported implicitly (such as `Prelude`) or if this module has a `hiding`
-- annotation, the previous import is removed.
evalImport :: GhcMonad m => String -> m ()
evalImport imports = do
  importDecl <- parseImportDecl imports
  context <- getContext

  -- If we've imported this implicitly, remove the old import.
  let noImplicit = filter (not . implicitImportOf importDecl) context

      -- If this is a `hiding` import, remove previous non-`hiding` imports.
      oldImps = if isHiddenImport importDecl
                then filter (not . importOf importDecl) context
                else noImplicit

  -- Replace the context.
  setContext $ IIDecl importDecl : oldImps

  where
    -- Check whether an import is the same as another import (same module).
    importOf :: ImportDecl RdrName -> InteractiveImport -> Bool
    importOf _ (IIModule _) = False
    importOf imp (IIDecl decl) = ((==) `on` (unLoc . ideclName)) decl imp

    -- Check whether an import is an *implicit* import of something.
    implicitImportOf :: ImportDecl RdrName -> InteractiveImport -> Bool
    implicitImportOf _ (IIModule _) = False
    implicitImportOf imp (IIDecl decl) = ideclImplicit decl && imp `importOf` IIDecl decl

    -- Check whether an import is hidden.
    isHiddenImport :: ImportDecl RdrName -> Bool
    isHiddenImport imp = case ideclHiding imp of
                           Just (True, _) -> True
                           _ -> False

-- | Evaluate a series of declarations.
-- Return all names which were bound by these declarations.
evalDeclarations :: GhcMonad m => String -> m [String]
evalDeclarations decl = do
  names <- runDecls decl
  flags <- getSessionDynFlags
  return $ map (replace ":Interactive." "" . showPpr flags) names

-- | Get the type of an expression and convert it to a string.
getType :: GhcMonad m => String -> m String
getType expr = do
  result <- exprType expr
  flags <- getSessionDynFlags
  let typeStr = showSDocUnqual flags $ ppr result
  return typeStr