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
|