{-# LANGUAGE TupleSections #-}
module Lower(lowerIR) where

import AST (Name)
import Data.List
import qualified Data.Map.Strict as Map
import Intermediate


-- Not yet regalloc'd
data AsmProgram' = AsmProgram' [(Label, [AsmInstr'])]

data AsmInstr'
    = Li Ref Int
    | Mv Ref Ref
    | Arith Arith Ref Ref Ref
    | Not Ref Ref
    | Call Ref Label
    | Jcc CCond Ref Label
    | JccR CCond Ref Ref
    | Load Int Ref Ref
    | Store Int Ref Ref

newtype Label = Label String

data Arith = Add | Sub | Mul | Div | Lt | Lte | And | Or | Xor | Sll | Slr | Sar

data CCond = CCZ | CCNZ

type GFDMap = Map.Map Name GlobFuncDef

type BId = Int


-- Calling convention:
-- Upon function entry, the stack should look as follows:
-- - Closure item 1
-- - Closure item 2
-- ...
-- - Closure item N
-- - Argument 1
-- - Arugment 2
-- ...
-- - Argument M
-- - Link register [pushed by callee]


lowerIR :: IRProgram -> AsmProgram'
lowerIR (IRProgram bbs gfds datatbl) =
    let argcmap = floodArgCount gfds bbs
    in AsmProgram' [(Label ("BB" ++ show bid), lowerBB bb gfds argcmap)
                   | bb@(BB bid _ _) <- bbs]

lowerBB :: BB -> GFDMap -> Map.Map BId Int -> [AsmInstr']
lowerBB (BB bid inss term) gfds argcmap = concatMap (\ins -> lowerIns ins gfds argcmap) inss

lowerIns :: Instruction -> GFDMap -> Map.Map BId Int -> [AsmInstr']
lowerIns (dest, instruction) gfds argcmap = case instruction of
    IAssign src -> [Mv dest src]
    IParam idx  ->  undefined

floodArgCount :: GFDMap -> [BB] -> Map.Map BId Int
floodArgCount gfds bbs = go Map.empty [(bid, n) | GlobFuncDef bid n _ <- Map.elems gfds]
  where
    bbMap :: Map.Map Int BB
    bbMap = Map.fromList [(bid, bb) | bb@(BB bid _ _) <- bbs]

    go :: Map.Map BId Int -> [(Int, Int)] -> Map.Map BId Int
    go result frontier =
        let result' = foldl' (\mp (bid, n) -> Map.insert bid n mp) result frontier
            frontier' = concat [map (,n) (nexts bid \\ Map.keys result')
                               | (bid, n) <- frontier]
        in go result' frontier'

    nexts :: BId -> [BId]
    nexts bid = case let BB _ _ term = bbMap Map.! bid in term of
        IBr _ a b -> [a, b]
        IJmp a    -> [a]
        IRet _    -> []
        IExit     -> []
        IUnknown  -> []