path: root/Lower.hs

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

import AST (Name)
import Control.Monad.State.Strict
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Intermediate


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

data AsmInstr'
    = Li ARef Int
    | Mv ARef ARef
    | Arith Arith ARef ARef ARef
    | Not ARef ARef
    | Call Label
    | Jcc CCond ARef Label
    | JccR CCond ARef ARef
    -- Load/Store: 'Int' is number of BYTES.
    | Load Int ARef ARef
    | Store Int ARef ARef
    -- Allocate n r: Allocate n bytes of memory, pointer in r
    | Alloc Int ARef
    -- Memcpy d s n: Copy n bytes from s to d
    | Memcpy ARef ARef ARef
    | SysExit
  deriving (Show)

data ARef
    = ASysReg Int  -- System register
    | AReg Int  -- Virtual register
    | ALabel Name  -- Function pointer for the label
    | ASClo Name  -- Static closure object of function
    | ANone
  deriving (Show)

newtype Label = Label String
  deriving (Show)

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

data CCond = CCZ | CCNZ
  deriving (Show)

type GFDMap = Map.Map Name GlobFuncDef

type BId = Int


data LowerState = LowerState
    { lsNextId :: Int }
  deriving (Show)

newtype LM a = LM {unLM :: State LowerState a}
  deriving (Functor, Applicative, Monad, MonadState LowerState)

initState :: IRProgram -> LowerState
initState (IRProgram bbs _ _) = LowerState (maxTemp + 1)
  where
    maxTemp = maximum (map maxTempB bbs)
    maxTempB (BB _ inss term) = max (maximum (map maxTempI inss)) (maxTempT term)
    maxTempI (r, ins) = max (maxTempR r) (maxTempI' ins)
    maxTempI' ins = case ins of
        IAssign r      -> maxTempR r
        IParam _       -> 0
        IClosure _     -> 0
        IData _        -> 0
        ICallC r rs    -> max (maxTempR r) (maximum (map maxTempR rs))
        IAllocClo _ rs -> maximum (map maxTempR rs)
        IDiscard r     -> maxTempR r
    maxTempT term = case term of
        IBr r _ _ -> maxTempR r
        IJmp _    -> 0
        IRet r    -> maxTempR r
        IExit     -> 0
        IUnknown  -> 0
    maxTempR r = case r of
        RConst _ -> 0
        RTemp n  -> n
        RSClo _  -> 0
        RNone    -> 0

runLM :: IRProgram -> LM a -> a
runLM initIR act = evalState (unLM act) (initState initIR)

genId :: LM Int
genId = state $ \s -> (lsNextId s, s {lsNextId = lsNextId s + 1})

genTemp :: LM ARef
genTemp = liftM AReg genId



-- Calling convention:
-- Upon function entry, the stack should look as follows:
-- - Closure item N-1
-- ...
-- - Closure item 1
-- - Closure item 0
-- - Argument M-1
-- ...
-- - Argument 1
-- - Argument 0
-- - Link register [pushed by callee]
-- Thus, the stack pointer is at the link entry.
--
-- A closure object is laid out as follows:
-- - Function pointer
-- - Number of closure items (N)
-- - Closure item 0
-- - Closure item 1
-- ...
-- - Closure item N-1


lowerIR :: IRProgram -> AsmProgram'
lowerIR origProgram =
    -- TODO: For each function, push and pop the link register
    let res1 = closuresAreParams origProgram
        IRProgram bbs gfds _ = res1

        blocks = sequence [(lab,) <$> lowerBB bb gfds
                          | bb@(BB bid _ _) <- bbs
                          , let lab = Label ("BB" ++ show bid)]
    in AsmProgram' (runLM origProgram blocks)

lowerBB :: BB -> GFDMap -> LM [AsmInstr']
lowerBB (BB _bid inss term) gfds = do
    res1 <- concat <$> mapM (\ins -> lowerIns ins gfds) inss
    res2 <- lowerTerm term
    return (res1 ++ res2)

lowerIns :: Instruction -> GFDMap -> LM [AsmInstr']
lowerIns (dest, instruction) gfds = case instruction of
    IAssign src -> do
        (inss, srcr) <- toARef src
        return $ inss ++ [Mv (toARefStore dest) srcr]

    IParam idx -> do
        offsetr <- genTemp
        return [Li offsetr (8 * (idx + 1))
               ,Arith Add offsetr offsetr (ASysReg 15)
               ,Load 8 (toARefStore dest) offsetr]

    IClosure _ -> error "Unexpected IClosure, why didn't closuresAreParams run?"

    IData _ -> error "Unsupported IData in Lower"

    ICallC closure args | Just act <- lowerBuiltin closure args -> act
                        | otherwise -> do
        (inss1, closurer) <- toARef closure
        (inss2, argrs) <- toARefs args
        eightr <- genTemp
        numitemsr <- genTemp
        numbytesr <- genTemp
        stackshiftr <- genTemp
        walkr <- genTemp
        funptrr <- genTemp
        return $
            inss1 ++ inss2 ++
            [Li eightr 8
            -- Copy the closure parameters
            ,Arith Add walkr closurer eightr
            ,Load 8 numitemsr walkr
            ,Arith Add walkr walkr eightr
            ,Arith Mul numbytesr numitemsr eightr
            ,Arith Sub (ASysReg 15) (ASysReg 15) numbytesr
            ,Memcpy (ASysReg 15) walkr numbytesr] ++
            -- Copy the function parameters
            concat [[Arith Sub (ASysReg 15) (ASysReg 15) eightr
                    ,Store 8 (ASysReg 15) ref]
                   | ref <- reverse argrs] ++
            [Load 8 funptrr closurer
            -- Do the call; don't separate these two instructions!
            ,Arith Add (ASysReg 14) (ASysReg 0) eightr
            ,JccR CCNZ (ASysReg 0) funptrr
            -- Clean up the stack
            ,Li stackshiftr (8 * length argrs)
            ,Arith Add stackshiftr stackshiftr numbytesr
            ,Arith Add (ASysReg 15) (ASysReg 15) stackshiftr
            ,Mv (toARefStore dest) (ASysReg 13)]

    IAllocClo name refs ->
        let GlobFuncDef _startBId _nargs closureSlots = gfds Map.! name
        in if | length refs /= length closureSlots ->
                    error $ "INTERNAL: Call to function '" ++ name ++ "' with " ++
                            show (length refs) ++ " closure params, while it expects " ++
                            show (length closureSlots)
              | otherwise -> do
                    let destr = toARefStore dest
                        GlobFuncDef initBId _ _ = gfds Map.! name
                    itemr <- genTemp
                    numitemsr <- genTemp
                    eightr <- genTemp
                    (inss, refs') <- toARefs refs
                    return $
                        inss ++
                        [Alloc (16 + 8 * length refs) destr
                        ,Store 8 destr (ALabel ("BB" ++ show initBId))
                        ,Li eightr 8
                        ,Arith Add itemr destr eightr
                        ,Li numitemsr (length refs')
                        ,Store 8 itemr numitemsr] ++
                        concat [[Arith Add itemr itemr eightr, Store 8 itemr ref]
                               | ref <- refs']

    IDiscard _  -> return []

lowerTerm :: Terminator -> LM [AsmInstr']
lowerTerm terminator = case terminator of
    IBr cond tg1 tg2 -> do
        (inss, condr) <- toARef cond
        return $
            inss ++
            [Jcc CCNZ condr (Label ("BB" ++ show tg1))
            ,Jcc CCNZ (ASysReg 0) (Label ("BB" ++ show tg2))]

    IJmp bid ->
        return [Jcc CCNZ (ASysReg 0) (Label ("BB" ++ show bid))]

    IRet val -> do
        (inss, valr) <- toARef val
        ptrr <- genTemp
        eightr <- genTemp
        return $
            inss ++
            [Mv (ASysReg 13) valr
            ,Load 8 ptrr (ASysReg 15)
            ,Li eightr 8
            ,Arith Add (ASysReg 15) (ASysReg 15) eightr
            ,JccR CCNZ (ASysReg 0) ptrr]

    IExit ->
        return [SysExit]

    IUnknown ->
        error "Unexpected IUnknown"

-- Currently useless because Compiler puts the RSClo in a temporary before
-- calling it.
lowerBuiltin :: Ref -> [Ref] -> Maybe (LM [AsmInstr'])
lowerBuiltin (RSClo "+") [arg1, arg2] = Just $ do
    (inss1, r1) <- toARef arg1
    (inss2, r2) <- toARef arg2
    return $ inss1 ++ inss2 ++ [Arith Add (ASysReg 13) r1 r2]
lowerBuiltin _ _ = Nothing

toARef :: Ref -> LM ([AsmInstr'], ARef)
toARef (RConst val) = genTemp >>= \r -> return ([Li r val], r)
toARef ref = return ([], toARefStore ref)

toARefs :: [Ref] -> LM ([AsmInstr'], [ARef])
toARefs refs =
    foldr (\(i, r) (inss, rs) -> (i ++ inss, r : rs)) ([], [])
        <$> mapM toARef refs

toARefStore :: Ref -> ARef
toARefStore (RConst _) = error "Cannot store to a constant immediate value"
toARefStore (RTemp r) = AReg r
toARefStore (RSClo name) = ASClo name
toARefStore RNone = ANone

closuresAreParams :: IRProgram -> IRProgram
closuresAreParams (IRProgram bbs gfds datatbl) =
    IRProgram (map convertBB bbs) gfds datatbl
  where
    argccount = floodArgCountAll gfds bbs

    convertBB (BB bid inss term) = BB bid (map (convert bid) inss) term

    convert bid (dest, IClosure n) = (dest, IParam (n + argccount Map.! bid))
    convert _ i = i

floodArgCountAll :: GFDMap -> [BB] -> Map.Map BId Int
floodArgCountAll gfds bbs = Map.unions [floodArgCount gfd bbs | gfd <- Map.elems gfds]

floodArgCount :: GlobFuncDef -> [BB] -> Map.Map BId Int
floodArgCount (GlobFuncDef initBId nargs _) bbs =
    Map.fromList [(bid, nargs) | bid <- Set.toList (reachable initBId)]
  where
    bbMap :: Map.Map BId BB
    bbMap = Map.fromList [(bid, bb) | bb@(BB bid _ _) <- bbs]

    reachable :: BId -> Set.Set BId
    reachable bid = go Set.empty [bid]
      where
        go :: Set.Set BId -> [BId] -> Set.Set BId
        go accum [] = accum
        go accum front =
            let front' = Set.fromList (concatMap nexts front) Set.\\ accum
            in go (Set.union accum (Set.fromList front)) (Set.toList front')

    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  -> []