CmmUtils: get rid of insertBlock

`Hoopl.Graph` has almost exactly the same function, so let's use that.
Also, use `IntMap.alter` to make it more efficient.

Also switch `Hoopl` to use strict maps.

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

Test Plan: ./validate

Reviewers: bgamari, simonmar

Reviewed By: bgamari

Subscribers: dfeuer, rwbarton, thomie, carter

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

5 files changed, 19 insertions, 24 deletions

diff --git a/compiler/cmm/CmmProcPoint.hs b/compiler/cmm/CmmProcPoint.hs
index e3eb1dc45d..bef8f384b8 100644
--- a/compiler/cmm/CmmProcPoint.hs
+++ b/compiler/cmm/CmmProcPoint.hs
@@ -242,11 +242,11 @@ splitAtProcPoints dflags entry_label callPPs procPoints procMap
                   (CmmProc (TopInfo {info_tbls = info_tbls})
                            top_l _ g@(CmmGraph {g_entry=entry})) =
   do -- Build a map from procpoints to the blocks they reach
-     let addBlock
+     let add_block
              :: LabelMap (LabelMap CmmBlock)
              -> CmmBlock
              -> LabelMap (LabelMap CmmBlock)
-         addBlock graphEnv b =
+         add_block graphEnv b =
            case mapLookup bid procMap of
              Just ProcPoint -> add graphEnv bid bid b
              Just (ReachedBy set) ->
@@ -265,7 +265,7 @@ splitAtProcPoints dflags entry_label callPPs procPoints procMap
                          regSetToList $
                          expectJust "ppLiveness" $ mapLookup pp liveness
 
-     graphEnv <- return $ foldlGraphBlocks addBlock mapEmpty g
+     graphEnv <- return $ foldlGraphBlocks add_block mapEmpty g
 
      -- Build a map from proc point BlockId to pairs of:
      --  * Labels for their new procedures
@@ -330,7 +330,7 @@ splitAtProcPoints dflags entry_label callPPs procPoints procMap
                   -- replace branches to procpoints with branches to jumps
                   blockEnv'' = toBlockMap $ replaceBranches jumpEnv $ ofBlockMap ppId blockEnv'
                   -- add the jump blocks to the graph
-                  blockEnv''' = foldl' (flip insertBlock) blockEnv'' jumpBlocks
+                  blockEnv''' = foldl' (flip addBlock) blockEnv'' jumpBlocks
               let g' = ofBlockMap ppId blockEnv'''
               -- pprTrace "g' pre jumps" (ppr g') $ do
               return (mapInsert ppId g' newGraphEnv)
diff --git a/compiler/cmm/CmmUtils.hs b/compiler/cmm/CmmUtils.hs
index aff16b3a19..53dbcddfbb 100644
--- a/compiler/cmm/CmmUtils.hs
+++ b/compiler/cmm/CmmUtils.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE CPP, GADTs, RankNTypes #-}
+{-# LANGUAGE GADTs, RankNTypes #-}
 
 -----------------------------------------------------------------------------
 --
@@ -56,7 +56,7 @@ module CmmUtils(
         -- * Operations that probably don't belong here
         modifyGraph,
 
-        ofBlockMap, toBlockMap, insertBlock,
+        ofBlockMap, toBlockMap,
         ofBlockList, toBlockList, bodyToBlockList,
         toBlockListEntryFirst, toBlockListEntryFirstFalseFallthrough,
         foldlGraphBlocks, mapGraphNodes, revPostorder, mapGraphNodes1,
@@ -65,8 +65,6 @@ module CmmUtils(
         blockTicks
   ) where
 
-#include "HsVersions.h"
-
 import GhcPrelude
 
 import TyCon    ( PrimRep(..), PrimElemRep(..) )
@@ -78,11 +76,9 @@ import BlockId
 import CLabel
 import Outputable
 import DynFlags
-import Util
 import CodeGen.Platform
 
 import Data.Word
-import Data.Maybe
 import Data.Bits
 import Hoopl.Graph
 import Hoopl.Label
@@ -495,12 +491,6 @@ toBlockMap (CmmGraph {g_graph=GMany NothingO body NothingO}) = body
 ofBlockMap :: BlockId -> LabelMap CmmBlock -> CmmGraph
 ofBlockMap entry bodyMap = CmmGraph {g_entry=entry, g_graph=GMany NothingO bodyMap NothingO}
 
-insertBlock :: CmmBlock -> LabelMap CmmBlock -> LabelMap CmmBlock
-insertBlock block map =
-  ASSERT(isNothing $ mapLookup id map)
-  mapInsert id block map
-  where id = entryLabel block
-
 toBlockList :: CmmGraph -> [CmmBlock]
 toBlockList g = mapElems $ toBlockMap g
 
diff --git a/compiler/cmm/Hoopl/Collections.hs b/compiler/cmm/Hoopl/Collections.hs
index b8072b37a7..ef7de4a078 100644
--- a/compiler/cmm/Hoopl/Collections.hs
+++ b/compiler/cmm/Hoopl/Collections.hs
@@ -12,7 +12,7 @@ module Hoopl.Collections
 
 import GhcPrelude
 
-import qualified Data.IntMap as M
+import qualified Data.IntMap.Strict as M
 import qualified Data.IntSet as S
 
 import Data.List (foldl', foldl1')
@@ -66,6 +66,7 @@ class IsMap map where
   mapInsert :: KeyOf map -> a -> map a -> map a
   mapInsertWith :: (a -> a -> a) -> KeyOf map -> a -> map a -> map a
   mapDelete :: KeyOf map -> map a -> map a
+  mapAlter :: (Maybe a -> Maybe a) -> KeyOf map -> map a -> map a
 
   mapUnion :: map a -> map a -> map a
   mapUnionWithKey :: (KeyOf map -> a -> a -> a) -> map a -> map a -> map a
@@ -143,6 +144,7 @@ instance IsMap UniqueMap where
   mapInsert k v (UM m) = UM (M.insert k v m)
   mapInsertWith f k v (UM m) = UM (M.insertWith f k v m)
   mapDelete k (UM m) = UM (M.delete k m)
+  mapAlter f k (UM m) = UM (M.alter f k m)
 
   mapUnion (UM x) (UM y) = UM (M.union x y)
   mapUnionWithKey f (UM x) (UM y) = UM (M.unionWithKey f x y)
diff --git a/compiler/cmm/Hoopl/Graph.hs b/compiler/cmm/Hoopl/Graph.hs
index df1ebe3ec1..0142f70c76 100644
--- a/compiler/cmm/Hoopl/Graph.hs
+++ b/compiler/cmm/Hoopl/Graph.hs
@@ -20,6 +20,7 @@ module Hoopl.Graph
 
 
 import GhcPrelude
+import Util
 
 import Hoopl.Label
 import Hoopl.Block
@@ -52,13 +53,14 @@ emptyBody = mapEmpty
 bodyList :: Body' block n -> [(Label,block n C C)]
 bodyList body = mapToList body
 
-addBlock :: NonLocal thing
-         => thing C C -> LabelMap (thing C C)
-         -> LabelMap (thing C C)
-addBlock b body
-  | mapMember lbl body = error $ "duplicate label " ++ show lbl ++ " in graph"
-  | otherwise          = mapInsert lbl b body
-  where lbl = entryLabel b
+addBlock
+    :: (NonLocal block, HasDebugCallStack)
+    => block C C -> LabelMap (block C C) -> LabelMap (block C C)
+addBlock block body = mapAlter add lbl body
+  where
+    lbl = entryLabel block
+    add Nothing = Just block
+    add _ = error $ "duplicate label " ++ show lbl ++ " in graph"
 
 
 -- ---------------------------------------------------------------------------
diff --git a/compiler/cmm/Hoopl/Label.hs b/compiler/cmm/Hoopl/Label.hs
index 8096fab073..6eae115779 100644
--- a/compiler/cmm/Hoopl/Label.hs
+++ b/compiler/cmm/Hoopl/Label.hs
@@ -87,6 +87,7 @@ instance IsMap LabelMap where
   mapInsert (Label k) v (LM m) = LM (mapInsert k v m)
   mapInsertWith f (Label k) v (LM m) = LM (mapInsertWith f k v m)
   mapDelete (Label k) (LM m) = LM (mapDelete k m)
+  mapAlter f (Label k) (LM m) = LM (mapAlter f k m)
 
   mapUnion (LM x) (LM y) = LM (mapUnion x y)
   mapUnionWithKey f (LM x) (LM y) = LM (mapUnionWithKey (f . mkHooplLabel) x y)