DListify `OccurAnal.oneShotGroup`

Also adds/renames combinators for DList.

4 files changed, 50 insertions, 25 deletions

diff --git a/compiler/simplCore/Exitify.hs b/compiler/simplCore/Exitify.hs
index 89c16e13c5..33a02dad08 100644
--- a/compiler/simplCore/Exitify.hs
+++ b/compiler/simplCore/Exitify.hs
@@ -154,7 +154,7 @@ exitifyRec in_scope pairs
     -- Case right hand sides are in tail-call position
     go captured (_, AnnCase scrut bndr ty alts) = do
         alts' <- forM alts $ \(dc, pats, rhs) -> do
-            rhs' <- go (DL.snoc captured bndr DL.++: pats) rhs
+            rhs' <- go (DL.snoc captured bndr DL.++. pats) rhs
             return (dc, pats, rhs')
         return $ Case (deAnnotate scrut) bndr ty alts'
 
@@ -163,7 +163,7 @@ exitifyRec in_scope pairs
         | AnnNonRec j rhs <- ann_bind
         , Just join_arity <- isJoinId_maybe j
         = do let (params, join_body) = collectNAnnBndrs join_arity rhs
-             join_body' <- go (captured DL.++: params) join_body
+             join_body' <- go (captured DL.++. params) join_body
              let rhs' = mkLams params join_body'
              body' <- go (DL.snoc captured j) body
              return $ Let (NonRec j rhs') body'
@@ -175,15 +175,15 @@ exitifyRec in_scope pairs
              pairs' <- forM pairs $ \(j,rhs) -> do
                  let join_arity = idJoinArity j
                      (params, join_body) = collectNAnnBndrs join_arity rhs
-                 join_body' <- go (captured DL.++: js DL.++: params) join_body
+                 join_body' <- go (captured DL.++. js DL.++. params) join_body
                  let rhs' = mkLams params join_body'
                  return (j, rhs')
-             body' <- go (captured DL.++: js) body
+             body' <- go (captured DL.++. js) body
              return $ Let (Rec pairs') body'
 
         -- normal Let, only the body is in tail-call position
         | otherwise
-        = do body' <- go (captured DL.++: bindersOf bind) body
+        = do body' <- go (captured DL.++. bindersOf bind) body
              return $ Let bind body'
       where bind = deAnnBind ann_bind
 
diff --git a/compiler/simplCore/OccurAnal.hs b/compiler/simplCore/OccurAnal.hs
index 500dc7a912..82776e0c4b 100644
--- a/compiler/simplCore/OccurAnal.hs
+++ b/compiler/simplCore/OccurAnal.hs
@@ -41,6 +41,7 @@ import Demand           ( argOneShots, argsOneShots )
 import Digraph          ( SCC(..), Node(..)
                         , stronglyConnCompFromEdgedVerticesUniq
                         , stronglyConnCompFromEdgedVerticesUniqR )
+import qualified DList as DL
 import Unique
 import UniqFM
 import UniqSet
@@ -76,11 +77,11 @@ occurAnalysePgm this_mod active_unf active_rule imp_rules binds
     init_env = initOccEnv { occ_rule_act = active_rule
                           , occ_unf_act  = active_unf }
 
-    (final_usage, occ_anald_binds) = go init_env binds
-    (_, occ_anald_glommed_binds)   = occAnalRecBind init_env TopLevel
-                                                    imp_rule_edges
-                                                    (flattenBinds occ_anald_binds)
-                                                    initial_uds
+    (final_usage, DL.toList -> occ_anald_binds) = go init_env binds
+    (_, occ_anald_glommed_binds)                = occAnalRecBind init_env TopLevel
+                                                                 imp_rule_edges
+                                                                 (flattenBinds occ_anald_binds)
+                                                                 initial_uds
           -- It's crucial to re-analyse the glommed-together bindings
           -- so that we establish the right loop breakers. Otherwise
           -- we can easily create an infinite loop (#9583 is an example)
@@ -99,11 +100,11 @@ occurAnalysePgm this_mod active_unf active_rule imp_rules binds
                                              `delVarSetList` ru_bndrs imp_rule
                             , arg <- ru_args imp_rule ]
 
-    go :: OccEnv -> [CoreBind] -> (UsageDetails, [CoreBind])
+    go :: OccEnv -> [CoreBind] -> (UsageDetails, DL.DList CoreBind)
     go _ []
-        = (initial_uds, [])
+        = (initial_uds, DL.empty)
     go env (bind:binds)
-        = (final_usage, bind' ++ binds')
+        = (final_usage, bind' DL..++ binds')
         where
            (bs_usage, binds')   = go env binds
            (final_usage, bind') = occAnalBind env TopLevel imp_rule_edges bind
@@ -2122,19 +2123,20 @@ oneShotGroup :: OccEnv -> [CoreBndr]
         -- the binder. This is useful to guide subsequent float-in/float-out tranformations
 
 oneShotGroup env@(OccEnv { occ_one_shots = ctxt }) bndrs
-  = go ctxt bndrs []
+  = fmap DL.toList (go ctxt bndrs DL.empty)
   where
+    go :: [OneShotInfo] -> [Var] -> DL.DList Id -> (OccEnv, DL.DList Var)
     go ctxt [] rev_bndrs
       = ( env { occ_one_shots = ctxt, occ_encl = OccVanilla }
-        , reverse rev_bndrs )
+        , rev_bndrs )
 
     go [] bndrs rev_bndrs
       = ( env { occ_one_shots = [], occ_encl = OccVanilla }
-        , reverse rev_bndrs ++ bndrs )
+        , DL.reverse rev_bndrs DL.++. bndrs )
 
     go ctxt@(one_shot : ctxt') (bndr : bndrs) rev_bndrs
-      | isId bndr = go ctxt' bndrs (bndr': rev_bndrs)
-      | otherwise = go ctxt  bndrs (bndr : rev_bndrs)
+      | isId bndr = go ctxt' bndrs (DL.cons bndr' rev_bndrs)
+      | otherwise = go ctxt  bndrs (DL.cons bndr rev_bndrs)
       where
         bndr' = updOneShotInfo bndr one_shot
                -- Use updOneShotInfo, not setOneShotInfo, as pre-existing
diff --git a/compiler/simplCore/SimplUtils.hs b/compiler/simplCore/SimplUtils.hs
index fbe462e000..0e0d6910bd 100644
--- a/compiler/simplCore/SimplUtils.hs
+++ b/compiler/simplCore/SimplUtils.hs
@@ -1417,7 +1417,7 @@ mkLam env bndrs body cont
         bad bndr = isCoVar bndr && bndr `elemVarSet` co_vars
 
     mkLam' dflags bndrs body@(Lam {})
-      = mkLam' dflags (bndrs DL.++: bndrs1) body1
+      = mkLam' dflags (bndrs DL.++. bndrs1) body1
       where
         (bndrs1, body1) = collectBinders body
 
diff --git a/compiler/utils/DList.hs b/compiler/utils/DList.hs
index 9505debf5e..2a8e94566c 100644
--- a/compiler/utils/DList.hs
+++ b/compiler/utils/DList.hs
@@ -37,7 +37,9 @@ module DList
   , snoc
   , append
   , (++)
-  , (++:)
+  , (++.)
+  , (.++)
+  , (.++.)
   , concat
   , replicate
   , list
@@ -46,11 +48,13 @@ module DList
   , unfoldr
   , foldr
   , map
+  , omap
+  , reverse
   , concatMap
   , concatMapA
   ) where
 
-import Prelude hiding (concat, foldr, map, head, tail, replicate, (++), concatMap)
+import Prelude hiding (concat, foldr, map, head, tail, replicate, (++), concatMap, reverse)
 import qualified Data.List as List
 import Control.Monad as M
 import Data.Function (on)
@@ -156,10 +160,21 @@ append xs ys = DL (unDL xs . unDL ys)
 infixr 5 ++
 
 -- | /O(1)/. Append a list to an existing DList.
-(++:)        :: DList a -> [a] -> DList a
-xs ++: ys    = DL (unDL xs . (List.++ ys))
-{-# INLINE (++:) #-}
-infixl 5 ++:
+(++.)        :: DList a -> [a] -> DList a
+xs ++. ys    = DL (unDL xs . (List.++ ys))
+{-# INLINE (++.) #-}
+infixl 5 ++.
+
+-- | /O(1)/. Append a list to an existing DList.
+(.++)        :: [a] -> DList a -> DList a
+xs .++ ys    = DL ((List.++ xs) . unDL ys)
+{-# INLINE (.++) #-}
+infixl 5 .++
+
+-- | /O(1)/. Make a DList that represents a list concatenation.
+(.++.) :: [a] -> [a] -> DList a
+xs .++. ys = DL ((List.++ xs) . (List.++ ys) )
+infixr 5 .++.
 
 -- | /O(spine)/. Concatenate dlists
 concat       :: [DList a] -> DList a
@@ -205,6 +220,14 @@ map          :: (a -> b) -> DList a -> DList b
 map f        = foldr (cons . f) empty
 {-# INLINE map #-}
 
+-- Monomorphic map over difference lists.
+omap         :: (a -> a) -> DList a -> DList a
+omap f xs    = DL (fmap f . unDL xs)
+
+-- | /O(1)/. reverse for difference lists.
+reverse :: DList a -> DList a
+reverse xs = DL (List.reverse . unDL xs)
+
 -- | /O(n)/. concatMap for difference lists.
 concatMap :: (Foldable t) => (a -> DList b) -> t a -> DList b
 concatMap f xs = F.foldr (append . f) empty xs