Hoopl: improve postorder calculation

- Fix the naming and comments to indicate that we are calculating
  *reverse* postorder (and not the standard postorder).

- Rewrite the calculation to avoid CPS code. I found it fairly
  difficult to understand and the new one seems faster (according to
  nofib, decreases compiler allocations by 0.2%)

- Remove `LabelsPtr`, which seems unnecessary and could be *really*
  confusing. For instance, previously:
  `postorder_dfs_from <block with label X>`
  and
  `postorder_dfs_from <label X>`
  would actually mean quite different things (and give different
  results).

- Change the `Dataflow` module to always use entry of the graph for
  reverse postorder calculation. This should be the only change in
  behavior of this commit.

  Previously, if the caller provided initial facts for some of the
  labels, we would use those labels for our postorder calculation.
  However, I don't think that's correct in general - if the initial
  facts did not contain the entry of the graph, we would never analyze
  the blocks reachable from the entry but unreachable from the labels
  provided with the initial facts. It seems that the only analysis that
  used this was proc-point analysis, which I think would always include
  the entry block (so I don't think there's any bug due to this).

Signed-off-by: Michal Terepeta <michal.terepeta@gmail.com>

Test Plan: ./validate

Reviewers: bgamari, simonmar

Reviewed By: simonmar

Subscribers: rwbarton, thomie, carter

Differential Revision: https://phabricator.haskell.org/D4464

2 files changed, 61 insertions, 86 deletions

diff --git a/compiler/cmm/Hoopl/Dataflow.hs b/compiler/cmm/Hoopl/Dataflow.hs
index 0b0434bb36..2538b70ee3 100644
--- a/compiler/cmm/Hoopl/Dataflow.hs
+++ b/compiler/cmm/Hoopl/Dataflow.hs
@@ -111,8 +111,7 @@ analyzeCmm dir lattice transfer cmmGraph initFact =
         blockMap =
             case hooplGraph of
                 GMany NothingO bm NothingO -> bm
-        entries = if mapNull initFact then [entry] else mapKeys initFact
-    in fixpointAnalysis dir lattice transfer entries blockMap initFact
+    in fixpointAnalysis dir lattice transfer entry blockMap initFact
 
 -- Fixpoint algorithm.
 fixpointAnalysis
@@ -120,16 +119,16 @@ fixpointAnalysis
        Direction
     -> DataflowLattice f
     -> TransferFun f
-    -> [Label]
+    -> Label
     -> LabelMap CmmBlock
     -> FactBase f
     -> FactBase f
-fixpointAnalysis direction lattice do_block entries blockmap = loop start
+fixpointAnalysis direction lattice do_block entry blockmap = loop start
   where
     -- Sorting the blocks helps to minimize the number of times we need to
     -- process blocks. For instance, for forward analysis we want to look at
     -- blocks in reverse postorder. Also, see comments for sortBlocks.
-    blocks     = sortBlocks direction entries blockmap
+    blocks     = sortBlocks direction entry blockmap
     num_blocks = length blocks
     block_arr  = {-# SCC "block_arr" #-} listArray (0, num_blocks - 1) blocks
     start      = {-# SCC "start" #-} IntSet.fromDistinctAscList
@@ -174,9 +173,8 @@ rewriteCmm dir lattice rwFun cmmGraph initFact = do
         blockMap1 =
             case hooplGraph of
                 GMany NothingO bm NothingO -> bm
-        entries = if mapNull initFact then [entry] else mapKeys initFact
     (blockMap2, facts) <-
-        fixpointRewrite dir lattice rwFun entries blockMap1 initFact
+        fixpointRewrite dir lattice rwFun entry blockMap1 initFact
     return (cmmGraph {g_graph = GMany NothingO blockMap2 NothingO}, facts)
 
 fixpointRewrite
@@ -184,16 +182,16 @@ fixpointRewrite
        Direction
     -> DataflowLattice f
     -> RewriteFun f
-    -> [Label]
+    -> Label
     -> LabelMap CmmBlock
     -> FactBase f
     -> UniqSM (LabelMap CmmBlock, FactBase f)
-fixpointRewrite dir lattice do_block entries blockmap = loop start blockmap
+fixpointRewrite dir lattice do_block entry blockmap = loop start blockmap
   where
     -- Sorting the blocks helps to minimize the number of times we need to
     -- process blocks. For instance, for forward analysis we want to look at
     -- blocks in reverse postorder. Also, see comments for sortBlocks.
-    blocks     = sortBlocks dir entries blockmap
+    blocks     = sortBlocks dir entry blockmap
     num_blocks = length blocks
     block_arr  = {-# SCC "block_arr_rewrite" #-}
                  listArray (0, num_blocks - 1) blocks
@@ -268,20 +266,15 @@ we'll propagate (x=4) to L4, and nuke the otherwise-good rewriting of L4.
 -- | Sort the blocks into the right order for analysis. This means reverse
 -- postorder for a forward analysis. For the backward one, we simply reverse
 -- that (see Note [Backward vs forward analysis]).
---
--- Note: We're using Hoopl's confusingly named `postorder_dfs_from` but AFAICS
--- it returns the *reverse* postorder of the blocks (it visits blocks in the
--- postorder and uses (:) to collect them, which gives the reverse of the
--- visitation order).
 sortBlocks
     :: NonLocal n
-    => Direction -> [Label] -> LabelMap (Block n C C) -> [Block n C C]
-sortBlocks direction entries blockmap =
+    => Direction -> Label -> LabelMap (Block n C C) -> [Block n C C]
+sortBlocks direction entry blockmap =
     case direction of
         Fwd -> fwd
         Bwd -> reverse fwd
   where
-    fwd = postorder_dfs_from blockmap entries
+    fwd = revPostorderFrom blockmap entry
 
 -- Note [Backward vs forward analysis]
 --
diff --git a/compiler/cmm/Hoopl/Graph.hs b/compiler/cmm/Hoopl/Graph.hs
index ca482ab4a8..df1ebe3ec1 100644
--- a/compiler/cmm/Hoopl/Graph.hs
+++ b/compiler/cmm/Hoopl/Graph.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE RankNTypes #-}
@@ -14,7 +15,7 @@ module Hoopl.Graph
     , labelsDefined
     , mapGraph
     , mapGraphBlocks
-    , postorder_dfs_from
+    , revPostorderFrom
     ) where
 
 
@@ -119,22 +120,10 @@ labelsDefined (GMany _ body x) = mapFoldlWithKey addEntry (exitLabel x) body
 
 ----------------------------------------------------------------
 
-class LabelsPtr l where
-  targetLabels :: l -> [Label]
-
-instance NonLocal n => LabelsPtr (n e C) where
-  targetLabels n = successors n
-
-instance LabelsPtr Label where
-  targetLabels l = [l]
-
-instance LabelsPtr LabelSet where
-  targetLabels = setElems
-
-instance LabelsPtr l => LabelsPtr [l] where
-  targetLabels = concatMap targetLabels
-
--- | This is the most important traversal over this data structure.  It drops
+-- | Returns a list of blocks reachable from the provided Labels in the reverse
+-- postorder.
+--
+-- This is the most important traversal over this data structure.  It drops
 -- unreachable code and puts blocks in an order that is good for solving forward
 -- dataflow problems quickly.  The reverse order is good for solving backward
 -- dataflow problems quickly.  The forward order is also reasonably good for
@@ -143,59 +132,52 @@ instance LabelsPtr l => LabelsPtr [l] where
 -- that you would need a more serious analysis, probably based on dominators, to
 -- identify loop headers.
 --
--- The ubiquity of 'postorder_dfs' is one reason for the ubiquity of the 'LGraph'
--- representation, when for most purposes the plain 'Graph' representation is
--- more mathematically elegant (but results in more complicated code).
---
--- Here's an easy way to go wrong!  Consider
+-- For forward analyses we want reverse postorder visitation, consider:
 -- @
 --      A -> [B,C]
 --      B -> D
 --      C -> D
 -- @
--- Then ordinary dfs would give [A,B,D,C] which has a back ref from C to D.
--- Better to get [A,B,C,D]
-
-
--- | Traversal: 'postorder_dfs' returns a list of blocks reachable
--- from the entry of enterable graph. The entry and exit are *not* included.
--- The list has the following property:
---
---      Say a "back reference" exists if one of a block's
---      control-flow successors precedes it in the output list
---
---      Then there are as few back references as possible
---
--- The output is suitable for use in
--- a forward dataflow problem.  For a backward problem, simply reverse
--- the list.  ('postorder_dfs' is sufficiently tricky to implement that
--- one doesn't want to try and maintain both forward and backward
--- versions.)
-
-postorder_dfs_from_except :: forall block e . (NonLocal block, LabelsPtr e)
-                          => LabelMap (block C C) -> e -> LabelSet -> [block C C]
-postorder_dfs_from_except blocks b visited =
- vchildren (get_children b) (\acc _visited -> acc) [] visited
- where
-   vnode :: block C C -> ([block C C] -> LabelSet -> a) -> [block C C] -> LabelSet -> a
-   vnode block cont acc visited =
-        if setMember id visited then
-            cont acc visited
-        else
-            let cont' acc visited = cont (block:acc) visited in
-            vchildren (get_children block) cont' acc (setInsert id visited)
-      where id = entryLabel block
-   vchildren :: forall a. [block C C] -> ([block C C] -> LabelSet -> a) -> [block C C] -> LabelSet -> a
-   vchildren bs cont acc visited = next bs acc visited
-      where next children acc visited =
-                case children of []     -> cont acc visited
-                                 (b:bs) -> vnode b (next bs) acc visited
-   get_children :: forall l. LabelsPtr l => l -> [block C C]
-   get_children block = foldr add_id [] $ targetLabels block
-   add_id id rst = case lookupFact id blocks of
-                      Just b -> b : rst
-                      Nothing -> rst
-
-postorder_dfs_from
-    :: (NonLocal block, LabelsPtr b) => LabelMap (block C C) -> b -> [block C C]
-postorder_dfs_from blocks b = postorder_dfs_from_except blocks b setEmpty
+-- Postorder: [D, C, B, A] (or [D, B, C, A])
+-- Reverse postorder: [A, B, C, D] (or [A, C, B, D])
+-- This matters for, e.g., forward analysis, because we want to analyze *both*
+-- B and C before we analyze D.
+revPostorderFrom
+  :: forall block.  (NonLocal block)
+  => LabelMap (block C C) -> Label -> [block C C]
+revPostorderFrom graph start = go start_worklist setEmpty []
+  where
+    start_worklist = lookup_for_descend start Nil
+
+    -- To compute the postorder we need to "visit" a block (mark as done)
+    -- *after* visiting all its successors. So we need to know whether we
+    -- already processed all successors of each block (and @NonLocal@ allows
+    -- arbitrary many successors). So we use an explicit stack with an extra bit
+    -- of information:
+    -- * @ConsTodo@ means to explore the block if it wasn't visited before
+    -- * @ConsMark@ means that all successors were already done and we can add
+    --   the block to the result.
+    --
+    -- NOTE: We add blocks to the result list in postorder, but we *prepend*
+    -- them (i.e., we use @(:)@), which means that the final list is in reverse
+    -- postorder.
+    go :: DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
+    go Nil                      !_           !result = result
+    go (ConsMark block rest)    !wip_or_done !result =
+        go rest wip_or_done (block : result)
+    go (ConsTodo block rest)    !wip_or_done !result
+        | entryLabel block `setMember` wip_or_done = go rest wip_or_done result
+        | otherwise =
+            let new_worklist =
+                    foldr lookup_for_descend
+                          (ConsMark block rest)
+                          (successors block)
+            in go new_worklist (setInsert (entryLabel block) wip_or_done) result
+
+    lookup_for_descend :: Label -> DfsStack (block C C) -> DfsStack (block C C)
+    lookup_for_descend label wl
+      | Just b <- mapLookup label graph = ConsTodo b wl
+      | otherwise =
+           error $ "Label that doesn't have a block?! " ++ show label
+
+data DfsStack a = ConsTodo a (DfsStack a) | ConsMark a (DfsStack a) | Nil