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)

-- TODO: Note about stackification. Temporaries need to be pushed before
-- a call if they're live after it and they could, conceivably, be wrongly
-- mutated by the called function otherwise. This has a couple of
-- interesting consequences:
-- 1. No temporaries ever need to be pushed from the global context, or
--    "main function", since they can never be wrongly mutated: that
--    requires re-entering the function in which they were defined, but one
--    can never re-enter the "main function".
-- 2. No temporaries ever need to be pushed before a tail call; since no
--    local variables are live after it (obviously). (Global variables are
--    covered by point (1.).)
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