path: root/Lower.hs

{-# LANGUAGE TupleSections, MultiWayIf, GeneralizedNewtypeDeriving #-}
module Lower(lowerIR,
             AsmProgram'(..), AsmInstr'(..), ARef(..), Immediate(..), Label(..), CCond(..)) where

import AST (Name)
import Control.Monad.State.Strict
import Data.List
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 Immediate
    | 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

    -- Macro's that will be instantiated with code to set up and clean up
    -- the register spill space. Will use R14 as temporary.
    | SetupSpill
    | CleanupSpill

    -- Semantics: CallR rFun r8 =^= add R14, R0, r8; jnzr R0, rFun
    -- r8 should contain the integer 8, rFun the function pointer
    | CallR ARef ARef

    -- Memcpy8 rDEST rSRC rNUM rTEMP
    -- Destroys all four registers, and requires that they have a PreventSpill annotation
    --
    --     jz rNUM, Lfinish
    -- Lloop:
    --     l64 rTEMP, rSRC
    --     s64 rDEST, rTEMP
    --     li rTEMP, 8
    --     add rDEST, rDEST, rTEMP
    --     add rSRC, rSRC, rTEMP
    --     li rTEMP, 1
    --     sub rNUM, rNUM, rTEMP
    --     jnz rNUM, Lloop
    -- Lfinish:
    | Memcpy8 ARef ARef ARef ARef

    -- Marks the virtual register as non-spillable. Should be an AReg.
    | PreventSpill ARef
  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)

data Immediate
    = Imm Int
    | FPOffset Int
  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
        IFunctionEntry -> 0
        IApplicationEntry -> 0
    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
--
-- Register usage:
-- - R15 = stack pointer
-- - R14 = link register
-- - R13 = return register
-- - Further registers: available for allocation
-- R13 and R14 can also be used as administration-temporary if they are not
-- otherwise occupied.


lowerIR :: IRProgram -> AsmProgram'
lowerIR origProgram =
    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 _ 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 (FPOffset (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
        closurebufr <- genTemp
        stackshiftr <- genTemp
        walkr <- genTemp
        funptrr <- genTemp
        memcpyTargetr <- genTemp
        memcpySourcer <- genTemp
        memcpyNumr <- genTemp
        memcpyTempr <- genTemp

        -- mov qword rcx, [closurer + 8]
        -- shl rcx, 3
        -- sub rsp, rcx
        -- lea rsi, [closurer + 16]
        -- mov rdi, rsp
        -- rep movsq
        -- push argN
        -- ...
        -- push arg1
        -- call [closurer]
        -- add rsp, (8 * N)
        -- add rsp, [closurer + 8]
        -- mov dest, rax

        return $
            inss1 ++ inss2 ++
            [Li eightr (Imm 8)
            -- Copy the closure parameters: first make space
            ,Arith Add walkr closurer eightr
            ,Load 8 numitemsr walkr
            ,Arith Add walkr walkr eightr
            ,Arith Mul numbytesr numitemsr eightr
            ,Arith Sub closurebufr (ASysReg 15) numbytesr
            -- Now do the copy
            ,Mv memcpyTargetr closurebufr
            ,Mv memcpySourcer walkr
            ,Mv memcpyNumr numitemsr
            ,PreventSpill memcpyTargetr
            ,PreventSpill memcpySourcer
            ,PreventSpill memcpyNumr
            ,PreventSpill memcpyTempr
            ,Memcpy8 memcpyTargetr memcpySourcer memcpyNumr memcpyTempr
            -- We're using R14 as temporary stack pointer here, so that the
            -- actual R15 doesn't move just yet
            ,Mv (ASysReg 14) closurebufr] ++
            -- Copy the function parameters
            concat [[Arith Sub (ASysReg 14) (ASysReg 14) eightr
                    ,Store 8 (ASysReg 14) ref]
                   | ref <- reverse argrs] ++
            [Load 8 funptrr closurer
            -- Do the call; these registers must not spill, since the
            -- stack references resulting from retrieval from the spill
            -- area would be invalid with the shifted stack pointer
            ,PreventSpill funptrr
            ,PreventSpill eightr
            ,Mv (ASysReg 15) (ASysReg 14)
            ,CallR funptrr eightr
            -- Clean up the stack
            ,PreventSpill stackshiftr
            ,PreventSpill numbytesr
            ,Li stackshiftr (Imm (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
                    allocsizer <- genTemp
                    (inss, refs') <- toARefs refs
                    return $
                        inss ++
                        [Li eightr (Imm 8)
                        -- First call malloc to obtain an allocation
                        ,Arith Sub (ASysReg 15) (ASysReg 15) eightr
                        ,Li allocsizer (Imm (16 + 8 * length refs))
                        ,Store 8 (ASysReg 15) allocsizer
                        ,Jcc CCNZ (ASysReg 0) (Label "__builtin_malloc")
                        ,Mv destr (ASysReg 13)
                        -- Then put the right data in the closure allocation
                        ,Store 8 destr (ALabel ("BB" ++ show initBId))
                        ,Arith Add itemr destr eightr
                        ,Li numitemsr (Imm (length refs'))
                        ,Store 8 itemr numitemsr] ++
                        concat [[Arith Add itemr itemr eightr, Store 8 itemr ref]
                               | ref <- refs']

    IDiscard _  ->
        return []

    IFunctionEntry -> do
        -- We need to push the link register and set up the spill space
        -- Note that the return register isn't used at this point, so we
        -- can abuse it as an 8 register
        return
            [Li (ASysReg 13) (Imm 8)
            ,Arith Sub (ASysReg 15) (ASysReg 15) (ASysReg 13)
            ,Store 8 (ASysReg 15) (ASysReg 14)
            ,SetupSpill]

    IApplicationEntry ->
        return [SetupSpill]

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
            ,CleanupSpill
            ,Load 8 ptrr (ASysReg 15)
            ,Li eightr (Imm 8)
            ,Arith Add (ASysReg 15) (ASysReg 15) eightr
            ,JccR CCNZ (ASysReg 0) ptrr]

    IExit ->
        return [Jcc CCNZ (ASysReg 0) (Label "__builtin_exit")]

    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 (Imm 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 = floodArgCount 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

floodArgCount :: GFDMap -> [BB] -> Map.Map BId Int
floodArgCount gfds bbs =
    Map.fromList [(bid, nargs)
                 | (name, bids) <- Map.assocs (snd $ funcBBPartition gfds bbs)
                 , bid <- bids
                 , let GlobFuncDef _ nargs _ = gfds Map.! name]

funcBBPartition :: GFDMap -> [BB] -> ([BId], Map.Map Name [BId])
funcBBPartition gfds bbs =
    let regions = toplevelCode : map snd spreads
    in if length (concat regions) /= length (nub $ concat regions)
        then error "Basic blocks belong to multiple functions in Lower"
        else (toplevelCode, Map.fromList spreads)
  where
    bbMap :: Map.Map BId BB
    bbMap = Map.fromList [(bid, bb) | bb@(BB bid _ _) <- bbs]

    toplevelCode :: [BId]
    toplevelCode = Set.toList (reachable 0)

    spreads :: [(Name, [BId])]
    spreads = [(name, Set.toList (reachable initBId))
              | (name, GlobFuncDef initBId _ _) <- Map.assocs gfds]

    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 termOf (bbMap Map.! bid) of
        IBr _ a b -> [a, b]
        IJmp a    -> [a]
        IRet _    -> []
        IExit     -> []
        IUnknown  -> []