aboutsummaryrefslogtreecommitdiff
path: root/X64Optimiser.hs
blob: 206529f60d0b31b0f4c7373383f2ffc2ca321f79 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
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