module Main where

import Control.Monad
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import System.Environment
import System.Exit

import Compiler
import CompilerMacros
import Intermediate
import Liveness
import Optimiser
import Parser
import VM


usage :: IO ()
usage = do
    progname <- getProgName
    putStrLn $ "Usage: " ++ progname ++ " [-h] [-ast] [-ir] [filename.lisp]"
    putStrLn $ "When no filename is given, will read from stdin."
    putStrLn $ "  -h     Show this help"
    putStrLn $ "  -ast   Print AST after compiler macro's"
    putStrLn $ "  -ir    Print IR after optimisation"

data Options = Options { optAST :: Bool, optIR :: Bool }

parseOptions' :: (Options, Maybe FilePath) -> [String] -> IO (Options, Maybe FilePath)
parseOptions' pair [] = return pair
parseOptions' _ ("-h":_) = usage >> exitSuccess
parseOptions' (opt, fp) ("-ast":as) = parseOptions' (opt { optAST = True }, fp) as
parseOptions' (opt, fp) ("-ir":as) = parseOptions' (opt { optIR = True }, fp) as
parseOptions' _ (('-':a):_) = putStrLn ("Unknown option '" ++ a ++ "'") >> usage >> exitFailure
parseOptions' (opt, Nothing) (f:as) = parseOptions' (opt, Just f) as
parseOptions' (_, Just _) (_:_) = putStrLn "At most one filename argument expected" >> usage >> exitFailure

parseOptions :: [String] -> IO (Options, Maybe FilePath)
parseOptions = parseOptions' (Options False False, Nothing)

liveness :: IRProgram -> Map.Map Int [Set.Set Int]
liveness (IRProgram bbs _ _) =
    let sets = livenessAnalysis bbs bidOf itemsOf outEdges fread fwrite
    in Map.fromList (zip (map bidOf bbs) sets)
  where
    itemsOf (BB _ inss term) = map Right inss ++ [Left term]
    fread (Right (_, IAssign r)) = collect [r]
    fread (Right (_, IParam _)) = []
    fread (Right (_, IClosure _)) = []
    fread (Right (_, IData _)) = []
    fread (Right (_, ICallC r rs)) = collect (r : rs)
    fread (Right (_, IAllocClo _ rs)) = collect rs
    fread (Right (_, IDiscard r)) = collect [r]
    fread (Left (IBr r _ _)) = collect [r]
    fread (Left (IJmp _)) = []
    fread (Left (IRet r)) = collect [r]
    fread (Left (ITailC r rs)) = collect (r : rs)
    fread (Left IExit) = []
    fread (Left IUnknown) = []
    fwrite (Right (r, _)) = collect [r]
    fwrite (Left _) = []
    collect rs = [i | RTemp i <- rs]

main :: IO ()
main = do
    (opts, mfname) <- getArgs >>= parseOptions
    source <- maybe getContents readFile mfname

    prog <- parseProgram mfname source >>= either (die . show) return
    let prog' = compilerMacros prog
    when (optAST opts) $ print prog'
    irprog <- either die return (compileProgram prog')
    let opt = optimise irprog
    when (optIR opts) $ print opt
    print (liveness opt)
    vmRun opt