CmmCommonBlockElim.hs 16.9 KB
Newer Older
1
{-# LANGUAGE GADTs, BangPatterns #-}
2
module CmmCommonBlockElim
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
3 4 5 6 7
  ( elimCommonBlocks
  )
where


8 9
import GhcPrelude hiding (iterate, succ, unzip, zip)

10
import BlockId
11
import Cmm
12
import CmmUtils
13
import CmmSwitch (eqSwitchTargetWith)
Simon Marlow's avatar
Simon Marlow committed
14
import CmmContFlowOpt
15
-- import PprCmm ()
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
16

17 18 19 20
import Hoopl.Block
import Hoopl.Graph
import Hoopl.Label
import Hoopl.Collections
21
import Data.Bits
Peter Wortmann's avatar
Peter Wortmann committed
22
import Data.Maybe (mapMaybe)
Ian Lynagh's avatar
Ian Lynagh committed
23
import qualified Data.List as List
24
import Data.Word
Peter Wortmann's avatar
Peter Wortmann committed
25
import qualified Data.Map as M
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
26
import Outputable
Ben Gamari's avatar
Ben Gamari committed
27
import DynFlags (DynFlags)
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
28
import UniqFM
niteria's avatar
niteria committed
29 30
import UniqDFM
import qualified TrieMap as TM
31 32
import Unique
import Control.Arrow (first, second)
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
33

34 35 36
-- -----------------------------------------------------------------------------
-- Eliminate common blocks

dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
37 38 39 40 41 42 43 44 45 46 47
-- If two blocks are identical except for the label on the first node,
-- then we can eliminate one of the blocks. To ensure that the semantics
-- of the program are preserved, we have to rewrite each predecessor of the
-- eliminated block to proceed with the block we keep.

-- The algorithm iterates over the blocks in the graph,
-- checking whether it has seen another block that is equal modulo labels.
-- If so, then it adds an entry in a map indicating that the new block
-- is made redundant by the old block.
-- Otherwise, it is added to the useful blocks.

48 49 50 51 52 53 54
-- To avoid comparing every block with every other block repeatedly, we group
-- them by
--   * a hash of the block, ignoring labels (explained below)
--   * the list of outgoing labels
-- The hash is invariant under relabeling, so we only ever compare within
-- the same group of blocks.
--
55 56
-- The list of outgoing labels is updated as we merge blocks (that is why they
-- are not included in the hash, which we want to calculate only once).
57 58 59 60 61
--
-- All in all, two blocks should never be compared if they have different
-- hashes, and at most once otherwise. Previously, we were slower, and people
-- rightfully complained: #10397

dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
62
-- TODO: Use optimization fuel
Ben Gamari's avatar
Ben Gamari committed
63 64
elimCommonBlocks :: DynFlags -> CmmGraph -> CmmGraph
elimCommonBlocks dflags g = replaceLabels env $ copyTicks env g
65
  where
Ben Gamari's avatar
Ben Gamari committed
66 67
     env = iterate dflags mapEmpty blocks_with_key
     groups = groupByInt (hash_block dflags) (postorderDfs g)
68 69 70 71 72 73
     blocks_with_key = [ [ (successors b, [b]) | b <- bs] | bs <- groups]

-- Invariant: The blocks in the list are pairwise distinct
-- (so avoid comparing them again)
type DistinctBlocks = [CmmBlock]
type Key = [Label]
74
type Subst = LabelMap BlockId
75 76

-- The outer list groups by hash. We retain this grouping throughout.
Ben Gamari's avatar
Ben Gamari committed
77 78
iterate :: DynFlags -> Subst -> [[(Key, DistinctBlocks)]] -> Subst
iterate dflags subst blocks
79
    | mapNull new_substs = subst
Ben Gamari's avatar
Ben Gamari committed
80
    | otherwise = iterate dflags subst' updated_blocks
81 82 83
  where
    grouped_blocks :: [[(Key, [DistinctBlocks])]]
    grouped_blocks = map groupByLabel blocks
84

85
    merged_blocks :: [[(Key, DistinctBlocks)]]
Ben Gamari's avatar
Ben Gamari committed
86 87
    (new_substs, merged_blocks) =
        List.mapAccumL (List.mapAccumL go) mapEmpty grouped_blocks
88 89 90
      where
        go !new_subst1 (k,dbs) = (new_subst1 `mapUnion` new_subst2, (k,db))
          where
Ben Gamari's avatar
Ben Gamari committed
91
            (new_subst2, db) = mergeBlockList dflags subst dbs
92

93 94
    subst' = subst `mapUnion` new_substs
    updated_blocks = map (map (first (map (lookupBid subst')))) merged_blocks
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
95

Ben Gamari's avatar
Ben Gamari committed
96 97 98 99
mergeBlocks :: DynFlags -> Subst
            -> DistinctBlocks -> DistinctBlocks
            -> (Subst, DistinctBlocks)
mergeBlocks dflags subst existing new = go new
100 101
  where
    go [] = (mapEmpty, existing)
Ben Gamari's avatar
Ben Gamari committed
102 103 104 105 106 107 108 109 110 111 112
    go (b:bs) =
        case List.find (eqBlockBodyWith dflags (eqBid subst) b) existing of
          -- This block is a duplicate. Drop it, and add it to the substitution
          Just b' -> first (mapInsert (entryLabel b) (entryLabel b')) $ go bs
          -- This block is not a duplicate, keep it.
          Nothing -> second (b:) $ go bs

mergeBlockList :: DynFlags -> Subst -> [DistinctBlocks]
               -> (Subst, DistinctBlocks)
mergeBlockList _      _     [] = pprPanic "mergeBlockList" empty
mergeBlockList dflags subst (b:bs) = go mapEmpty b bs
113 114 115 116
  where
    go !new_subst1 b [] = (new_subst1, b)
    go !new_subst1 b1 (b2:bs) = go new_subst b bs
      where
Ben Gamari's avatar
Ben Gamari committed
117
        (new_subst2, b) =  mergeBlocks dflags subst b1 b2
118
        new_subst = new_subst1 `mapUnion` new_subst2
119 120 121 122 123 124 125


-- -----------------------------------------------------------------------------
-- Hashing and equality on blocks

-- Below here is mostly boilerplate: hashing blocks ignoring labels,
-- and comparing blocks modulo a label mapping.
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
126

127
-- To speed up comparisons, we hash each basic block modulo jump labels.
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
128 129
-- The hashing is a bit arbitrary (the numbers are completely arbitrary),
-- but it should be fast and good enough.
130

131 132 133 134
-- We want to get as many small buckets as possible, as comparing blocks is
-- expensive. So include as much as possible in the hash. Ideally everything
-- that is compared with (==) in eqBlockBodyWith.

135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171
{-
Note [Equivalence up to local registers in CBE]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

CBE treats two blocks which are equivalent up to alpha-renaming of locally-bound
local registers as equivalent. This was not always the case (see #14226) but is
quite important for effective CBE. For instance, consider the blocks,

    c2VZ: // global
        _c2Yd::I64 = _s2Se::I64 + 1;
        _s2Sx::I64 = _c2Yd::I64;
        _s2Se::I64 = _s2Sx::I64;
        goto c2TE;

    c2VY: // global
        _c2Yb::I64 = _s2Se::I64 + 1;
        _s2Sw::I64 = _c2Yb::I64;
        _s2Se::I64 = _s2Sw::I64;
        goto c2TE;

These clearly implement precisely the same logic, differing only register
naming. This happens quite often in the code produced by GHC.

This alpha-equivalence relation must be accounted for in two places:

 1. the block hash function (hash_block), which we use for approximate "binning"
 2. the exact block comparison function, which computes pair-wise equivalence

In (1) we maintain a de Bruijn numbering of each block's locally-bound local
registers and compute the hash relative to this numbering.

For (2) we maintain a substitution which maps the local registers of one block
onto those of the other. We then compare local registers modulo this
substitution.

-}

172 173
type HashCode = Int

174 175 176 177 178 179 180 181 182 183
type LocalRegEnv a = UniqFM a
type DeBruijn = Int

-- | Maintains a de Bruijn numbering of local registers bound within a block.
--
-- See Note [Equivalence up to local registers in CBE]
data HashEnv = HashEnv { localRegHashEnv :: !(LocalRegEnv DeBruijn)
                       , nextIndex       :: !DeBruijn
                       }

Ben Gamari's avatar
Ben Gamari committed
184 185
hash_block :: DynFlags -> CmmBlock -> HashCode
hash_block dflags block =
186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204
  --pprTrace "hash_block" (ppr (entryLabel block) $$ ppr hash)
  hash
  where hash_fst _ (env, h) = (env, h)
        hash_mid m (env, h) = let (env', h') = hash_node env m
                              in (env', h' + h `shiftL` 1)
        hash_lst m (env, h) = let (env', h') = hash_node env m
                              in (env', h' + h `shiftL` 1)

        hash =
            let (_, raw_hash) =
                    foldBlockNodesF3 (hash_fst, hash_mid, hash_lst)
                                     block
                                     (emptyEnv, 0 :: Word32)
                emptyEnv = HashEnv mempty 0
            in fromIntegral (raw_hash .&. (0x7fffffff :: Word32))
               -- UniqFM doesn't like negative Ints

        hash_node :: HashEnv -> CmmNode O x -> (HashEnv, Word32)
        hash_node env n =
Ben Gamari's avatar
Ben Gamari committed
205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222
            (env', hash)
          where
            hash =
              case n of
                n | dont_care n -> 0  -- don't care
                -- don't include register as it is a binding occurrence
                CmmAssign (CmmLocal _) e -> hash_e env e
                CmmAssign r e   -> hash_reg env r + hash_e env e
                CmmStore e e'   -> hash_e env e + hash_e env e'
                CmmUnsafeForeignCall t _ as
                                -> hash_tgt env t + hash_list (hash_e env) as
                CmmBranch _     ->  23 -- NB. ignore the label
                CmmCondBranch p _ _ _ -> hash_e env p
                CmmCall e _ _ _ _ _   -> hash_e env e
                CmmForeignCall t _ _ _ _ _ _ -> hash_tgt env t
                CmmSwitch e _   -> hash_e env e
                _               -> error "hash_node: unknown Cmm node!"
            env' = foldLocalRegsDefd dflags (flip bind_local_reg) env n
223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238

        hash_reg :: HashEnv -> CmmReg -> Word32
        hash_reg env (CmmLocal localReg)
          | Just idx <- lookupUFM (localRegHashEnv env) localReg
          = fromIntegral idx
          | otherwise
          = hash_unique localReg -- important for performance, see #10397
        hash_reg _  (CmmGlobal _)    = 19

        hash_e :: HashEnv -> CmmExpr -> Word32
        hash_e _   (CmmLit l) = hash_lit l
        hash_e env (CmmLoad e _) = 67 + hash_e env e
        hash_e env (CmmReg r) = hash_reg env r
        hash_e env (CmmMachOp _ es) = hash_list (hash_e env) es -- pessimal - no operator check
        hash_e env (CmmRegOff r i) = hash_reg env r + cvt i
        hash_e _   (CmmStackSlot _ _) = 13
239

240
        hash_lit :: CmmLit -> Word32
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
241 242
        hash_lit (CmmInt i _) = fromInteger i
        hash_lit (CmmFloat r _) = truncate r
243
        hash_lit (CmmVec ls) = hash_list hash_lit ls
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
244
        hash_lit (CmmLabel _) = 119 -- ugh
245 246
        hash_lit (CmmLabelOff _ i) = cvt $ 199 + i
        hash_lit (CmmLabelDiffOff _ _ i) = cvt $ 299 + i
247
        hash_lit (CmmBlock _) = 191 -- ugh
248
        hash_lit (CmmHighStackMark) = cvt 313
249

250 251 252
        hash_tgt :: HashEnv -> ForeignTarget -> Word32
        hash_tgt env (ForeignTarget e _) = hash_e env e
        hash_tgt _   (PrimTarget _) = 31 -- lots of these
253

254
        hash_list f = List.foldl' (\z x -> f x + z) (0::Word32)
255

256
        cvt = fromInteger . toInteger
Peter Wortmann's avatar
Peter Wortmann committed
257

258 259 260 261 262 263 264
        bind_local_reg :: LocalReg -> HashEnv -> HashEnv
        bind_local_reg reg env =
            env { localRegHashEnv =
                    addToUFM (localRegHashEnv env) reg (nextIndex env)
                , nextIndex = nextIndex env + 1
                }

265 266 267
        hash_unique :: Uniquable a => a -> Word32
        hash_unique = cvt . getKey . getUnique

Peter Wortmann's avatar
Peter Wortmann committed
268 269 270 271
-- | Ignore these node types for equality
dont_care :: CmmNode O x -> Bool
dont_care CmmComment {}  = True
dont_care CmmTick {}     = True
272
dont_care CmmUnwind {}   = True
Peter Wortmann's avatar
Peter Wortmann committed
273 274
dont_care _other         = False

dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
275 276 277
-- Utilities: equality and substitution on the graph.

-- Given a map ``subst'' from BlockID -> BlockID, we define equality.
278
eqBid :: LabelMap BlockId -> BlockId -> BlockId -> Bool
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
279
eqBid subst bid bid' = lookupBid subst bid == lookupBid subst bid'
280
lookupBid :: LabelMap BlockId -> BlockId -> BlockId
281
lookupBid subst bid = case mapLookup bid subst of
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
282 283 284
                        Just bid  -> lookupBid subst bid
                        Nothing -> bid

285 286 287 288 289
-- | Maps the local registers of one block to those of another
--
-- See Note [Equivalence up to local registers in CBE]
type LocalRegMapping = LocalRegEnv LocalReg

290 291 292 293
-- Middle nodes and expressions can contain BlockIds, in particular in
-- CmmStackSlot and CmmBlock, so we have to use a special equality for
-- these.
--
Ben Gamari's avatar
Ben Gamari committed
294 295
eqMiddleWith :: DynFlags
             -> (BlockId -> BlockId -> Bool)
296 297 298
             -> LocalRegMapping
             -> CmmNode O O -> CmmNode O O
             -> (LocalRegMapping, Bool)
Ben Gamari's avatar
Ben Gamari committed
299
eqMiddleWith dflags eqBid env a b =
300
  case (a, b) of
Ben Gamari's avatar
Ben Gamari committed
301 302
     -- registers aren't compared since they are binding occurrences
    (CmmAssign (CmmLocal _) e1,  CmmAssign (CmmLocal _) e2) ->
303 304 305 306 307 308
        let eq = eqExprWith eqBid env e1 e2
        in (env', eq)

    (CmmAssign r1 e1,  CmmAssign r2 e2) ->
        let eq = r1 == r2
              && eqExprWith eqBid env e1 e2
Ben Gamari's avatar
Ben Gamari committed
309
        in (env', eq)
310 311 312 313

    (CmmStore l1 r1,  CmmStore l2 r2) ->
        let eq = eqExprWith eqBid env l1 l2
              && eqExprWith eqBid env r1 r2
Ben Gamari's avatar
Ben Gamari committed
314
        in (env', eq)
315

Ben Gamari's avatar
Ben Gamari committed
316 317
     -- result registers aren't compared since they are binding occurrences
    (CmmUnsafeForeignCall t1 _ a1,  CmmUnsafeForeignCall t2 _ a2) ->
318
        let eq = t1 == t2
Ben Gamari's avatar
Ben Gamari committed
319
              && eqLists (eqExprWith eqBid env) a1 a2
Ben Gamari's avatar
Ben Gamari committed
320
        in (env', eq)
321 322

    _ -> (env, False)
Ben Gamari's avatar
Ben Gamari committed
323 324 325 326 327
  where
    env' = List.foldl' (\acc (ra,rb) -> addToUFM acc ra rb) emptyUFM
           $ List.zip defd_a defd_b
    defd_a = foldLocalRegsDefd dflags (flip (:)) [] a
    defd_b = foldLocalRegsDefd dflags (flip (:)) [] b
328

Ben Gamari's avatar
Ben Gamari committed
329 330 331 332 333
eqLists :: (a -> b -> Bool) -> [a] -> [b] -> Bool
eqLists f (a:as) (b:bs) = f a b && eqLists f as bs
eqLists _ []     []     = True
eqLists _ _      _      = False

334
eqExprWith :: (BlockId -> BlockId -> Bool)
335
           -> LocalRegMapping
Ben Gamari's avatar
Ben Gamari committed
336 337
           -> CmmExpr -> CmmExpr
           -> Bool
338
eqExprWith eqBid env = eq
339 340 341
 where
  CmmLit l1          `eq` CmmLit l2          = eqLit l1 l2
  CmmLoad e1 _       `eq` CmmLoad e2 _       = e1 `eq` e2
342 343
  CmmReg r1          `eq` CmmReg r2          = r1 `eqReg` r2
  CmmRegOff r1 i1    `eq` CmmRegOff r2 i2    = r1 `eqReg` r2 && i1==i2
344 345 346 347
  CmmMachOp op1 es1  `eq` CmmMachOp op2 es2  = op1==op2 && es1 `eqs` es2
  CmmStackSlot a1 i1 `eq` CmmStackSlot a2 i2 = eqArea a1 a2 && i1==i2
  _e1                `eq` _e2                = False

Ben Gamari's avatar
Ben Gamari committed
348
  xs `eqs` ys = eqLists eq xs ys
349

350 351 352 353 354 355
  -- See Note [Equivalence up to local registers in CBE]
  CmmLocal a `eqReg` CmmLocal b
    | Just a' <- lookupUFM env a
    = a' == b
  a `eqReg` b = a == b

356 357 358 359 360 361 362
  eqLit (CmmBlock id1) (CmmBlock id2) = eqBid id1 id2
  eqLit l1 l2 = l1 == l2

  eqArea Old Old = True
  eqArea (Young id1) (Young id2) = eqBid id1 id2
  eqArea _ _ = False

Simon Marlow's avatar
Simon Marlow committed
363 364
-- Equality on the body of a block, modulo a function mapping block
-- IDs to block IDs.
Ben Gamari's avatar
Ben Gamari committed
365 366 367 368
eqBlockBodyWith :: DynFlags
                -> (BlockId -> BlockId -> Bool)
                -> CmmBlock -> CmmBlock -> Bool
eqBlockBodyWith dflags eqBid block block'
369 370 371 372
  {-
  | equal     = pprTrace "equal" (vcat [ppr block, ppr block']) True
  | otherwise = pprTrace "not equal" (vcat [ppr block, ppr block']) False
  -}
Ben Gamari's avatar
Ben Gamari committed
373
  = equal
Simon Marlow's avatar
Simon Marlow committed
374
  where (_,m,l)   = blockSplit block
Peter Wortmann's avatar
Peter Wortmann committed
375
        nodes     = filter (not . dont_care) (blockToList m)
Simon Marlow's avatar
Simon Marlow committed
376
        (_,m',l') = blockSplit block'
Peter Wortmann's avatar
Peter Wortmann committed
377
        nodes'    = filter (not . dont_care) (blockToList m')
378

379 380 381 382 383 384 385 386 387
        eqMids :: LocalRegMapping -> [CmmNode O O] -> [CmmNode O O] -> Bool
        eqMids env (a:as) (b:bs)
          | eq = eqMids env' as bs
          where
            (env', eq) = eqMiddleWith dflags eqBid env a b
        eqMids env [] [] = eqLastWith eqBid env l l'
        eqMids _ _ _ = False

        equal = eqMids emptyUFM nodes nodes'
388 389 390 391 392


eqLastWith :: (BlockId -> BlockId -> Bool) -> LocalRegMapping
           -> CmmNode O C -> CmmNode O C -> Bool
eqLastWith eqBid env a b =
Ben Gamari's avatar
Ben Gamari committed
393 394 395 396 397 398 399 400 401 402 403 404 405 406
    case (a, b) of
      (CmmBranch bid1, CmmBranch bid2) -> eqBid bid1 bid2
      (CmmCondBranch c1 t1 f1 l1, CmmCondBranch c2 t2 f2 l2) ->
          eqExprWith eqBid env c1 c2 && l1 == l2 && eqBid t1 t2 && eqBid f1 f2
      (CmmCall t1 c1 g1 a1 r1 u1, CmmCall t2 c2 g2 a2 r2 u2) ->
             t1 == t2
          && eqMaybeWith eqBid c1 c2
          && a1 == a2 && r1 == r2 && u1 == u2 && g1 == g2
      (CmmSwitch e1 ids1, CmmSwitch e2 ids2) ->
          eqExprWith eqBid env e1 e2 && eqSwitchTargetWith eqBid ids1 ids2
      -- result registers aren't compared since they are binding occurrences
      (CmmForeignCall t1 _ a1 s1 ret_args1 ret_off1 intrbl1,
       CmmForeignCall t2 _ a2 s2 ret_args2 ret_off2 intrbl2) ->
             t1 == t2
Ben Gamari's avatar
Ben Gamari committed
407
          && eqLists (eqExprWith eqBid env) a1 a2
Ben Gamari's avatar
Ben Gamari committed
408 409 410 411 412
          && s1 == s2
          && ret_args1 == ret_args2
          && ret_off1 == ret_off2
          && intrbl1 == intrbl2
      _ -> False
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
413 414 415 416 417

eqMaybeWith :: (a -> b -> Bool) -> Maybe a -> Maybe b -> Bool
eqMaybeWith eltEq (Just e) (Just e') = eltEq e e'
eqMaybeWith _ Nothing Nothing = True
eqMaybeWith _ _ _ = False
Peter Wortmann's avatar
Peter Wortmann committed
418 419 420 421 422

-- | Given a block map, ensure that all "target" blocks are covered by
-- the same ticks as the respective "source" blocks. This not only
-- means copying ticks, but also adjusting tick scopes where
-- necessary.
423
copyTicks :: LabelMap BlockId -> CmmGraph -> CmmGraph
Peter Wortmann's avatar
Peter Wortmann committed
424 425
copyTicks env g
  | mapNull env = g
Peter Wortmann's avatar
Peter Wortmann committed
426 427 428
  | otherwise   = ofBlockMap (g_entry g) $ mapMap copyTo blockMap
  where -- Reverse block merge map
        blockMap = toBlockMap g
429 430
        revEnv = mapFoldlWithKey insertRev M.empty env
        insertRev m k x = M.insertWith (const (k:)) x [k] m
Peter Wortmann's avatar
Peter Wortmann committed
431 432
        -- Copy ticks and scopes into the given block
        copyTo block = case M.lookup (entryLabel block) revEnv of
Peter Wortmann's avatar
Peter Wortmann committed
433
          Nothing -> block
Peter Wortmann's avatar
Peter Wortmann committed
434 435 436 437 438 439 440
          Just ls -> foldr copy block $ mapMaybe (flip mapLookup blockMap) ls
        copy from to =
          let ticks = blockTicks from
              CmmEntry  _   scp0        = firstNode from
              (CmmEntry lbl scp1, code) = blockSplitHead to
          in CmmEntry lbl (combineTickScopes scp0 scp1) `blockJoinHead`
             foldr blockCons code (map CmmTick ticks)
441 442 443

-- Group by [Label]
groupByLabel :: [(Key, a)] -> [(Key, [a])]
niteria's avatar
niteria committed
444
groupByLabel = go (TM.emptyTM :: TM.ListMap UniqDFM a)
445
  where
niteria's avatar
niteria committed
446 447
    go !m [] = TM.foldTM (:) m []
    go !m ((k,v) : entries) = go (TM.alterTM k' adjust m) entries
448 449 450 451 452
      where k' = map getUnique k
            adjust Nothing       = Just (k,[v])
            adjust (Just (_,vs)) = Just (k,v:vs)


453
groupByInt :: (a -> Int) -> [a] -> [[a]]
niteria's avatar
niteria committed
454 455
groupByInt f xs = nonDetEltsUFM $ List.foldl' go emptyUFM xs
  -- See Note [Unique Determinism and code generation]
456
  where go m x = alterUFM (Just . maybe [x] (x:)) m (f x)