module Optimiser(optimise) where

import Data.Either
import Data.Function
import Data.List
import Data.Maybe
import qualified Data.Map.Strict as Map
import Debug.Trace

import Defs
import Intermediate
import Pretty
import ReplaceRefs
import Utils


type Optimisation = IRProgram -> IRProgram
type FuncOptimisation = IRFunc -> IRFunc

optimise :: IRProgram -> Error IRProgram
optimise prog =
    let optlist = [trace "-- OPT PASS --" {-, \p -> trace (pretty p) p-}] ++ optimisations
        reslist = scanl (flip ($)) prog $ cycle optlist
        passreslist = map fst $ filter (\(_, i) -> i `mod` length optlist == 0) $ zip reslist [0..]
        applyFinalOpts p = foldl (flip ($)) p finaloptimisations
    in if True
        then return $ applyFinalOpts $ 
                fst $ fromJust $ find (uncurry (==)) $ zip passreslist (tail passreslist)
        else return $ reslist !! 5
  where
    optimisations = map funcopt
        [chainJumps, mergeTerminators, looseJumps,
         removeUnusedBlocks, removeDuplicateBlocks,
         identityOps,
         constantPropagate, movPush,
         arithPush, removeUnusedInstructions,
         evaluateInstructions, evaluateTerminators]
    finaloptimisations = map funcopt
        [reorderBlocks, flipJccs, invertJccs]


funcopt :: FuncOptimisation -> Optimisation
funcopt fo (IRProgram vars funcs) = IRProgram vars (map fo funcs)


chainJumps :: FuncOptimisation
chainJumps (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    bbs' = snd $ last $ takeWhile fst $ iterate (mergeChain . snd) (True, bbs)
    
    mergeChain :: [BB] -> (Bool, [BB])
    mergeChain [] = (False, [])
    mergeChain bbs2 = case findIndex isSuitable bbs2 of
        Nothing -> (False, bbs2)
        Just idx ->
            let (BB bid1 inss1 (IJmp target), rest) =
                    (bbs2 !! idx, take idx bbs2 ++ drop (idx+1) bbs2)
                [BB _ inss2 term2] = filter (\(BB bid _ _) -> bid == target) rest
                merged = BB bid1 (inss1 ++ inss2) term2
            in (True, merged : rest)
      where
        hasJmpTo :: Id -> BB -> Bool
        hasJmpTo i (BB _ _ (IJmp i')) = i == i'
        hasJmpTo i (BB _ _ (IJcc _ _ _ i1 i2)) = i == i1 || i == i2
        hasJmpTo _ _ = False

        isSuitable :: BB -> Bool
        isSuitable (BB _ _ (IJmp target)) = sum (map (fromEnum . hasJmpTo target) bbs2) == 1
        isSuitable _ = False

mergeTerminators :: FuncOptimisation
mergeTerminators (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    bbs' = flip map bbs $ \bb@(BB bid inss term) -> case term of
        IJmp i -> case find ((== i) . fst) singles of
            Just (_, t) -> BB bid inss t
            Nothing -> bb
        _ -> bb

    singles = map (\(BB i _ t) -> (i, t)) $ filter (\(BB _ inss _) -> null inss) bbs

looseJumps :: FuncOptimisation
looseJumps (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    bbs' = flip map bbs $ \bb@(BB bid inss term) -> case term of
        IJmp i -> BB bid inss (IJmp (translate i))
        IJcc ct r1 r2 i j -> BB bid inss (IJcc ct r1 r2 (translate i) (translate j))
        _ -> bb

    translate i = fromMaybe i $ Map.lookup i transmap

    transmap = Map.fromList $ catMaybes $ flip map bbs $ \bb -> case bb of
        BB bid [] (IJmp i) -> Just (bid, i)
        _ -> Nothing

removeUnusedBlocks :: FuncOptimisation
removeUnusedBlocks (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    bbs' = filter isReachable bbs

    isReachable :: BB -> Bool
    isReachable (BB bid _ _)
        | bid == sid = True
        | otherwise = isJust $ flip find bbs $ \(BB _ _ term) -> case term of
            IJcc _ _ _ i1 i2 -> i1 == bid || i2 == bid
            IJmp i -> i == bid
            _ -> False

removeDuplicateBlocks :: FuncOptimisation
removeDuplicateBlocks (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    bbs' = let (bbspre, repls) = foldr foldfunc ([], []) bbs
           in foldl (\l (from, to) -> replaceBBIds from to l) bbspre repls

    foldfunc bb@(BB bid inss term) (l, repls) =
        case find (\(BB _ inss' term') -> inss == inss' && term == term') l of
            Nothing -> (bb : l, repls)
            Just (BB bid' _ _) -> (l, (bid, bid') : repls)

    replaceBBIds :: Id -> Id -> [BB] -> [BB]
    replaceBBIds from to = map $ \(BB bid inss term) -> BB bid inss $ case term of
        IJcc ct r1 r2 i1 i2 -> IJcc ct r1 r2 (trans from to i1) (trans from to i2)
        IJmp i -> IJmp (trans from to i)
        IRet -> IRet
        IRetr r -> IRetr r
        IUnreachable -> IUnreachable
        ITermNone -> undefined

    trans :: (Eq a) => a -> a -> a -> a
    trans a b c | a == c = b
                | otherwise = c

identityOps :: FuncOptimisation
identityOps (IRFunc rt name al bbs sid) = IRFunc rt name al (map go bbs) sid
  where
    go :: BB -> BB
    go (BB bid inss term) = BB bid (catMaybes $ map goI inss) term

    goI :: IRIns -> Maybe IRIns
    goI (IAri AAdd d s (Constant _ 0)) = Just $ IMov d s
    goI (IAri AAdd d (Constant _ 0) s) = Just $ IMov d s
    goI (IAri ASub d s (Constant _ 0)) = Just $ IMov d s
    goI (IAri AMul d s (Constant _ 1)) = Just $ IMov d s
    goI (IAri AMul d (Constant _ 1) s) = Just $ IMov d s
    goI (IAri ADiv d s (Constant _ 1)) = Just $ IMov d s
    goI (IMov d s) | d == s = Nothing
    goI i = Just i

constantPropagate :: FuncOptimisation
constantPropagate (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    alltemps = findAllTempsBBList bbs
    consttemps = catMaybes $ flip map alltemps $ \ref ->
        let locs = findMutations' bbs ref
            loc = head locs
            ins = insAt bbs loc
            readlocs = findMentions' bbs ref \\ locs
            readinss = map (insAt' bbs) readlocs
            allimov = all (maybe False isIMov) readinss
        in if length locs == 1 && (isIMov ins || ((isILoad ins || isIAri ins || isIResize ins) && allimov))
               then Just (loc, ins)
               else Nothing

    bbs' = case consttemps of
                [] -> bbs
                ((loc, IMov ref value) : _) ->
                    replaceRefsBBList ref value (nopifyInsAt bbs loc)
                ((loc, ILoad ref s) : _) ->
                    replaceMovs ref (\r' -> ILoad r' s) (nopifyInsAt bbs loc)
                ((loc, IAri at ref s1 s2) : _) ->
                    replaceMovs ref (\r' -> IAri at r' s1 s2) (nopifyInsAt bbs loc)
                ((loc, IResize ref s) : _) ->
                    replaceMovs ref (\r' -> IResize r' s) (nopifyInsAt bbs loc)
                _ -> undefined

    replaceMovs :: Ref -> (Ref -> IRIns) -> [BB] -> [BB]
    replaceMovs srcref insb = map $ \(BB bid inss term) -> BB bid (map go inss) term
      where
        go :: IRIns -> IRIns
        go (IMov d src) | src == srcref = insb d
        go ins = ins

movPush :: FuncOptimisation
movPush (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs) sid
  where
    goBB :: BB -> BB
    goBB (BB bid inss term) =
        let inss' = go inss term
            term' = if null inss' then term else pushT (last inss) term
        in BB bid inss' term'

    go :: [IRIns] -> IRTerm -> [IRIns]
    go [] _ = []
    go (IMov d s : rest) term
        | isJust (find (== d) (findAllRefsInss rest ++ findAllRefsTerm term)) =
            push (d, s) rest term
    go (ins : rest) term = ins : go rest term

    push :: (Ref, Ref) -> [IRIns] -> IRTerm -> [IRIns]
    push (d, s) [] _ = [IMov d s]
    push (d, s) l _ | d == s = l
    push mov@(d, s) (IMov d' s' : rest) term
        | d' == d = if d' == s' then push mov rest term else push (d', s') rest term
        | d' == s = IMov d s : push (d', replaceRef d s s') rest term
        | otherwise = IMov d' (replaceRef d s s') : push mov rest term
    push mov@(d, _) (ILea d' n : rest) term
        | d' == d = ILea d' n : go rest term
        | otherwise = ILea d' n : push mov rest term
    push mov@(d, s) (IResize d' s' : rest) term
        | d' == d = IResize d' (replaceRef d s s') : go rest term
        | d' == s = IMov d s : IResize d' (replaceRef d s s') : go rest term
        | otherwise = IResize d' (replaceRef d s s') : push mov rest term
    push mov@(d, s) (ILoad d' s' : rest) term
        | d' == d = ILoad d' (replaceRef d s s') : go rest term
        | d' == s = IMov d s : ILoad d' (replaceRef d s s') : go rest term
        | otherwise = ILoad d' (replaceRef d s s') : push mov rest term
    push mov@(d, s) (ISet d' n' s' : rest) term
        | d' == d = ISet d' n' (replaceRef d s s') : go rest term
        | d' == s = IMov d s : ISet d' n' s' : go rest term
        | otherwise = ISet d' n' (replaceRef d s s') : push mov rest term
    push mov@(d, s) (IGet d' s' n' : rest) term
        | d' == d = IGet d' (replaceRef d s s') n' : go rest term
        | d' == s = IMov d s : IGet d' s' n' : go rest term
        | otherwise = IGet d' (replaceRef d s s') n' : push mov rest term
    push mov@(d, s) (IAri at d' s1' s2' : rest) term
        | d' == d = IAri at d' (replaceRef d s s1') (replaceRef d s s2') : go rest term
        | d' == s = IMov d s : IAri at d' (replaceRef d s s1') (replaceRef d s s2') : go rest term
        | otherwise = IAri at d' (replaceRef d s s1') (replaceRef d s s2') : push mov rest term
    -- I don't trust going past calls because globals might change. Might be able to
    -- catch that case, but that will go wrong when more stuff gets added.
    -- push mov@(d, s) (ins@(ICallr d' _ _) : rest) term
    --     | d' == d = IMov d s : ins : go rest term
    --     | otherwise = replaceRefsIns d s ins : push mov rest term
    -- push mov@(d, s) (ins@(ICall _ _) : rest) term = replaceRefsIns d s ins : push mov rest term
    push (d, s) l@(ICallr _ _ _ : _) term = IMov d s : go l term
    push (d, s) l@(ICall _ _ : _) term = IMov d s : go l term
    push mov@(d, s) (ins@(IStore _ _) : rest) term = replaceRefsIns d s ins : push mov rest term
    push (d, s) l@(IDebugger : _) term = IMov d s : go l term
    push mov (INop : rest) term = push mov rest term

    pushT :: IRIns -> IRTerm -> IRTerm
    pushT (IMov d s) term = replaceRefsTerm d s term
    pushT _ term = term

arithPush :: FuncOptimisation
arithPush (IRFunc rt name al allbbs sid) = IRFunc rt name al resbbs sid
  where
    resbbs = foldl (\bbs i -> goBB (blockById i bbs) bbs) allbbs (map blockIdOf allbbs)

    goBB :: BB -> [BB] -> [BB]
    goBB bb@(BB bid _ _) bbs =
        let (mari, (inss', [Right term'])) = fmap (span isLeft) $ go (bbToList bb)
            resbbs1 = replaceBlock bid (BB bid (map (fromLeft undefined) inss') term') bbs
        in case mari of
            Nothing -> resbbs1
            Just ari ->
                let tgs = map (flip blockById bbs) $
                            filter (\b -> length (originBlocks b) == 1) $ jumpTargets term'
                in foldl (\bbs' tg -> propagateContinue ari tg bbs') resbbs1 tgs

    propagateContinue :: IRIns -> BB -> [BB] -> [BB]
    -- propagateContinue ari bb _ | traceShow (ari, bb) False = undefined
    propagateContinue ari@(IAri at d s1 s2) bb@(BB bid _ _) bbs =
        let (cont, (inss', [Right term'])) = fmap (span isLeft) $ propagate (at, d, s1, s2) (bbToList bb)
            resbbs1 = replaceBlock bid (BB bid (map (fromLeft undefined) inss') term') bbs
        in if cont
            then let tgs = map (flip blockById bbs) $
                             filter (\b -> length (originBlocks b) == 1) $ jumpTargets term'
                 in foldl (\bbs' tg -> propagateContinue ari tg bbs') resbbs1 tgs
            else resbbs1
    propagateContinue _ _ _ = undefined

    blockById :: Id -> [BB] -> BB
    blockById i bbs = head $ filter (\(BB bid _ _) -> bid == i) bbs

    originBlocks :: Id -> [BB]
    originBlocks i = filter (\(BB _ _ term) -> i `elem` jumpTargets term) allbbs

    replaceBlock :: Id -> BB -> [BB] -> [BB]
    replaceBlock _ _ [] = []
    replaceBlock bid bb (bb'@(BB bid' _ _) : rest)
        | bid' == bid = bb : rest
        | otherwise = bb' : replaceBlock bid bb rest

    go :: [Either IRIns IRTerm] -> (Maybe IRIns, [Either IRIns IRTerm])
    go [] = (Nothing, [])
    go (Left ari@(IAri at d s1 s2) : rest) = case propagate (at, d, s1, s2) rest of
        (False, res) -> fmap (Left ari :) $ go res
        (True, res) -> (Just ari, Left ari : res)
    go (ins : rest) = fmap (ins :) $ go rest

    bbToList :: BB -> [Either IRIns IRTerm]
    bbToList (BB _ inss term) = map Left inss ++ [Right term]

    propagate :: (ArithType, Ref, Ref, Ref) -> [Either IRIns IRTerm] -> (Bool, [Either IRIns IRTerm])
    propagate _ [] = (True, [])
    propagate ari@(_, d, s1, s2) (Left ins@(IMov md _) : rest)
        | d /= md && md /= s1 && md /= s2 = fmap (Left ins :) $ propagate ari rest
        | otherwise = (False, Left ins : rest)
    propagate ari (Left ins@(IStore _ _) : rest) =
        fmap (Left ins :) $ propagate ari rest
    propagate ari@(_, d, s1, s2) (Left ins@(ILoad md _) : rest)
        | d /= md && md /= s1 && md /= s2 = fmap (Left ins :) $ propagate ari rest
        | otherwise = (False, Left ins : rest)
    propagate ari@(_, d, s1, s2) (Left ins@(ISet md _ _) : rest)
        | d /= md && md /= s1 && md /= s2 = fmap (Left ins :) $ propagate ari rest
        | otherwise = (False, Left ins : rest)
    propagate ari@(_, d, s1, s2) (Left ins@(IGet md _ _) : rest)
        | d /= md && md /= s1 && md /= s2 = fmap (Left ins :) $ propagate ari rest
        | otherwise = (False, Left ins : rest)
    propagate ari@(at, d, s1, s2) (Left ins@(IAri mat md ms1 ms2) : rest)
        | d /= md && (at, s1, s2) == (mat, ms1, ms2) = fmap (Left (IMov md d) :) $ propagate ari rest
        | d /= md && md /= s1 && md /= s2 = fmap (Left ins :) $ propagate ari rest
        | otherwise = fmap (Left ins :) $ propagate (mat, md, ms1, ms2) rest
    -- I don't trust going past calls because globals might change. Might be able to
    -- catch that case, but that will go wrong when more stuff gets added.
    -- propagate ari@(_, d, s1, s2) (Left ins@(ICall _ mal) : rest)
    --     | null (intersect [d] mal) = fmap (Left ins :) $ propagate ari rest
    --     | otherwise = (False, Left ins : rest)
    -- propagate ari@(_, d, s1, s2) (Left ins@(ICallr md _ mal) : rest)
    --     | null (intersect [d,s1,s2] (md : mal)) = fmap (Left ins :) $ propagate ari rest
    --     | otherwise = (False, Left ins : rest)
    propagate ari@(_, d, s1, s2) (Left ins@(IResize md _) : rest)
        | d /= md && md /= s1 && md /= s2 = fmap (Left ins :) $ propagate ari rest
        | otherwise = (False, Left ins : rest)
    propagate ari@(_, d, s1, s2) (Left ins@(ILea md _) : rest)
        | d /= md && md /= s1 && md /= s2 = fmap (Left ins :) $ propagate ari rest
        | otherwise = (False, Left ins : rest)
    propagate ari (Left INop : rest) = propagate ari rest
    propagate (at, d, s1, s2) (Right term@(IJcc ct r1 r2 i1 i2) : rest)
        | (r1 == d || r2 == d) &&
          (isConstant r1 || isConstant r2) &&
          at `elem` [AEq, ANeq, AGt, ALt, AGeq, ALeq] =
            let ct' = if isConstant r2 then ct else flipCmpType ct
                conref = if isConstant r2 then r2 else r1
                (ct'', con) = case (ct', conref) of
                    (CEq, Constant _ c) -> (CEq, if c `elem` [0, 1] then c else (-1))
                    (CNeq, Constant _ c) -> (CNeq, if c `elem` [0, 1] then c else (-1))
                    (CGt, Constant _ c) | c < 0 -> (CNeq, (-1))
                                        | c == 0 -> (CEq, 1)
                                        | otherwise -> (CEq, (-1))
                    (CLt, Constant _ c) | c > 1 -> (CNeq, (-1))
                                        | c == 1 -> (CEq, 0)
                                        | otherwise -> (CEq, (-1))
                    (CGeq, Constant _ c) | c <= 0 -> (CNeq, (-1))
                                         | c == 1 -> (CEq, 1)
                                         | otherwise -> (CEq, (-1))
                    (CLeq, Constant _ c) | c >= 1 -> (CNeq, (-1))
                                         | c == 0 -> (CEq, 0)
                                         | otherwise -> (CEq, (-1))
                    _ -> undefined
                resterm = case (ct'', con) of
                    (CEq, 0) -> IJcc (invertCmpType (arithTypeToCmpType at)) s1 s2 i1 i2
                    (CEq, 1) -> IJcc (arithTypeToCmpType at) s1 s2 i1 i2
                    (CEq, _) -> IJmp i2
                    (CNeq, 0) -> IJcc (arithTypeToCmpType at) s1 s2 i1 i2
                    (CNeq, 1) -> IJcc (invertCmpType (arithTypeToCmpType at)) s1 s2 i1 i2
                    (CNeq, _) -> IJmp i1
                    _ -> undefined
            in (True, Right resterm : rest)
        | otherwise = (True, Right term : rest)
    propagate _ l@(Left (ICall _ _) : _) = (False, l)
    propagate _ l@(Left (ICallr _ _ _) : _) = (False, l)
    propagate _ l@(Left IDebugger : _) = (False, l)
    propagate _ l@(Right (IJmp _) : _) = (True, l)
    propagate _ l@(Right IRet : _) = (False, l)
    propagate _ l@(Right (IRetr _) : _) = (False, l)
    propagate _ l@(Right IUnreachable : _) = (False, l)
    propagate _ (Right ITermNone : _) = undefined

flipCmpType :: CmpType -> CmpType
flipCmpType CEq = CEq
flipCmpType CNeq = CNeq
flipCmpType CGt = CLt
flipCmpType CLt = CGt
flipCmpType CGeq = CLeq
flipCmpType CLeq = CGeq
flipCmpType CUGt = CULt
flipCmpType CULt = CUGt
flipCmpType CUGeq = CULeq
flipCmpType CULeq = CUGeq

invertCmpType :: CmpType -> CmpType
invertCmpType CEq = CNeq
invertCmpType CNeq = CEq
invertCmpType CGt = CLeq
invertCmpType CLt = CGeq
invertCmpType CGeq = CLt
invertCmpType CLeq = CGt
invertCmpType CUGt = CULeq
invertCmpType CULt = CUGeq
invertCmpType CUGeq = CULt
invertCmpType CULeq = CUGt

arithTypeToCmpType :: ArithType -> CmpType
arithTypeToCmpType AEq = CEq
arithTypeToCmpType ANeq = CNeq
arithTypeToCmpType AGt = CGt
arithTypeToCmpType ALt = CLt
arithTypeToCmpType AGeq = CGeq
arithTypeToCmpType ALeq = CLeq
arithTypeToCmpType _ = undefined

removeUnusedInstructions :: FuncOptimisation
removeUnusedInstructions (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs) sid
  where
    goBB :: BB -> BB
    goBB (BB bid inss term) = BB bid (catMaybes $ map goI inss) term

    goI :: IRIns -> Maybe IRIns
    goI ins@(IMov d _) = pureInstruction d ins
    goI ins@(ILea d _) = pureInstruction d ins
    goI ins@(IStore _ _) = Just ins
    goI ins@(ILoad d _) = pureInstruction d ins
    goI ins@(ISet _ _ _) = Just ins
    goI ins@(IGet d _ _) = pureInstruction d ins
    goI ins@(IAri _ d _ _) = pureInstruction d ins
    goI ins@(ICall _ _) = Just ins
    goI ins@(ICallr _ _ _) = Just ins
    goI ins@(IResize d _) = pureInstruction d ins
    goI IDebugger = Just IDebugger
    goI INop = Nothing

    pureInstruction :: Ref -> IRIns -> Maybe IRIns
    pureInstruction d ins = if length (findMentions' bbs d) == 1 then Nothing else Just ins

evaluateInstructions :: FuncOptimisation
evaluateInstructions (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs) sid
  where
    goBB :: BB -> BB
    goBB (BB bid inss term) = BB bid (map goI inss) term

    goI :: IRIns -> IRIns
    goI (IAri at ref (Constant _ v1) (Constant _ v2)) =
        IMov ref $ Constant (refSize ref) $ truncValue (refSize ref) $ evaluateArith at v1 v2
    goI (IResize ref (Constant _ v)) =
        IMov ref $ Constant (refSize ref) $ truncValue (refSize ref) v
    goI ins = ins

    truncValue :: Size -> Value -> Value
    truncValue sz v = fromIntegral $ (fromIntegral v :: Integer) `mod` (2 ^ (8 * sz))

evaluateTerminators :: FuncOptimisation
evaluateTerminators (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    bbs' = map (\(BB bid inss term) -> BB bid inss (go term)) bbs

    go :: IRTerm -> IRTerm
    go term@(IJcc ct (Constant sza a) (Constant szb b) i1 i2)
        | sza /= szb = error $ "Inconsistent sizes in " ++ show term
        | evaluateCmp ct a b = IJmp i1
        | otherwise = IJmp i2
    go term = term

flipJccs :: FuncOptimisation
flipJccs (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs) sid
  where
    goBB :: BB -> BB
    goBB (BB bid inss term) = BB bid inss (goT term)

    goT :: IRTerm -> IRTerm
    goT (IJcc ct r1@(Constant _ _) r2 i1 i2) = IJcc (flipCmpType ct) r2 r1 i1 i2
    goT term = term

reorderBlocks :: FuncOptimisation
reorderBlocks (IRFunc rt name al allbbs sid) = IRFunc rt name al resbbs sid
  where
    resbbs = buildResult (allChainsFrom allbbs sid) allbbs

    allChains :: [BB] -> [[Id]]
    allChains bbs = concatMap (allChainsFrom bbs . blockIdOf) bbs

    allChainsFrom :: [BB] -> Id -> [[Id]]
    allChainsFrom b start = go b [] start
      where
        go :: [BB] -> [Id] -> Id -> [[Id]]
        go bbs chain at =
            let ((BB _ _ term), rest) = takeBlock at bbs
                chain' = chain ++ [at]
            in case intersect (jumpTargets term) (map blockIdOf rest) of
                   [] -> [chain']
                   tgs -> flip concatMap tgs $ \tg ->
                            if hasUnreachable (fst $ takeBlock tg bbs)
                                then []
                                else go rest chain' tg

    buildResult :: [[Id]] -> [BB] -> [BB]
    buildResult _ [] = []
    buildResult chains bbs =
        let chain = maximumBy (compare `on` length) chains
            (chainbbs', newbbs) = partition ((`elem` chain) . blockIdOf) bbs
            chainbbs = sortBy (compare `on` (\(BB i _ _) -> fromJust $ findIndex (== i) chain)) chainbbs'
            newchains = allChains newbbs
        in chainbbs ++ buildResult newchains newbbs

    takeBlock :: Id -> [BB] -> (BB, [BB])
    takeBlock _ [] = undefined
    takeBlock target (bb@(BB bid _ _) : rest)
        | bid == target = (bb, rest)
        | otherwise = fmap (bb :) $ takeBlock target rest

    hasUnreachable :: BB -> Bool
    hasUnreachable (BB _ _ IUnreachable) = True
    hasUnreachable _ = False

invertJccs :: FuncOptimisation
invertJccs (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
  where
    bbs' = map goBB (zip bbs (tail bbs)) ++ [last bbs]

    goBB :: (BB, BB) -> BB
    goBB (BB bid inss term, BB nextbid _ _) = BB bid inss (goT term nextbid)

    goT :: IRTerm -> Id -> IRTerm
    goT (IJcc ct r1 r2 i1 i2) next | i1 == next = IJcc (invertCmpType ct) r1 r2 i2 i1
    goT term _ = term


insAt :: [BB] -> (Int, Int) -> IRIns
insAt bbs (i, j) =
    let (BB _ inss _) = bbs !! i
    in inss !! j

insAt' :: [BB] -> (Int, Int) -> Maybe IRIns
insAt' bbs (i, j) = do
    (BB _ inss _) <- if i >= length bbs then Nothing else Just (bbs !! i)
    if j >= length inss then Nothing else Just (inss !! j)

nopifyInsAt :: [BB] -> (Int, Int) -> [BB]
nopifyInsAt bbs (i, j) =
    let (pre, BB bid inss term : post) = splitAt i bbs
        (ipre, _ : ipost) = splitAt j inss
    in pre ++ BB bid (ipre ++ INop : ipost) term : post

findMutations :: BB -> Ref -> [Int]
findMutations (BB _ inss _) ref =
    catMaybes $ flip map (zip inss [0..]) $ \(ins, idx) -> case ins of
        (IMov r _) | r == ref -> Just idx
        (ILoad r _) | r == ref -> Just idx
        (IAri _ r _ _) | r == ref -> Just idx
        (ICallr r _ _) | r == ref -> Just idx
        (IResize r _) | r == ref -> Just idx
        _ -> Nothing

findMutations' :: [BB] -> Ref -> [(Int, Int)]
findMutations' bbs ref =
    [(i, j) | (bb, i) <- zip bbs [0..], j <- findMutations bb ref]

findMentions :: BB -> Ref -> [Int]
findMentions (BB _ inss term) ref = insres ++ termres
  where
    insres = catMaybes $ flip map (zip inss [0..]) $ \(ins, idx) ->
        if ref `elem` findAllRefsIns ins
            then Just idx
            else Nothing
    termres = if ref `elem` findAllRefsTerm term
                then [length inss]
                else []

findMentions' :: [BB] -> Ref -> [(Int, Int)]
findMentions' bbs ref =
    [(i, j) | (bb, i) <- zip bbs [0..], j <- findMentions bb ref]

-- findMentionsIns :: BB -> Ref -> [IRIns]
-- findMentionsIns (BB _ inss term) ref = insres ++ termres
--   where
--     insres = catMaybes $ flip map (zip inss [0..]) $ \(ins, idx) ->
--         if ref `elem` findAllRefsIns ins
--             then Just ins
--             else Nothing
--     termres = if ref `elem` findAllRefsTerm term
--                 then [term]
--                 else []

-- findMentionsIns' :: [BB] -> Ref -> [IRIns]
-- findMentionsIns' bbs ref = concatMap (flip findMentionsIns ref) bbs

findAllTemps :: BB -> [Ref]
findAllTemps bb = flip filter (findAllRefs bb) $ \ref -> case ref of
    (Temp _ _) -> True
    _ -> False

findAllTempsBBList :: [BB] -> [Ref]
findAllTempsBBList = concatMap findAllTemps