module X64Optimiser(x64Optimise) where

import Data.List
import Data.Maybe

import Defs hiding (Offset)
import X64


x64Optimise :: Asm -> Error Asm
x64Optimise asm =
    return $
    funcopt optSimpleInstructions $
    funcopt optDoubleAdd $
    funcopt optMergeRSP $
    funcopt optMergeRSP $  -- #HACK (sometimes needed to eliminate all rsp arithmetic)
    optUnnecessaryJumps $
    funcopt optSimpleInstructions $
    asm

funcopt :: (Func -> Func) -> Asm -> Asm
funcopt f (Asm funcs) = Asm (map f funcs)

optUnnecessaryJumps :: Asm -> Asm
optUnnecessaryJumps (Asm funcs) = Asm $ map goF (zip funcs (tail funcs)) ++ [last funcs]
  where
    goF :: (Func, Func) -> Func
    goF (f1@(_, f1i), (f2n, _)) = case last f1i of
        JMP n | n == f2n -> fmap init f1
        _ -> f1

optSimpleInstructions :: Func -> Func
optSimpleInstructions (name, inss) = (name, concat $ map goI inss)
  where
    goI :: Ins -> [Ins]
    goI (MOV (RegMem a) (RegMemImm b)) | a == b = []
    goI (MOV (RegMem (XReg 8 r)) (RegMemImm (XImm 0))) = [XOR (RegMem (XReg 4 r)) (RegMemImm (XReg 4 r))]
    goI (MOV (RegMem a@(XReg _ _)) (RegMemImm (XImm 0))) = [XOR (RegMem a) (RegMemImm a)]
    goI (MOVi (Reg (XReg 8 r)) (Imm (XImm 0))) = [XOR (RegMem (XReg 4 r)) (RegMemImm (XReg 4 r))]
    goI (MOVi (Reg a) (Imm (XImm 0))) = [XOR (RegMem a) (RegMemImm a)]
    goI (MOVSX (Reg a) (RegMem b)) | a == b = []
    goI (ADD _ (RegMemImm (XImm 0))) = []
    goI (SUB _ (RegMemImm (XImm 0))) = []
    goI ins = [ins]

optMergeRSP :: Func -> Func
optMergeRSP (name, inss) = (name, go inss)
  where
    go :: [Ins] -> [Ins]
    go [] = []
    go (add@(ADD (RegMem (XReg 8 RSP)) (RegMemImm (XImm n))) : rest) =
        let midx = flip findIndex rest $ \ins -> case ins of
                SUB (RegMem (XReg 8 RSP)) (RegMemImm (XImm n')) | n' == n -> True
                _ -> False
        in case midx of
                Nothing -> add : go rest
                Just idx -> case mapM (shiftRSP n) (take idx rest) of
                    Nothing -> add : go rest
                    Just shifted -> shifted ++ go (drop (idx + 1) rest)
    go (MOV (RegMem (XMem 8 (Just RSP) (0, _) Nothing (-8))) (RegMemImm r@(XReg 8 _)) :
        SUB (RegMem (XReg 8 RSP)) (RegMemImm (XImm 8)) :
        rest) =
            PUSH (RegMemImm r) : go rest
    go (ins : rest) = ins : go rest

    isNonLinear :: Ins -> Bool
    isNonLinear (CALL _) = True
    isNonLinear (JMP _) = True
    isNonLinear (JCC _ _) = True
    isNonLinear RET = True
    isNonLinear _ = False

    shiftRSP :: Offset -> Ins -> Maybe Ins
    shiftRSP _ ins | isNonLinear ins = Nothing
    shiftRSP off ins = flip xrefMapM ins $ \thexref -> case thexref of
        XMem sz (Just RSP) (0, zero) lbl o -> Just $ XMem sz (Just RSP) (0, zero) lbl (o + off)
        XMem sz Nothing (c, RSP) lbl o -> Just $ XMem sz Nothing (c, RSP) lbl (o + (fromIntegral c) * off)
        XMem sz (Just RSP) (c, RSP) lbl o -> Just $ XMem sz (Just RSP) (c, RSP) lbl (o + (fromIntegral c + 1) * off)
        x@(XImm _) -> Just x
        XReg _ RSP -> Nothing
        XMem _ (Just RSP) _ _ _ -> Nothing
        XMem _ _ (_, RSP) _ _ -> Nothing
        x@(XReg _ _) -> Just x
        x@(XMem _ _ _ _ _) -> Just x

optDoubleAdd :: Func -> Func
optDoubleAdd (name, inss) = (name, go inss)
  where
    go :: [Ins] -> [Ins]
    go [] = []
    go (add@(ADD (RegMem xreg@(XReg _ xregReg)) (RegMemImm (XImm _))) : rest) = start xreg xregReg add rest
    go (sub@(SUB (RegMem xreg@(XReg _ xregReg)) (RegMemImm (XImm _))) : rest) = start xreg xregReg sub rest
    go (ins : rest) = ins : go rest

    start :: XRef -> Register -> Ins -> [Ins] -> [Ins]
    start xreg xregReg addsub rest =
        let midx = flip findIndex rest $ \ins -> case ins of
                ADD (RegMem xreg2@(XReg _ _)) (RegMemImm (XImm _)) | xreg == xreg2 -> True
                SUB (RegMem xreg2@(XReg _ _)) (RegMemImm (XImm _)) | xreg == xreg2 -> True
                _ -> False
        in case midx of
                Nothing -> addsub : go rest
                Just idx -> if all (canSkip xregReg) (take idx rest)
                    then go $ merge addsub (rest !! idx) : take idx rest ++ drop (idx + 1) rest
                    else addsub : go rest

    canSkip :: Register -> Ins -> Bool
    canSkip _ (CALL _) = False
    canSkip _ (JMP _) = False
    canSkip _ (JCC _ _) = False
    canSkip _ RET = False
    canSkip reg ins =
        isJust $ xrefMapM (\y -> if y `containsReg` reg then Nothing else Just y) ins

    containsReg :: XRef -> Register -> Bool
    containsReg (XReg _ r) reg | r == reg = True
    containsReg (XMem _ (Just r) _ _ _) reg | r == reg = True
    containsReg (XMem _ _ (s, r) _ _) reg | s /= 0 && r == reg = True
    containsReg _ _ = False

    merge :: Ins -> Ins -> Ins
    merge ins1 ins2 =
        let e1 = effectOf ins1
            e2 = effectOf ins2
            dst1 = destOf ins1
            dst2 = destOf ins2
        in if dst1 == dst2
            then (if e1 + e2 < 0 then SUB else ADD) (RegMem dst1) (RegMemImm $ XImm $ abs $ e1 + e2)
            else undefined

    effectOf :: Ins -> Offset
    effectOf (ADD _ (RegMemImm (XImm i))) = i
    effectOf (SUB _ (RegMemImm (XImm i))) = -i
    effectOf _ = undefined

    destOf :: Ins -> XRef
    destOf (ADD (RegMem d) _) = d
    destOf (SUB (RegMem d) _) = d
    destOf _ = undefined