module Bril.CFG (BasicBlock, CFG, createCFG) where

import Bril
import qualified Data.Map.Strict as Map

data BasicBlock = BasicBlock {
        instrs :: [Instruction],
        incoming :: [String]
    }
   deriving (Show)

data CFG = CFG {
        cfgBlocks :: Map.Map String BasicBlock,
        cfgStart :: String
    }
   deriving (Show)

createCFG :: [Instruction] -> CFG
createCFG instrs =
    let
        (initCFG, firstLabel) = initializeCFG $ addFallthroughJmp $ splitBBs $ addInitialLabel instrs
        cfgBlocks' = foldr processBB (cfgBlocks initCFG) (Map.keys (cfgBlocks initCFG))
    in
        CFG cfgBlocks' firstLabel
    where
        processBB :: String -> Map.Map String BasicBlock -> Map.Map String BasicBlock
        processBB label m = addBBJumps label (outgoingCFGLabels m label) m
        addBBJumps :: String -> [String] -> Map.Map String BasicBlock -> Map.Map String BasicBlock
        addBBJumps bblockLabel incomingLabels m = Map.adjust (addLabels incomingLabels) bblockLabel m
        addLabels :: [String] -> BasicBlock -> BasicBlock
        addLabels labels (BasicBlock instrs incoming) = (BasicBlock instrs (incoming ++ labels))

initializeCFG :: [([Instruction], String)] -> (CFG, String)
initializeCFG instrs =
    let mappedBlocks = foldr mapBlock Map.empty instrs -- Map.fromList $ map swap instrs
        firstLabel = snd . head $ instrs
    in (CFG mappedBlocks firstLabel, firstLabel)
    where
        mapBlock (instrs, label) m = Map.insert label (BasicBlock instrs []) m

addFallthroughJmp :: [([Instruction], String)] -> [([Instruction], String)]
addFallthroughJmp blocks =
    let blocks' = map addJumpToNextBB $ zip blocks $ tail $ map snd blocks
    in blocks' ++ [last blocks]
    where
        addJumpToNextBB (block@(instrs, label), nextLabel)
            -- TODO: handle empty instrs
            | isJump (last instrs) = block
            | otherwise = (instrs ++ [EffectInstr Jmp [] [] [nextLabel]], label)

splitBBs :: [Instruction] -> [([Instruction], String)]
splitBBs instrs = zip (tail basicBlocks) labels
    where
        (basicBlocks, labels) = foldr splitBBs' ([[]], []) instrs
        splitBBs' instr (bblocks, labels)
                | isLabel instr = ([] : bblocks, (labelName instr) : labels)
                | isJump instr = ([instr] : (tail bblocks) , labels)
                | otherwise = ((instr : (head bblocks)) : (tail bblocks), labels)

addInitialLabel :: [Instruction] -> [Instruction]
addInitialLabel instrs@(firstInstr : instrs')
    | isLabel firstInstr = instrs
    | otherwise = (Label "@start") : instrs
addInitialLabel [] = []

labelName (Label name) = name
isLabel instr = case instr of
    (Label _) -> True
    _ -> False
    
isJump instr = case instr of
    (EffectInstr op _ _ _) -> op == Jmp || op == Br
    _ -> False

outgoingCFGLabels :: Map.Map String BasicBlock -> String -> [String]
outgoingCFGLabels m label =
    case Map.lookup label m of
        Just bb -> outgoingLabels bb
        _ -> []

outgoingLabels :: BasicBlock -> [String]
outgoingLabels (BasicBlock instrs _) =
    let (EffectInstr lastOp _ _ labels) = last instrs
    in labels