asic-1.2: src/Main.hs

-- | Instrumenter entry point
module Main(main) where

import System.IO
import Control.Monad
import Parser
import qualified PrettyCode as PC
import qualified PrettyTree as PT
import Options
import qualified Data.ByteString.Lazy as B
import ProgInfo
import ByteCode
import TrfInjectAbc
import TrfInjectRefl
import SymbolTables
import CFG
import PrettyUtil
import GenInstrLib
import SymInfo
import Data.Monoid
import ExtractAbc


-- | Describes different types of outputs
data Result
  = ResSwf !SwfFile
  | ResDoc !Doc
  | ResStr !B.ByteString
  | ResNothing

-- | Takes an swf file as parameter and processes it depending on the commandline options.
main :: IO ()
main = do
  opts <- commandlineArgs
  let source = optSourceFile opts
  when (source /= "") $ do
    swf <- parseAnyAsSwf source
    let tbls = symInfoSwf swf
    when (optDumpSym opts) (hPutStrLn stderr $ show tbls)
    when (optDumpAst opts) $ do
      B.hPut stderr $ renderBytes $ PT.ppSwf opts tbls swf
      hPutStrLn stderr ""  -- trailing newline

    syms <- case optEnv opts of
      [] -> return mempty
      ps -> extractSymInfos ps

    res <- process opts syms tbls swf
    let outp = case optOutputFile opts of
                 Nothing   -> B.hPut stdout
                 Just dest -> withFile dest WriteMode . flip B.hPut
    case res of
      ResSwf swf' -> outp $ PC.ppSwf swf'
      ResDoc doc  -> outp $ renderBytes doc
      ResStr str  -> outp str
      ResNothing  -> return ()

-- | Processes the swf file
process :: Options -> SymInfo -> [SymbolTables] -> SwfFile -> IO Result
process opts gEnv tbls file = case optGenLib opts of
  Just name -> return $ ResDoc $ genInstrLib name gEnv tbls
  Nothing -> pipeline file [injRefl opts tbls, injAbc opts]

-- | Sequences transformations on the swf file
pipeline :: SwfFile -> [SwfFile -> IO (Maybe SwfFile)] -> IO Result
pipeline file = fmap wrap . foldl (>>=) (return $ Left file) . map act where
  wrap = either (const ResNothing) ResSwf
  act f inp = fmap (maybe inp Right) $ f (either id id inp)

-- | Wrapper around the reflection injection
injRefl :: Options -> [SymbolTables] -> SwfFile -> IO (Maybe SwfFile)
injRefl opts tbls file
  | optInjectRefl opts = return $ Just $ injectReflSwf tbls file
  | otherwise          = return Nothing

-- | Wrapper around the abc injection
injAbc :: Options -> SwfFile -> IO (Maybe SwfFile)
injAbc opts file = case optInjectAbc opts of
    []    -> return Nothing
    paths -> do
      swfs <- mapM parseAnyAsSwf paths
      return $! Just $! foldl inj file $ concatMap extractAbcFiles swfs
  where inj = flip (injectAbc "asic-injected-code")