Better block ordering

1 files changed, 26 insertions, 47 deletions

diff --git a/Optimiser.hs b/Optimiser.hs
index 637ec4c..2923af6 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -1,6 +1,7 @@
 module Optimiser(optimise) where
 
 import Data.Either
+import Data.Function
 import Data.List
 import Data.Maybe
 import qualified Data.Map.Strict as Map
@@ -369,30 +370,32 @@ flipJccs (IRFunc rt name al bbs sid) = IRFunc rt name al (map goBB bbs) sid
     goT term = term
 
 reorderBlocks :: FuncOptimisation
-reorderBlocks (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
+reorderBlocks (IRFunc rt name al allbbs sid) = IRFunc rt name al resbbs sid
   where
-    bbs' = uncurry (buildResult [] (foldl foldfunc Map.empty bbs)) $ takeBlock sid bbs
-
-    foldfunc m (BB bid _ _) =
-        let candidates = map blockIdOf $ flip filter bbs $ \(BB bid' _ term') ->
-                term' `canJumpTo` bid && isNothing (Map.lookup bid' m)
-        in case candidates of
-            [cand] -> let m' = Map.insert cand bid m
-                      in if hasCycle m' cand then m else m'
-            _ -> m
-
-    buildResult :: [BB] -> Map.Map Id Id -> BB -> [BB] -> [BB]
-    buildResult pre _ bb [] = pre ++ [bb]
-    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
+    resbbs = buildResult (allChainsFrom allbbs sid) allbbs
+
+    allChains :: [BB] -> [[Id]]
+    allChains bbs = concatMap (allChainsFrom bbs . blockIdOf) bbs
+
+    allChainsFrom :: [BB] -> Id -> [[Id]]
+    allChainsFrom b start = go b [] start
+      where
+        go :: [BB] -> [Id] -> Id -> [[Id]]
+        go bbs chain at =
+            let ((BB _ _ term), rest) = takeBlock at bbs
+                chain' = chain ++ [at]
+            in case intersect (jumpTargets term) (map blockIdOf rest) of
+                   [] -> [chain']
+                   tgs -> concatMap (go rest chain') tgs
+
+    buildResult :: [[Id]] -> [BB] -> [BB]
+    buildResult _ [] = []
+    buildResult chains bbs =
+        let chain = maximumBy (compare `on` length) chains
+            (chainbbs', newbbs) = partition ((`elem` chain) . blockIdOf) bbs
+            chainbbs = sortBy (compare `on` (\(BB i _ _) -> fromJust $ findIndex (== i) chain)) chainbbs'
+            newchains = allChains newbbs
+        in chainbbs ++ buildResult newchains newbbs
 
     takeBlock :: Id -> [BB] -> (BB, [BB])
     takeBlock _ [] = undefined
@@ -400,30 +403,6 @@ reorderBlocks (IRFunc rt name al bbs sid) = IRFunc rt name al bbs' sid
         | bid == target = (bb, rest)
         | otherwise = fmap (bb :) $ takeBlock target rest
 
-    findChainHead :: Ord a => Map.Map a a -> Maybe a
-    findChainHead mp = case Map.keys mp \\ Map.elems mp of
-        [] -> 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)
-        in go [] values
-      where
-        go _ [] = False
-        go l (x:xs) | l `contains` x = True
-                    | otherwise = go (x:l) xs
-
-    canJumpTo :: IRTerm -> Id -> Bool
-    canJumpTo (IJcc _ _ _ i1 i2) i = i1 == i || i2 == i
-    canJumpTo (IJmp i') i = i' == i
-    canJumpTo _ _ = False
-
     jumpTargets :: IRTerm -> [Id]
     jumpTargets (IJcc _ _ _ i2 i1) = [i1, i2]
     jumpTargets (IJmp i) = [i]