module Optimiser(optimise) where

import Data.List

import Intermediate


optimise :: IRProgram -> IRProgram
optimise (IRProgram bbs gfds datas) = IRProgram (mergeBlocks bbs) gfds datas

mergeBlocks :: [BB] -> [BB]
mergeBlocks [] = []
mergeBlocks allbbs@(BB startb _ _ : _) =
    uncurry (++) (partition ((== startb) . bidOf) (go allbbs (length allbbs)))
  where
    go [] _ = []
    go bbs 0 = bbs
    go (bb@(BB bid inss term) : bbs) n = case partition (hasJumpTo bid . termOf) bbs of
        ([], _) -> go (bbs ++ [bb]) (n - 1)
        ([BB bid' inss' _], rest) -> go (BB bid' (inss' ++ inss) term : rest) n
        _ -> go (bbs ++ [bb]) (n - 1)

    hasJumpTo bid (IJmp a) = a == bid
    hasJumpTo _ _ = False