ghc-core-0.1: ghc-core.hs

{-# LANGUAGE PatternGuards #-}
--------------------------------------------------------------------
-- |
-- Module    : ghc-core
-- Copyright : (c) Galois, Inc. 2008
-- License   : BSD3
--
-- Maintainer: Don Stewart <dons@galois.com>
-- Stability : provisional
-- Portability: Needs a few libraries from hackage.
--
--------------------------------------------------------------------
--
-- Inspect the optimised core and assembly produce by GHC.
--
-- Examples:
--
-- > ghc-core zipwith
--
-- > ghc-core -fvia-C zipwith 
--

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

import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.List
import System.Cmd
import System.Directory
import System.Environment
import System.Exit
import System.FilePath
import System.IO
import System.Process
import Text.Printf
import qualified Control.Exception as C

import Text.Regex.PCRE.Light.Char8

import Language.Haskell.HsColour
import Language.Haskell.HsColour.Colourise

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

-- default hscolours
defaultColourPrefs :: ColourPrefs
defaultColourPrefs = ColourPrefs
  { keyword  = [Foreground Green,Underscore]
  , keyglyph = [Foreground Red]
  , layout   = [Foreground Cyan]
  , comment  = [Foreground Blue]
  , conid    = [Normal]
  , varid    = [Normal]
  , conop    = [Foreground Red,Bold]
  , varop    = [Foreground Cyan]
  , string   = [Foreground Magenta]
  , char     = [Foreground Magenta]
  , number   = [Foreground Magenta]
  , cpp      = [Foreground Magenta,Dim]
  , selection = [Bold, Foreground Magenta]
  , variantselection = [Dim, Foreground Red, Underscore]
  , definition = [Foreground Blue]
  }

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

main :: IO ()
main = do
    args  <- getArgs

    mv <- getEnvMaybe "PAGER"
    let less = case mv of Just s -> s ; _ -> "less"

    strs' <- polish `fmap` compileWithCore args
    (strs,tmps) <- do
        x <- readProcess "sh" ["-c","ls /tmp/ghc*/*.s | head -1"] []
        case x of
            Left _ -> return (strs', Nothing)
            Right s -> if "-fvia-C" `elem` args
                            then do asm <- readFile (init s)
                                    return ((strs' ++ asm), Just $ takeDirectory s)
                            else return (strs', Just $ takeDirectory s)

    let nice = hscolour TTY defaultColourPrefs False True False strs

    bracket
        (openTempFile "/tmp" "ghc-core-XXXX.hcr")
        (\(f,h) -> do hClose h
                      removeFile f
                      case tmps of
                        Just g -> system ("rm -rf " ++ g) >> return ()
                        _      -> return ()
        )
        (\(f,h) -> do
            hPutStrLn h nice >> hFlush h
            e <- system $ less ++ " -r " ++ f
            exitWith e)

--  
-- Clean up the output with some regular expressions.  
--  
polish :: String -> String
polish = unlines . dups . map polish' . lines
    where
        polish' [] = []
        polish' s
            | Just [_,a,b] <- match name  s [] =  polish' (a ++ b)
            | Just [_,a,b] <- match local s [] =  polish' (a ++ b)
            | Just _ <- match core  s [] =  "------------------------------- Core -----------------------------------"
            | Just _ <- match asm   s [] =  "------------------------------- Assembly -------------------------------"

            | Just _       <- match junk  s [] =  []
            | Just _       <- match junk2 s [] =  []
            | Just _       <- match junk3 s [] =  []
            | Just _       <- match junk4 s [] =  []
            | Just _       <- match junk5 s [] =  []
            | Just _       <- match junk6 s [] =  []
            | Just _       <- match junk7 s [] =  []

            | otherwise = s

        -- simplify some qualified names
        name  = compile
                 "^(.*)GHC\\.[^\\.]*\\.(.*)$"
                 [ungreedy]

        local = compile
                 "^(.*)Main\\.(.*)$"
                 [ungreedy]

        -- remove boring things
        core  = compile "Tidy Core" [ungreedy]
        asm   = compile "Asm code" [ungreedy]

        junk  = compile "^.GlobalId" []
        junk2 = compile "^.Arity .*" []
        junk3 = compile "^Rec {|^end Rec" []
        junk4 = compile "DmdType" []
        junk5 = compile "NoCafRefs" []
        junk6 = compile "^\\[\\]$" []
        junk7 = compile "==========" []

        -- remove duplicate blank lines
        dups []         = []
        dups ([]:[]:xs) = dups ([]:xs)
        dups (x:xs) = x : dups xs

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

compileWithCore :: [String] -> IO String
compileWithCore opts = do
    let ghc  = "ghc"
        args = words "-O2 -keep-tmp-files -ddump-simpl -ddump-asm -ddump-simpl-stats -no-recomp --make"

    x <- readProcess ghc (opts ++ args) []
    case x of
         Left (err,str,std) -> do
            mapM_ putStrLn (lines str)
            mapM_ putStrLn (lines std)
            printf "GHC failed to compile %s\n" (show err)
            exitWith (ExitFailure 1) -- fatal

         Right str      -> return str

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

--
-- Strict process reading 
--
readProcess :: FilePath                              -- ^ command to run
            -> [String]                              -- ^ any arguments
            -> String                                -- ^ standard input
            -> IO (Either (ExitCode,String,String) String)  -- ^ either the stdout, or an exitcode and any output

readProcess cmd args input = C.handle (return . handler) $ do
    (inh,outh,errh,pid) <- runInteractiveProcess cmd args Nothing Nothing

    output  <- hGetContents outh
    outMVar <- newEmptyMVar
    forkIO $ (C.evaluate (length output) >> putMVar outMVar ())

    errput  <- hGetContents errh
    errMVar <- newEmptyMVar
    forkIO $ (C.evaluate (length errput) >> putMVar errMVar ())

    when (not (null input)) $ hPutStr inh input
    takeMVar outMVar
    takeMVar errMVar
    ex     <- C.catch (waitForProcess pid) (\_e -> return ExitSuccess)
    hClose outh
    hClose inh          -- done with stdin
    hClose errh         -- ignore stderr

    return $ case ex of
        ExitSuccess   -> Right output
        ExitFailure _ -> Left (ex, errput, output)

  where
    handler (C.ExitException e) = Left (e,"","")
    handler e                   = Left (ExitFailure 1, show e, "")


-- Safe wrapper for getEnv
getEnvMaybe :: String -> IO (Maybe String)
getEnvMaybe name = handle (const $ return Nothing) (Just `fmap` getEnv name)