mergeBlocks should not be quadratic

2 files changed, 38 insertions, 10 deletions

diff --git a/Intermediate.hs b/Intermediate.hs
index 0feeb13..62a5d4a 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -51,6 +51,9 @@ data Terminator
 bidOf :: BB -> Int
 bidOf (BB i _ _) = i
 
+inssOf :: BB -> [Instruction]
+inssOf (BB _ i _) = i
+
 termOf :: BB -> Terminator
 termOf (BB _ _ t) = t
 
diff --git a/Optimiser.hs b/Optimiser.hs
index b6dd902..b874fc6 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE TupleSections #-}
 module Optimiser(optimise) where
 
+import Data.Function (on)
 import Data.List
 import qualified Data.Map.Strict as Map
 import Data.Maybe
@@ -27,17 +28,32 @@ optimise prog =
 mergeBlocks :: [BB] -> [BB]
 mergeBlocks [] = []
 mergeBlocks allbbs@(BB startb _ _ : _) =
-    uncurry (++) (partition ((== startb) . bidOf) (go allbbs (length allbbs)))
+    let bbmap = Map.fromList [(bidOf bb, bb) | bb <- allbbs]
+        cfg = Map.fromList [(bidOf bb, outEdges bb) | bb <- allbbs]
+        revcfg = oppositeGraph cfg
+        resbbs = go bbmap cfg revcfg Set.empty (map bidOf allbbs)
+    in uncurry (++) (partition ((== startb) . bidOf) resbbs)
   where
-    go [] _ = []
-    go bbs 0 = bbs
-    go (bb@(BB bid inss term) : bbs) n = case partition (hasJumpTo bid . termOf) bbs of
-        ([], _) -> go (bbs ++ [bb]) (n - 1)
-        ([BB bid' inss' _], rest) -> go (BB bid' (inss' ++ inss) term : rest) n
-        _ -> go (bbs ++ [bb]) (n - 1)
-
-    hasJumpTo bid (IJmp a) = a == bid
-    hasJumpTo _ _ = False
+    go bbmap _ _ _ [] = Map.elems bbmap
+    go bbmap cfg revcfg seen (curid:rest)
+        | curid `Set.member` seen = go bbmap cfg revcfg seen rest
+        | otherwise =
+            let topid = walkBack cfg revcfg curid
+                (ids, bb') = walkForward bbmap cfg revcfg topid
+                bbmap' = Map.insert topid bb' (foldr Map.delete bbmap ids)
+                seen' = seen <> Set.fromList ids
+            in go bbmap' cfg revcfg seen' rest
+
+    walkBack cfg revcfg curid = fromMaybe curid $ do
+        [v] <- Map.lookup curid revcfg
+        [_] <- Map.lookup v cfg
+        return (walkBack cfg revcfg v)
+
+    walkForward bbmap cfg revcfg curid = fromMaybe ([curid], bbmap Map.! curid) $ do
+        [v] <- Map.lookup curid cfg
+        [_] <- Map.lookup v revcfg
+        let (ids, BB _ inss term) = walkForward bbmap cfg revcfg v
+        return (curid : ids, BB curid (inssOf (bbmap Map.! curid) ++ inss) term)
 
 mergeRets :: [BB] -> [BB]
 mergeRets bbs =
@@ -181,3 +197,12 @@ uniq :: Eq a => [a] -> [a]
 uniq (x:y:zs) | x == y = uniq (y:zs)
               | otherwise = x : uniq (y:zs)
 uniq l = l
+
+oppositeGraph :: (Show a, Ord a) => Map.Map a [a] -> Map.Map a [a]
+oppositeGraph graph =
+    let nodes = concat [k : vs | (k, vs) <- Map.assocs graph]
+        edges = map ((,) <$> fst . head <*> map snd)
+                . groupBy ((==) `on` fst)
+                . sortOn fst
+                $ [(to, from) | (from, tos) <- Map.assocs graph, to <- tos]
+    in Map.fromList (map (,[]) nodes ++ edges)