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
|
module X64Optimiser(x64Optimise) where
import Data.List
import Defs
import X64
x64Optimise :: Asm -> Error Asm
x64Optimise asm =
return $
funcopt optCoalesceInstructions $
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 ins = [ins]
optCoalesceInstructions :: Func -> Func
optCoalesceInstructions (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
|