{-# LANGUAGE TupleSections, MultiWayIf, GeneralizedNewtypeDeriving #-} module Lower(lowerIR) 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 -- - R12 = administration-temporary (TODO can this be eliminated) -- - 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 -> []