From ad683ea5fb455b6a321af9cdee5a313e9961def8 Mon Sep 17 00:00:00 2001
From: tomsmeding <tom.smeding@gmail.com>
Date: Sun, 20 Aug 2017 17:14:52 +0200
Subject: Fifth

---
 Optimiser.hs | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/Optimiser.hs b/Optimiser.hs
index c31b6dd..637ec4c 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -383,12 +383,16 @@ reorderBlocks (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
 
     buildResult :: [BB] -> Map.Map Id Id -> BB -> [BB] -> [BB]
     buildResult pre _ bb [] = pre ++ [bb]
-    buildResult pre succmap bb@(BB bid _ _) rest = case Map.lookup bid succmap of
-        Nothing ->   buildResult (pre ++ [bb]) succmap hbb hrest
-                         where i = fromMaybe (blockIdOf (head rest)) $ findChainHead succmap
-                               (hbb, hrest) = takeBlock i rest
-        Just next -> buildResult (pre ++ [bb]) (Map.delete bid succmap) hbb hrest
-                         where (hbb, hrest) = takeBlock next rest
+    buildResult pre succmap bb@(BB bid _ term) rest = case Map.lookup bid succmap of
+        Nothing ->   buildResult (pre ++ [bb]) succmap nbb nrest
+                         where (nbb, nrest) = takeBlock next rest
+                               next1 = fromMaybe (blockIdOf (head rest)) $ findChainHead succmap
+                               chlen = chainLength next1 succmap
+                               next = case intersect (jumpTargets term) (map blockIdOf rest) of
+                                          [] -> next1
+                                          (t:_) -> if chlen > 1 then next1 else t
+        Just next -> buildResult (pre ++ [bb]) (Map.delete bid succmap) nbb nrest
+                         where (nbb, nrest) = takeBlock next rest
 
     takeBlock :: Id -> [BB] -> (BB, [BB])
     takeBlock _ [] = undefined
@@ -401,6 +405,11 @@ reorderBlocks (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
         [] -> Nothing
         (x:_) -> Just x
 
+    chainLength :: Ord a => a -> Map.Map a a -> Int
+    chainLength v mp = case Map.lookup v mp of
+        Nothing -> 0
+        Just x -> 1 + chainLength x mp
+
     hasCycle :: (Show a, Ord a) => Map.Map a a -> a -> Bool
     hasCycle mp from =
         let values = map fromJust $ takeWhile isJust $ iterate (>>= (\x -> Map.lookup x mp)) (Just from)
@@ -415,6 +424,11 @@ reorderBlocks (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
     canJumpTo (IJmp i') i = i' == i
     canJumpTo _ _ = False
 
+    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
-- 
cgit v1.2.3-54-g00ecf