Seventh

6 files changed, 209 insertions, 89 deletions

diff --git a/CodeGen.hs b/CodeGen.hs
index 905bee5..8d7cb78 100644
--- a/CodeGen.hs
+++ b/CodeGen.hs
@@ -177,9 +177,9 @@ mkxref r m = fromJust $ Map.lookup r m
 mkmov :: XRef -> XRef -> X64.Ins
 mkmov a@(XReg _ _) b@(XReg _ _) = MOV (xref a) (xref b)
 mkmov a@(XReg _ _) b@(XMem _ _ _ _ _) = MOV (xref a) (xref b)
-mkmov a@(XReg _ _) b@(XImm _) = MOVi64 (xref a) (xref b)
+mkmov a@(XReg _ _) b@(XImm _) = MOVi (xref a) (xref b)
 mkmov a@(XMem _ _ _ _ _) b@(XReg _ _) = MOV (xref a) (xref b)
-mkmov a@(XMem _ _ _ _ _) b@(XImm v) | v < 2 ^ (32 :: Int) = MOVi (xref a) (xref b)
+mkmov a@(XMem _ _ _ _ _) b@(XImm v) | v < 2 ^ (32 :: Int) = MOV (xref a) (xref b)
 mkmov a b = CALL $ "Invalid mkmov: " ++ show a ++ "; " ++ show b
 -- mkmov a b = error $ "Invalid mkmov: " ++ show a ++ "; " ++ show b
 
diff --git a/Intermediate.hs b/Intermediate.hs
index ad3cb89..c395f55 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -200,6 +200,19 @@ isIMov :: IRIns -> Bool
 isIMov (IMov _ _) = True
 isIMov _ = False
 
+isILoad :: IRIns -> Bool
+isILoad (ILoad _ _) = True
+isILoad _ = False
+
 isIAri :: IRIns -> Bool
 isIAri (IAri _ _ _ _) = True
 isIAri _ = False
+
+isIResize :: IRIns -> Bool
+isIResize (IResize _ _) = True
+isIResize _ = False
+
+jumpTargets :: IRTerm -> [Id]
+jumpTargets (IJcc _ _ _ i2 i1) = [i1, i2]
+jumpTargets (IJmp i) = [i]
+jumpTargets _ = []
diff --git a/Optimiser.hs b/Optimiser.hs
index 2923af6..61834f8 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -33,10 +33,10 @@ optimise prog =
          removeUnusedBlocks, removeDuplicateBlocks,
          identityOps,
          constantPropagate, movPush,
-         arithPush False, removeUnusedInstructions,
+         arithPush, removeUnusedInstructions,
          evaluateInstructions, evaluateTerminators]
     finaloptimisations = map funcopt
-        [arithPush True, reorderBlocks, flipJccs, invertJccs]
+        [reorderBlocks, flipJccs, invertJccs]
 
 
 funcopt :: FuncOptimisation -> Optimisation
@@ -153,7 +153,10 @@ constantPropagate (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
         let locs = findMutations' bbs ref
             loc = head locs
             ins = insAt bbs loc
-        in if length locs == 1 && isIMov ins
+            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
 
@@ -161,8 +164,21 @@ constantPropagate (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
                 [] -> 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
@@ -212,46 +228,80 @@ movPush (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs) sid
     pushT (IMov d s) term = replaceRefsTerm d s term
     pushT _ term = term
 
-arithPush :: Bool -> FuncOptimisation
-arithPush ariari (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs) sid
+arithPush :: FuncOptimisation
+arithPush (IRFunc rt name al allbbs sid) = IRFunc rt name al resbbs sid
   where
-    goBB :: BB -> BB
-    goBB (BB bid inss term) =
-        let (inss', [Right term']) = span isLeft $ go (map Left inss ++ [Right term])
-        in BB bid (map (fromLeft undefined) inss') term'
-
-    go :: [Either IRIns IRTerm] -> [Either IRIns IRTerm]
-    go [] = []
-    go (Left ari@(IAri _ _ _ _) : rest) = Left ari : go (propagate ari rest)
-    go (ins : rest) = ins : go rest
-
-    propagate :: IRIns -> [Either IRIns IRTerm] -> [Either IRIns IRTerm]
-    propagate _ [] = []
-    propagate ari@(IAri at d s1 s2) (Left ins@(IMov md ms) : rest)
-        | d == ms = Left (IAri at md s1 s2) : (if d /= md then propagate ari rest else rest)
-        | d /= md && md /= s1 && md /= s2 = Left ins : propagate ari rest
-        | otherwise = Left ins : rest
-    propagate ari@(IAri _ d _ _) (Left ins@(IStore md ms) : rest)
-        | null (intersect [d] [md,ms]) = Left ins : propagate ari rest
-        | otherwise = Left ins : rest
-    propagate ari@(IAri _ d s1 s2) (Left ins@(ILoad md ms) : rest)
-        | null (intersect [d] [md,ms] ++ intersect [s1,s2] [md]) = Left ins : propagate ari rest
-        | otherwise = Left ins : rest
+    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 bb@(BB bid _ _) bbs =
+        let (cont, (inss', [Right term'])) = fmap (span isLeft) $ propagate ari (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
+
+    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 _ _ _ _) : rest) = case propagate ari 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 :: IRIns -> [Either IRIns IRTerm] -> (Bool, [Either IRIns IRTerm])
+    propagate _ [] = (True, [])
+    propagate ari@(IAri _ 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@(IAri _ _ _ _) (Left ins@(IStore _ _) : rest) =
+        fmap (Left ins :) $ propagate ari rest
+    propagate ari@(IAri _ 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@(IAri at d s1 s2) (Left ins@(IAri mat md ms1 ms2) : rest)
-        | ariari && d /= md && (at, s1, s2) == (mat, ms1, ms2) = Left (IMov md d) : propagate ari rest
-        | null (intersect [d] [md,ms1,ms2] ++ intersect [s1,s2] [md]) = Left ins : propagate ari rest
-        | otherwise = Left ins : propagate ins 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 ins 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@(IAri _ d s1 s2) (Left ins@(ICall _ mal) : rest)
-    --     | null (intersect [d] mal) = Left ins : propagate ari rest
-    --     | otherwise = Left ins : rest
+    --     | null (intersect [d] mal) = fmap (Left ins :) $ propagate ari rest
+    --     | otherwise = (False, Left ins : rest)
     -- propagate ari@(IAri _ d s1 s2) (Left ins@(ICallr md _ mal) : rest)
-    --     | null (intersect [d,s1,s2] (md : mal)) = Left ins : propagate ari rest
-    --     | otherwise = Left ins : rest
-    propagate ari@(IAri _ d s1 s2) (Left ins@(IResize md ms) : rest)
-        | null (intersect [d] [md,ms] ++ intersect [s1,s2] [md]) = Left ins : propagate ari rest
-        | otherwise = Left ins : rest
+    --     | null (intersect [d,s1,s2] (md : mal)) = fmap (Left ins :) $ propagate ari rest
+    --     | otherwise = (False, Left ins : rest)
+    propagate ari@(IAri _ 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@(IAri _ _ _ _) (Left INop : rest) = propagate ari rest
     propagate (IAri at d s1 s2) (Right term@(IJcc ct r1 r2 i1 i2) : rest)
         | (r1 == d || r2 == d) &&
@@ -283,9 +333,9 @@ arithPush ariari (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs)
                     (CNeq, 1) -> IJcc (invertCmpType (arithTypeToCmpType at)) s1 s2 i1 i2
                     (CNeq, _) -> IJmp i1
                     _ -> undefined
-            in Right resterm : rest
-        | otherwise = Right term : rest
-    propagate _ l = l
+            in (True, Right resterm : rest)
+        | otherwise = (True, Right term : rest)
+    propagate _ l = (False, l)
 
 flipCmpType :: CmpType -> CmpType
 flipCmpType CEq = CEq
@@ -403,11 +453,6 @@ reorderBlocks (IRFunc rt name al allbbs sid) = IRFunc rt name al resbbs sid
         | bid == target = (bb, rest)
         | otherwise = fmap (bb :) $ takeBlock target rest
 
-    jumpTargets :: IRTerm -> [Id]
-    jumpTargets (IJcc _ _ _ i2 i1) = [i1, i2]
-    jumpTargets (IJmp i) = [i]
-    jumpTargets _ = []
-
 invertJccs :: FuncOptimisation
 invertJccs (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
   where
@@ -426,6 +471,11 @@ 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
@@ -436,6 +486,7 @@ 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
@@ -460,6 +511,20 @@ 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
+
 findAllRefs :: BB -> [Ref]
 findAllRefs (BB _ inss _) = findAllRefsInss inss
 
diff --git a/X64.hs b/X64.hs
index 42a8857..222d1cd 100644
--- a/X64.hs
+++ b/X64.hs
@@ -1,6 +1,7 @@
 module X64 where
 
 import Control.Monad
+import Data.Functor.Identity
 import Data.Char
 import Data.Int
 import Data.List
@@ -28,7 +29,7 @@ data CondCode = CCA | CCAE | CCB | CCBE | CCC | CCE | CCG | CCGE | CCL | CCLE |
   deriving (Show, Eq)
 
 data Ins
-    = MOV RegMem RegMem | MOVi RegMem Imm | MOVi64 Reg Imm
+    = MOV RegMem RegMemImm | MOVi Reg Imm
     | MOVSX Reg RegMem
     | ADD RegMem RegMemImm
     | SUB RegMem RegMemImm
@@ -80,9 +81,8 @@ verify :: Asm -> Either String ()
 verify (Asm funcs) = mapM_ (\(_, inss) -> mapM_ goI inss) funcs
   where
     goI :: Ins -> Either String ()
-    goI (MOV (RegMem a) (RegMem b)) = ckRegMem a >> ckRegMem b >> ck2mem a b >> ckSizes a b
-    goI (MOVi (RegMem a) (Imm b)) = ckRegMem a >> ckImm b >> ckSizes a b
-    goI (MOVi64 (Reg a) (Imm b)) = ckReg a >> ckImm b >> ckSizes64 a b
+    goI (MOV (RegMem a) (RegMemImm b)) = ckRegMem a >> ckRegMemImm b >> ck2mem a b >> ckSizes a b
+    goI (MOVi (Reg a) (Imm b)) = ckReg a >> ckImm b >> ckSizes64 a b
     goI (MOVSX (Reg a) (RegMem b)) = ckReg a >> ckRegMem b >> ckMovsx a b
     goI (ADD (RegMem a) (RegMemImm b)) = ckRegMem a >> ckRegMemImm b >> ck2mem a b >> ckSizes a b
     goI (SUB (RegMem a) (RegMemImm b)) = ckRegMem a >> ckRegMemImm b >> ck2mem a b >> ckSizes a b
@@ -178,7 +178,7 @@ instance Stringifiable XRef where
     stringify (XMem _ _ (mult, _) _ _) | not (mult `elem` [0,1,2,4,8]) =
         error $ "Register multiplier has invalid value " ++ show mult ++ " in XMem"
     stringify (XMem sz mr pair lab off) =
-        let res = intercalate "+" $ catMaybes [goR1 mr, goPair pair, goLab lab, goOff off]
+        let res = intercalate "+" (catMaybes [goR1 mr, goPair pair, goLab lab]) ++ goOff off
         in szword sz ++ " " ++ if null res then "[0]" else "[" ++ res ++ "]"
       where
         szword 1 = "byte"
@@ -191,8 +191,9 @@ instance Stringifiable XRef where
         goPair (0, _) = Nothing
         goPair (mult, r) = Just $ show mult ++ "*" ++ stringify (XReg 8 r)
         goLab = id
-        goOff 0 = Nothing
-        goOff o = Just $ show o
+        goOff o | o > 0 = '+' : show o
+                | o < 0 = show o
+                | otherwise = ""
 
     stringify (XImm imm) = show imm
 
@@ -227,7 +228,6 @@ instance Stringifiable CondCode where
 instance Stringifiable Ins where
     stringify (MOV a b) = "mov " ++ stringify a ++ ", " ++ stringify b
     stringify (MOVi a b) = "mov " ++ stringify a ++ ", " ++ stringify b
-    stringify (MOVi64 a b) = "mov " ++ stringify a ++ ", " ++ stringify b
     stringify (MOVSX a b@(RegMem bx)) = case compare (xrefGetSize bx) 4 of
         EQ -> "movsxd " ++ stringify a ++ ", " ++ stringify b
         LT -> "movsx " ++ stringify a ++ ", " ++ stringify b
@@ -280,3 +280,31 @@ isXMem _ = False
 isXImm :: XRef -> Bool
 isXImm (XImm _) = True
 isXImm _ = False
+
+xrefMapM :: Monad m => (XRef -> m XRef) -> Ins -> m Ins
+xrefMapM f (MOV (RegMem x) (RegMemImm y)) = MOV <$> (RegMem <$> f x) <*> (RegMemImm <$> f y)
+xrefMapM f (MOVi (Reg x) (Imm y)) = MOVi <$> (Reg <$> f x) <*> (Imm <$> f y)
+xrefMapM f (MOVSX (Reg x) (RegMem y)) = MOVSX <$> (Reg <$> f x) <*> (RegMem <$> f y)
+xrefMapM f (ADD (RegMem x) (RegMemImm y)) = ADD <$> (RegMem <$> f x) <*> (RegMemImm <$> f y)
+xrefMapM f (SUB (RegMem x) (RegMemImm y)) = SUB <$> (RegMem <$> f x) <*> (RegMemImm <$> f y)
+xrefMapM f (AND (RegMem x) (RegMemImm y)) = AND <$> (RegMem <$> f x) <*> (RegMemImm <$> f y)
+xrefMapM f (OR (RegMem x) (RegMemImm y)) = OR <$> (RegMem <$> f x) <*> (RegMemImm <$> f y)
+xrefMapM f (XOR (RegMem x) (RegMemImm y)) = XOR <$> (RegMem <$> f x) <*> (RegMemImm <$> f y)
+xrefMapM f (IMULDA (RegMem x)) = IMULDA <$> (RegMem <$> f x)
+xrefMapM f (IMUL (Reg x) (RegMem y)) = IMUL <$> (Reg <$> f x) <*> (RegMem <$> f y)
+xrefMapM f (IMUL3 (Reg x) (RegMem y) (Imm z)) = IMUL3<$>(Reg <$> f x)<*>(RegMem <$> f y)<*>(Imm <$> f z)
+xrefMapM f (MULDA (RegMem x)) = MULDA <$> (RegMem <$> f x)
+xrefMapM f (IDIVDA (RegMem x)) = IDIVDA <$> (RegMem <$> f x)
+xrefMapM f (DIVDA (RegMem x)) = DIVDA <$> (RegMem <$> f x)
+xrefMapM f (CMP (RegMem x) (RegMem y)) = CMP <$> (RegMem <$> f x) <*> (RegMem <$> f y)
+xrefMapM f (CMPi (RegMem x) (Imm y)) = CMPi <$> (RegMem <$> f x) <*> (Imm <$> f y)
+xrefMapM f (SETCC c (RegMem x)) = SETCC c <$> (RegMem <$> f x)
+xrefMapM _ i@(CALL _) = return i
+xrefMapM f (PUSH (RegMemImm x)) = PUSH <$> (RegMemImm <$> f x)
+xrefMapM f (POP (RegMem x)) = POP <$> (RegMem <$> f x)
+xrefMapM _ i@(JMP _) = return i
+xrefMapM _ i@(JCC _ _) = return i
+xrefMapM _ i@RET = return i
+
+xrefMap :: (XRef -> XRef) -> Ins -> Ins
+xrefMap f i = runIdentity $ xrefMapM (return . f) i
diff --git a/X64Optimiser.hs b/X64Optimiser.hs
index 746d88f..fa5d113 100644
--- a/X64Optimiser.hs
+++ b/X64Optimiser.hs
@@ -1,5 +1,7 @@
 module X64Optimiser(x64Optimise) where
 
+import Data.List
+
 import Defs
 import X64
 
@@ -7,6 +9,7 @@ import X64
 x64Optimise :: Asm -> Error Asm
 x64Optimise asm =
     return $
+    funcopt optCoalesceInstructions $
     optUnnecessaryJumps $
     funcopt optSimpleInstructions $
     asm
@@ -26,10 +29,50 @@ optSimpleInstructions :: Func -> Func
 optSimpleInstructions (name, inss) = (name, concat $ map goI inss)
   where
     goI :: Ins -> [Ins]
-    goI (MOV (RegMem a) (RegMem b)) | a == b = []
-    goI (MOVi (RegMem (XReg 8 r)) (Imm (XImm 0))) = [XOR (RegMem (XReg 4 r)) (RegMemImm (XReg 4 r))]
-    goI (MOVi (RegMem a@(XReg _ _)) (Imm (XImm 0))) = [XOR (RegMem a) (RegMemImm a)]
-    goI (MOVi64 (Reg (XReg 8 r)) (Imm (XImm 0))) = [XOR (RegMem (XReg 4 r)) (RegMemImm (XReg 4 r))]
-    goI (MOVi64 (Reg a) (Imm (XImm 0))) = [XOR (RegMem a) (RegMemImm a)]
+    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
diff --git a/putstr.lang b/putstr.lang
index 7485d8f..f132bda 100644
--- a/putstr.lang
+++ b/putstr.lang
@@ -17,32 +17,3 @@ func int main() {
 	putstr(str);
 	return 0;
 }
-
-
-
-/*
-irfunc putstr(char[] str)
-  {{{(0)
-    mov t5Q <- 0Q
-    jmp 7
-  }}}
-  {{{(7)
-    add t17Q <- t5Q, 8Q
-    add t18Q <- astrQ, t17Q
-    load t19B <- *t18Q
-    neq t22Q <- t19B, 0B
-    jne t22Q, 0Q -> 9 | 6
-  }}}
-  {{{(9)
-    add t30Q <- t5Q, 8Q
-    add t31Q <- astrQ, t30Q
-    load t32B <- *t31Q
-    call putc (t32B)
-    add t39Q <- t5Q, 1Q
-    mov t5Q <- t39Q
-    jmp 7
-  }}}
-  {{{(6)
-    ret
-  }}}
-*/