{-# LANGUAGE GADTs, BangPatterns, ScopedTypeVariables #-}

module GHC.Cmm.CommonBlockElim
  ( elimCommonBlocks
  )
where


import GHC.Prelude hiding (iterate, succ, unzip, zip)

import GHC.Cmm.BlockId
import GHC.Cmm
import GHC.Cmm.Utils
import GHC.Cmm.Switch (eqSwitchTargetWith)
import GHC.Cmm.ContFlowOpt

import GHC.Cmm.Dataflow.Block
import GHC.Cmm.Dataflow.Graph
import GHC.Cmm.Dataflow.Label
import GHC.Cmm.Dataflow.Collections
import Data.Functor.Classes (liftEq)
import Data.Maybe (mapMaybe)
import qualified Data.List as List
import Data.Word
import qualified Data.Map as M
import qualified GHC.Data.TrieMap as TM
import GHC.Types.Unique.FM
import GHC.Types.Unique
import Control.Arrow (first, second)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE

-- -----------------------------------------------------------------------------
-- Eliminate common blocks

-- 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.

-- 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.
--
-- 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).
--
-- 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

-- TODO: Use optimization fuel
elimCommonBlocks :: CmmGraph -> CmmGraph
elimCommonBlocks :: CmmGraph -> CmmGraph
elimCommonBlocks CmmGraph
g = LabelMap BlockId -> CmmGraph -> CmmGraph
replaceLabels LabelMap BlockId
env (CmmGraph -> CmmGraph) -> CmmGraph -> CmmGraph
forall a b. (a -> b) -> a -> b
$ LabelMap BlockId -> CmmGraph -> CmmGraph
copyTicks LabelMap BlockId
env CmmGraph
g
  where
     env :: LabelMap BlockId
env = LabelMap BlockId -> [[(Key, DistinctBlocks)]] -> LabelMap BlockId
iterate LabelMap BlockId
forall a. LabelMap a
forall (map :: * -> *) a. IsMap map => map a
mapEmpty [[(Key, DistinctBlocks)]]
blocks_with_key
     -- The order of blocks doesn't matter here. While we could use
     -- revPostorder which drops unreachable blocks this is done in
     -- ContFlowOpt already which runs before this pass. So we use
     -- toBlockList since it is faster.
     groups :: [DistinctBlocks]
groups = (CmmBlock -> Int) -> DistinctBlocks -> [DistinctBlocks]
forall a. (a -> Int) -> [a] -> [[a]]
groupByInt CmmBlock -> Int
hash_block (CmmGraph -> DistinctBlocks
toBlockList CmmGraph
g) :: [[CmmBlock]]
     blocks_with_key :: [[(Key, DistinctBlocks)]]
blocks_with_key = [ [ (CmmBlock -> Key
forall (e :: Extensibility). Block CmmNode e C -> Key
forall (thing :: Extensibility -> Extensibility -> *)
       (e :: Extensibility).
NonLocal thing =>
thing e C -> Key
successors CmmBlock
b, [CmmBlock
b]) | CmmBlock
b <- DistinctBlocks
bs] | DistinctBlocks
bs <- [DistinctBlocks]
groups]

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

-- The outer list groups by hash. We retain this grouping throughout.
iterate :: Subst -> [[(Key, DistinctBlocks)]] -> Subst
iterate :: LabelMap BlockId -> [[(Key, DistinctBlocks)]] -> LabelMap BlockId
iterate LabelMap BlockId
subst [[(Key, DistinctBlocks)]]
blocks
    | LabelMap BlockId -> Bool
forall a. LabelMap a -> Bool
forall (map :: * -> *) a. IsMap map => map a -> Bool
mapNull LabelMap BlockId
new_substs = LabelMap BlockId
subst
    | Bool
otherwise = LabelMap BlockId -> [[(Key, DistinctBlocks)]] -> LabelMap BlockId
iterate LabelMap BlockId
subst' [[(Key, DistinctBlocks)]]
updated_blocks
  where
    grouped_blocks :: [[(Key, NonEmpty DistinctBlocks)]]
    grouped_blocks :: [[(Key, NonEmpty DistinctBlocks)]]
grouped_blocks = ([(Key, DistinctBlocks)] -> [(Key, NonEmpty DistinctBlocks)])
-> [[(Key, DistinctBlocks)]] -> [[(Key, NonEmpty DistinctBlocks)]]
forall a b. (a -> b) -> [a] -> [b]
map [(Key, DistinctBlocks)] -> [(Key, NonEmpty DistinctBlocks)]
groupByLabel [[(Key, DistinctBlocks)]]
blocks

    merged_blocks :: [[(Key, DistinctBlocks)]]
    (LabelMap BlockId
new_substs, [[(Key, DistinctBlocks)]]
merged_blocks) = (LabelMap BlockId
 -> [(Key, NonEmpty DistinctBlocks)]
 -> (LabelMap BlockId, [(Key, DistinctBlocks)]))
-> LabelMap BlockId
-> [[(Key, NonEmpty DistinctBlocks)]]
-> (LabelMap BlockId, [[(Key, DistinctBlocks)]])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL ((LabelMap BlockId
 -> (Key, NonEmpty DistinctBlocks)
 -> (LabelMap BlockId, (Key, DistinctBlocks)))
-> LabelMap BlockId
-> [(Key, NonEmpty DistinctBlocks)]
-> (LabelMap BlockId, [(Key, DistinctBlocks)])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL LabelMap BlockId
-> (Key, NonEmpty DistinctBlocks)
-> (LabelMap BlockId, (Key, DistinctBlocks))
go) LabelMap BlockId
forall a. LabelMap a
forall (map :: * -> *) a. IsMap map => map a
mapEmpty [[(Key, NonEmpty DistinctBlocks)]]
grouped_blocks
      where
        go :: LabelMap BlockId
-> (Key, NonEmpty DistinctBlocks)
-> (LabelMap BlockId, (Key, DistinctBlocks))
go !LabelMap BlockId
new_subst1 (Key
k,NonEmpty DistinctBlocks
dbs) = (LabelMap BlockId
new_subst1 LabelMap BlockId -> LabelMap BlockId -> LabelMap BlockId
forall a. LabelMap a -> LabelMap a -> LabelMap a
forall (map :: * -> *) a. IsMap map => map a -> map a -> map a
`mapUnion` LabelMap BlockId
new_subst2, (Key
k,DistinctBlocks
db))
          where
            (LabelMap BlockId
new_subst2, DistinctBlocks
db) = LabelMap BlockId
-> NonEmpty DistinctBlocks -> (LabelMap BlockId, DistinctBlocks)
mergeBlockList LabelMap BlockId
subst NonEmpty DistinctBlocks
dbs

    subst' :: LabelMap BlockId
subst' = LabelMap BlockId
subst LabelMap BlockId -> LabelMap BlockId -> LabelMap BlockId
forall a. LabelMap a -> LabelMap a -> LabelMap a
forall (map :: * -> *) a. IsMap map => map a -> map a -> map a
`mapUnion` LabelMap BlockId
new_substs
    updated_blocks :: [[(Key, DistinctBlocks)]]
updated_blocks = ([(Key, DistinctBlocks)] -> [(Key, DistinctBlocks)])
-> [[(Key, DistinctBlocks)]] -> [[(Key, DistinctBlocks)]]
forall a b. (a -> b) -> [a] -> [b]
map (((Key, DistinctBlocks) -> (Key, DistinctBlocks))
-> [(Key, DistinctBlocks)] -> [(Key, DistinctBlocks)]
forall a b. (a -> b) -> [a] -> [b]
map ((Key -> Key) -> (Key, DistinctBlocks) -> (Key, DistinctBlocks)
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first ((BlockId -> BlockId) -> Key -> Key
forall a b. (a -> b) -> [a] -> [b]
map (LabelMap BlockId -> BlockId -> BlockId
lookupBid LabelMap BlockId
subst')))) [[(Key, DistinctBlocks)]]
merged_blocks

-- Combine two lists of blocks.
-- While they are internally distinct they can still share common blocks.
mergeBlocks :: Subst -> DistinctBlocks -> DistinctBlocks -> (Subst, DistinctBlocks)
mergeBlocks :: LabelMap BlockId
-> DistinctBlocks
-> DistinctBlocks
-> (LabelMap BlockId, DistinctBlocks)
mergeBlocks LabelMap BlockId
subst DistinctBlocks
existing DistinctBlocks
new = DistinctBlocks -> (LabelMap BlockId, DistinctBlocks)
go DistinctBlocks
new
  where
    go :: DistinctBlocks -> (LabelMap BlockId, DistinctBlocks)
go [] = (LabelMap BlockId
forall a. LabelMap a
forall (map :: * -> *) a. IsMap map => map a
mapEmpty, DistinctBlocks
existing)
    go (CmmBlock
b:DistinctBlocks
bs) = case (CmmBlock -> Bool) -> DistinctBlocks -> Maybe CmmBlock
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find ((BlockId -> BlockId -> Bool) -> CmmBlock -> CmmBlock -> Bool
eqBlockBodyWith (LabelMap BlockId -> BlockId -> BlockId -> Bool
eqBid LabelMap BlockId
subst) CmmBlock
b) DistinctBlocks
existing of
        -- This block is a duplicate. Drop it, and add it to the substitution
        Just CmmBlock
b' -> (LabelMap BlockId -> LabelMap BlockId)
-> (LabelMap BlockId, DistinctBlocks)
-> (LabelMap BlockId, DistinctBlocks)
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (KeyOf LabelMap -> BlockId -> LabelMap BlockId -> LabelMap BlockId
forall a. KeyOf LabelMap -> a -> LabelMap a -> LabelMap a
forall (map :: * -> *) a.
IsMap map =>
KeyOf map -> a -> map a -> map a
mapInsert (CmmBlock -> BlockId
forall (x :: Extensibility). Block CmmNode C x -> BlockId
forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> BlockId
entryLabel CmmBlock
b) (CmmBlock -> BlockId
forall (x :: Extensibility). Block CmmNode C x -> BlockId
forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> BlockId
entryLabel CmmBlock
b')) ((LabelMap BlockId, DistinctBlocks)
 -> (LabelMap BlockId, DistinctBlocks))
-> (LabelMap BlockId, DistinctBlocks)
-> (LabelMap BlockId, DistinctBlocks)
forall a b. (a -> b) -> a -> b
$ DistinctBlocks -> (LabelMap BlockId, DistinctBlocks)
go DistinctBlocks
bs
        -- This block is not a duplicate, keep it.
        Maybe CmmBlock
Nothing -> (DistinctBlocks -> DistinctBlocks)
-> (LabelMap BlockId, DistinctBlocks)
-> (LabelMap BlockId, DistinctBlocks)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (CmmBlock
bCmmBlock -> DistinctBlocks -> DistinctBlocks
forall a. a -> [a] -> [a]
:) ((LabelMap BlockId, DistinctBlocks)
 -> (LabelMap BlockId, DistinctBlocks))
-> (LabelMap BlockId, DistinctBlocks)
-> (LabelMap BlockId, DistinctBlocks)
forall a b. (a -> b) -> a -> b
$ DistinctBlocks -> (LabelMap BlockId, DistinctBlocks)
go DistinctBlocks
bs

mergeBlockList :: Subst -> NonEmpty DistinctBlocks -> (Subst, DistinctBlocks)
mergeBlockList :: LabelMap BlockId
-> NonEmpty DistinctBlocks -> (LabelMap BlockId, DistinctBlocks)
mergeBlockList LabelMap BlockId
subst (DistinctBlocks
b:|[DistinctBlocks]
bs) = LabelMap BlockId
-> DistinctBlocks
-> [DistinctBlocks]
-> (LabelMap BlockId, DistinctBlocks)
go LabelMap BlockId
forall a. LabelMap a
forall (map :: * -> *) a. IsMap map => map a
mapEmpty DistinctBlocks
b [DistinctBlocks]
bs
  where
    go :: LabelMap BlockId
-> DistinctBlocks
-> [DistinctBlocks]
-> (LabelMap BlockId, DistinctBlocks)
go !LabelMap BlockId
new_subst1 DistinctBlocks
b [] = (LabelMap BlockId
new_subst1, DistinctBlocks
b)
    go !LabelMap BlockId
new_subst1 DistinctBlocks
b1 (DistinctBlocks
b2:[DistinctBlocks]
bs) = LabelMap BlockId
-> DistinctBlocks
-> [DistinctBlocks]
-> (LabelMap BlockId, DistinctBlocks)
go LabelMap BlockId
new_subst DistinctBlocks
b [DistinctBlocks]
bs
      where
        (LabelMap BlockId
new_subst2, DistinctBlocks
b) =  LabelMap BlockId
-> DistinctBlocks
-> DistinctBlocks
-> (LabelMap BlockId, DistinctBlocks)
mergeBlocks LabelMap BlockId
subst DistinctBlocks
b1 DistinctBlocks
b2
        new_subst :: LabelMap BlockId
new_subst = LabelMap BlockId
new_subst1 LabelMap BlockId -> LabelMap BlockId -> LabelMap BlockId
forall a. LabelMap a -> LabelMap a -> LabelMap a
forall (map :: * -> *) a. IsMap map => map a -> map a -> map a
`mapUnion` LabelMap BlockId
new_subst2


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

-- Below here is mostly boilerplate: hashing blocks ignoring labels,
-- and comparing blocks modulo a label mapping.

-- To speed up comparisons, we hash each basic block modulo jump labels.
-- The hashing is a bit arbitrary (the numbers are completely arbitrary),
-- but it should be fast and good enough.

-- 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.

type HashCode = Int

hash_block :: CmmBlock -> HashCode
hash_block :: CmmBlock -> Int
hash_block CmmBlock
block =
  Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((CmmNode C O -> Word32 -> Word32, CmmNode O O -> Word32 -> Word32,
 CmmNode O C -> Word32 -> Word32)
-> forall (e :: Extensibility) (x :: Extensibility).
   Block CmmNode e x
   -> IndexedCO x Word32 Word32 -> IndexedCO e Word32 Word32
forall (n :: Extensibility -> Extensibility -> *) a b c.
(n C O -> b -> c, n O O -> b -> b, n O C -> a -> b)
-> forall (e :: Extensibility) (x :: Extensibility).
   Block n e x -> IndexedCO x a b -> IndexedCO e c b
foldBlockNodesB3 (CmmNode C O -> Word32 -> Word32
forall {p} {p}. p -> p -> p
hash_fst, CmmNode O O -> Word32 -> Word32
forall {x :: Extensibility}. CmmNode O x -> Word32 -> Word32
hash_mid, CmmNode O C -> Word32 -> Word32
forall {x :: Extensibility}. CmmNode O x -> Word32 -> Word32
hash_lst) CmmBlock
block (Word32
0 :: Word32) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. (Word32
0x7fffffff :: Word32))
  -- UniqFM doesn't like negative Ints
  where hash_fst :: p -> p -> p
hash_fst p
_ p
h = p
h
        hash_mid :: CmmNode O x -> Word32 -> Word32
hash_mid CmmNode O x
m Word32
h = CmmNode O x -> Word32
forall (x :: Extensibility). CmmNode O x -> Word32
hash_node CmmNode O x
m Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
h Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
1
        hash_lst :: CmmNode O x -> Word32 -> Word32
hash_lst CmmNode O x
m Word32
h = CmmNode O x -> Word32
forall (x :: Extensibility). CmmNode O x -> Word32
hash_node CmmNode O x
m Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
h Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
1

        hash_node :: CmmNode O x -> Word32
        hash_node :: forall (x :: Extensibility). CmmNode O x -> Word32
hash_node CmmNode O x
n | CmmNode O x -> Bool
forall (x :: Extensibility). CmmNode O x -> Bool
dont_care CmmNode O x
n = Word32
0 -- don't care
        hash_node (CmmAssign CmmReg
r CmmExpr
e) = CmmReg -> Word32
hash_reg CmmReg
r Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ CmmExpr -> Word32
hash_e CmmExpr
e
        hash_node (CmmStore CmmExpr
e CmmExpr
e' AlignmentSpec
_) = CmmExpr -> Word32
hash_e CmmExpr
e Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ CmmExpr -> Word32
hash_e CmmExpr
e'
        hash_node (CmmUnsafeForeignCall ForeignTarget
t [LocalReg]
_ [CmmExpr]
as) = ForeignTarget -> Word32
hash_tgt ForeignTarget
t Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ (CmmExpr -> Word32) -> [CmmExpr] -> Word32
forall {t :: * -> *} {t}.
Foldable t =>
(t -> Word32) -> t t -> Word32
hash_list CmmExpr -> Word32
hash_e [CmmExpr]
as
        hash_node (CmmBranch BlockId
_) = Word32
23 -- NB. ignore the label
        hash_node (CmmCondBranch CmmExpr
p BlockId
_ BlockId
_ Maybe Bool
_) = CmmExpr -> Word32
hash_e CmmExpr
p
        hash_node (CmmCall CmmExpr
e Maybe BlockId
_ [GlobalReg]
_ Int
_ Int
_ Int
_) = CmmExpr -> Word32
hash_e CmmExpr
e
        hash_node (CmmForeignCall ForeignTarget
t [LocalReg]
_ [CmmExpr]
_ BlockId
_ Int
_ Int
_ Bool
_) = ForeignTarget -> Word32
hash_tgt ForeignTarget
t
        hash_node (CmmSwitch CmmExpr
e SwitchTargets
_) = CmmExpr -> Word32
hash_e CmmExpr
e
        hash_node CmmNode O x
_ = [Char] -> Word32
forall a. HasCallStack => [Char] -> a
error [Char]
"hash_node: unknown Cmm node!"

        hash_reg :: CmmReg -> Word32
        hash_reg :: CmmReg -> Word32
hash_reg   (CmmLocal LocalReg
localReg) = LocalReg -> Word32
forall a. Uniquable a => a -> Word32
hash_unique LocalReg
localReg -- important for performance, see #10397
        hash_reg   (CmmGlobal GlobalReg
_)    = Word32
19

        hash_e :: CmmExpr -> Word32
        hash_e :: CmmExpr -> Word32
hash_e (CmmLit CmmLit
l) = CmmLit -> Word32
hash_lit CmmLit
l
        hash_e (CmmLoad CmmExpr
e CmmType
_ AlignmentSpec
_) = Word32
67 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ CmmExpr -> Word32
hash_e CmmExpr
e
        hash_e (CmmReg CmmReg
r) = CmmReg -> Word32
hash_reg CmmReg
r
        hash_e (CmmMachOp MachOp
_ [CmmExpr]
es) = (CmmExpr -> Word32) -> [CmmExpr] -> Word32
forall {t :: * -> *} {t}.
Foldable t =>
(t -> Word32) -> t t -> Word32
hash_list CmmExpr -> Word32
hash_e [CmmExpr]
es -- pessimal - no operator check
        hash_e (CmmRegOff CmmReg
r Int
i) = CmmReg -> Word32
hash_reg CmmReg
r Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Int -> Word32
cvt Int
i
        hash_e (CmmStackSlot Area
_ Int
_) = Word32
13

        hash_lit :: CmmLit -> Word32
        hash_lit :: CmmLit -> Word32
hash_lit (CmmInt Integer
i Width
_) = Integer -> Word32
forall a. Num a => Integer -> a
fromInteger Integer
i
        hash_lit (CmmFloat Rational
r Width
_) = Rational -> Word32
forall b. Integral b => Rational -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate Rational
r
        hash_lit (CmmVec [CmmLit]
ls) = (CmmLit -> Word32) -> [CmmLit] -> Word32
forall {t :: * -> *} {t}.
Foldable t =>
(t -> Word32) -> t t -> Word32
hash_list CmmLit -> Word32
hash_lit [CmmLit]
ls
        hash_lit (CmmLabel CLabel
_) = Word32
119 -- ugh
        hash_lit (CmmLabelOff CLabel
_ Int
i) = Int -> Word32
cvt (Int -> Word32) -> Int -> Word32
forall a b. (a -> b) -> a -> b
$ Int
199 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i
        hash_lit (CmmLabelDiffOff CLabel
_ CLabel
_ Int
i Width
_) = Int -> Word32
cvt (Int -> Word32) -> Int -> Word32
forall a b. (a -> b) -> a -> b
$ Int
299 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i
        hash_lit (CmmBlock BlockId
_) = Word32
191 -- ugh
        hash_lit (CmmLit
CmmHighStackMark) = Int -> Word32
cvt Int
313

        hash_tgt :: ForeignTarget -> Word32
hash_tgt (ForeignTarget CmmExpr
e ForeignConvention
_) = CmmExpr -> Word32
hash_e CmmExpr
e
        hash_tgt (PrimTarget CallishMachOp
_) = Word32
31 -- lots of these

        hash_list :: (t -> Word32) -> t t -> Word32
hash_list t -> Word32
f = (Word32 -> t -> Word32) -> Word32 -> t t -> Word32
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Word32
z t
x -> t -> Word32
f t
x Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
z) (Word32
0::Word32)

        cvt :: Int -> Word32
cvt = Integer -> Word32
forall a. Num a => Integer -> a
fromInteger (Integer -> Word32) -> (Int -> Integer) -> Int -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a. Integral a => a -> Integer
toInteger

        hash_unique :: Uniquable a => a -> Word32
        hash_unique :: forall a. Uniquable a => a -> Word32
hash_unique = Int -> Word32
cvt (Int -> Word32) -> (a -> Int) -> a -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Unique -> Int
getKey (Unique -> Int) -> (a -> Unique) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Unique
forall a. Uniquable a => a -> Unique
getUnique

-- | Ignore these node types for equality
dont_care :: CmmNode O x -> Bool
dont_care :: forall (x :: Extensibility). CmmNode O x -> Bool
dont_care CmmComment {}  = Bool
True
dont_care CmmTick {}     = Bool
True
dont_care CmmUnwind {}   = Bool
True
dont_care CmmNode O x
_other         = Bool
False

-- Utilities: equality and substitution on the graph.

-- Given a map ``subst'' from BlockID -> BlockID, we define equality.
eqBid :: LabelMap BlockId -> BlockId -> BlockId -> Bool
eqBid :: LabelMap BlockId -> BlockId -> BlockId -> Bool
eqBid LabelMap BlockId
subst BlockId
bid BlockId
bid' = LabelMap BlockId -> BlockId -> BlockId
lookupBid LabelMap BlockId
subst BlockId
bid BlockId -> BlockId -> Bool
forall a. Eq a => a -> a -> Bool
== LabelMap BlockId -> BlockId -> BlockId
lookupBid LabelMap BlockId
subst BlockId
bid'
lookupBid :: LabelMap BlockId -> BlockId -> BlockId
lookupBid :: LabelMap BlockId -> BlockId -> BlockId
lookupBid LabelMap BlockId
subst BlockId
bid = case KeyOf LabelMap -> LabelMap BlockId -> Maybe BlockId
forall a. KeyOf LabelMap -> LabelMap a -> Maybe a
forall (map :: * -> *) a.
IsMap map =>
KeyOf map -> map a -> Maybe a
mapLookup KeyOf LabelMap
BlockId
bid LabelMap BlockId
subst of
                        Just BlockId
bid  -> LabelMap BlockId -> BlockId -> BlockId
lookupBid LabelMap BlockId
subst BlockId
bid
                        Maybe BlockId
Nothing -> BlockId
bid

-- Middle nodes and expressions can contain BlockIds, in particular in
-- CmmStackSlot and CmmBlock, so we have to use a special equality for
-- these.
--
eqMiddleWith :: (BlockId -> BlockId -> Bool)
             -> CmmNode O O -> CmmNode O O -> Bool
eqMiddleWith :: (BlockId -> BlockId -> Bool) -> CmmNode O O -> CmmNode O O -> Bool
eqMiddleWith BlockId -> BlockId -> Bool
eqBid (CmmAssign CmmReg
r1 CmmExpr
e1) (CmmAssign CmmReg
r2 CmmExpr
e2)
  = CmmReg
r1 CmmReg -> CmmReg -> Bool
forall a. Eq a => a -> a -> Bool
== CmmReg
r2 Bool -> Bool -> Bool
&& (BlockId -> BlockId -> Bool) -> CmmExpr -> CmmExpr -> Bool
eqExprWith BlockId -> BlockId -> Bool
eqBid CmmExpr
e1 CmmExpr
e2
eqMiddleWith BlockId -> BlockId -> Bool
eqBid (CmmStore CmmExpr
l1 CmmExpr
r1 AlignmentSpec
_) (CmmStore CmmExpr
l2 CmmExpr
r2 AlignmentSpec
_)
  = (BlockId -> BlockId -> Bool) -> CmmExpr -> CmmExpr -> Bool
eqExprWith BlockId -> BlockId -> Bool
eqBid CmmExpr
l1 CmmExpr
l2 Bool -> Bool -> Bool
&& (BlockId -> BlockId -> Bool) -> CmmExpr -> CmmExpr -> Bool
eqExprWith BlockId -> BlockId -> Bool
eqBid CmmExpr
r1 CmmExpr
r2
eqMiddleWith BlockId -> BlockId -> Bool
eqBid (CmmUnsafeForeignCall ForeignTarget
t1 [LocalReg]
r1 [CmmExpr]
a1)
                   (CmmUnsafeForeignCall ForeignTarget
t2 [LocalReg]
r2 [CmmExpr]
a2)
  = ForeignTarget
t1 ForeignTarget -> ForeignTarget -> Bool
forall a. Eq a => a -> a -> Bool
== ForeignTarget
t2 Bool -> Bool -> Bool
&& [LocalReg]
r1 [LocalReg] -> [LocalReg] -> Bool
forall a. Eq a => a -> a -> Bool
== [LocalReg]
r2 Bool -> Bool -> Bool
&& (CmmExpr -> CmmExpr -> Bool) -> [CmmExpr] -> [CmmExpr] -> Bool
forall a b. (a -> b -> Bool) -> [a] -> [b] -> Bool
forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq ((BlockId -> BlockId -> Bool) -> CmmExpr -> CmmExpr -> Bool
eqExprWith BlockId -> BlockId -> Bool
eqBid) [CmmExpr]
a1 [CmmExpr]
a2
eqMiddleWith BlockId -> BlockId -> Bool
_ CmmNode O O
_ CmmNode O O
_ = Bool
False

eqExprWith :: (BlockId -> BlockId -> Bool)
           -> CmmExpr -> CmmExpr -> Bool
eqExprWith :: (BlockId -> BlockId -> Bool) -> CmmExpr -> CmmExpr -> Bool
eqExprWith BlockId -> BlockId -> Bool
eqBid = CmmExpr -> CmmExpr -> Bool
eq
 where
  CmmLit CmmLit
l1          eq :: CmmExpr -> CmmExpr -> Bool
`eq` CmmLit CmmLit
l2          = CmmLit -> CmmLit -> Bool
eqLit CmmLit
l1 CmmLit
l2
  CmmLoad CmmExpr
e1 CmmType
t1 AlignmentSpec
a1   `eq` CmmLoad CmmExpr
e2 CmmType
t2 AlignmentSpec
a2   = CmmType
t1 CmmType -> CmmType -> Bool
`cmmEqType` CmmType
t2 Bool -> Bool -> Bool
&& CmmExpr
e1 CmmExpr -> CmmExpr -> Bool
`eq` CmmExpr
e2 Bool -> Bool -> Bool
&& AlignmentSpec
a1AlignmentSpec -> AlignmentSpec -> Bool
forall a. Eq a => a -> a -> Bool
==AlignmentSpec
a2
  CmmReg CmmReg
r1          `eq` CmmReg CmmReg
r2          = CmmReg
r1CmmReg -> CmmReg -> Bool
forall a. Eq a => a -> a -> Bool
==CmmReg
r2
  CmmRegOff CmmReg
r1 Int
i1    `eq` CmmRegOff CmmReg
r2 Int
i2    = CmmReg
r1CmmReg -> CmmReg -> Bool
forall a. Eq a => a -> a -> Bool
==CmmReg
r2 Bool -> Bool -> Bool
&& Int
i1Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
i2
  CmmMachOp MachOp
op1 [CmmExpr]
es1  `eq` CmmMachOp MachOp
op2 [CmmExpr]
es2  = MachOp
op1MachOp -> MachOp -> Bool
forall a. Eq a => a -> a -> Bool
==MachOp
op2 Bool -> Bool -> Bool
&& (CmmExpr -> CmmExpr -> Bool) -> [CmmExpr] -> [CmmExpr] -> Bool
forall a b. (a -> b -> Bool) -> [a] -> [b] -> Bool
forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq CmmExpr -> CmmExpr -> Bool
eq [CmmExpr]
es1 [CmmExpr]
es2
  CmmStackSlot Area
a1 Int
i1 `eq` CmmStackSlot Area
a2 Int
i2 = Area -> Area -> Bool
eqArea Area
a1 Area
a2 Bool -> Bool -> Bool
&& Int
i1Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
i2
  CmmExpr
_e1                `eq` CmmExpr
_e2                = Bool
False

  eqLit :: CmmLit -> CmmLit -> Bool
eqLit (CmmBlock BlockId
id1) (CmmBlock BlockId
id2) = BlockId -> BlockId -> Bool
eqBid BlockId
id1 BlockId
id2
  eqLit CmmLit
l1 CmmLit
l2 = CmmLit
l1 CmmLit -> CmmLit -> Bool
forall a. Eq a => a -> a -> Bool
== CmmLit
l2

  eqArea :: Area -> Area -> Bool
eqArea Area
Old Area
Old = Bool
True
  eqArea (Young BlockId
id1) (Young BlockId
id2) = BlockId -> BlockId -> Bool
eqBid BlockId
id1 BlockId
id2
  eqArea Area
_ Area
_ = Bool
False

-- Equality on the body of a block, modulo a function mapping block
-- IDs to block IDs.
eqBlockBodyWith :: (BlockId -> BlockId -> Bool) -> CmmBlock -> CmmBlock -> Bool
eqBlockBodyWith :: (BlockId -> BlockId -> Bool) -> CmmBlock -> CmmBlock -> Bool
eqBlockBodyWith BlockId -> BlockId -> Bool
eqBid CmmBlock
block CmmBlock
block'
  {-
  | equal     = pprTrace "equal" (vcat [ppr block, ppr block']) True
  | otherwise = pprTrace "not equal" (vcat [ppr block, ppr block']) False
  -}
  = Bool
equal
  where (CmmNode C O
_,Block CmmNode O O
m,CmmNode O C
l)   = CmmBlock -> (CmmNode C O, Block CmmNode O O, CmmNode O C)
forall (n :: Extensibility -> Extensibility -> *).
Block n C C -> (n C O, Block n O O, n O C)
blockSplit CmmBlock
block
        nodes :: [CmmNode O O]
nodes     = (CmmNode O O -> Bool) -> [CmmNode O O] -> [CmmNode O O]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (CmmNode O O -> Bool) -> CmmNode O O -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmmNode O O -> Bool
forall (x :: Extensibility). CmmNode O x -> Bool
dont_care) (Block CmmNode O O -> [CmmNode O O]
forall (n :: Extensibility -> Extensibility -> *).
Block n O O -> [n O O]
blockToList Block CmmNode O O
m)
        (CmmNode C O
_,Block CmmNode O O
m',CmmNode O C
l') = CmmBlock -> (CmmNode C O, Block CmmNode O O, CmmNode O C)
forall (n :: Extensibility -> Extensibility -> *).
Block n C C -> (n C O, Block n O O, n O C)
blockSplit CmmBlock
block'
        nodes' :: [CmmNode O O]
nodes'    = (CmmNode O O -> Bool) -> [CmmNode O O] -> [CmmNode O O]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (CmmNode O O -> Bool) -> CmmNode O O -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmmNode O O -> Bool
forall (x :: Extensibility). CmmNode O x -> Bool
dont_care) (Block CmmNode O O -> [CmmNode O O]
forall (n :: Extensibility -> Extensibility -> *).
Block n O O -> [n O O]
blockToList Block CmmNode O O
m')

        equal :: Bool
equal = (CmmNode O O -> CmmNode O O -> Bool)
-> [CmmNode O O] -> [CmmNode O O] -> Bool
forall a b. (a -> b -> Bool) -> [a] -> [b] -> Bool
forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq ((BlockId -> BlockId -> Bool) -> CmmNode O O -> CmmNode O O -> Bool
eqMiddleWith BlockId -> BlockId -> Bool
eqBid) [CmmNode O O]
nodes [CmmNode O O]
nodes' Bool -> Bool -> Bool
&&
                (BlockId -> BlockId -> Bool) -> CmmNode O C -> CmmNode O C -> Bool
eqLastWith BlockId -> BlockId -> Bool
eqBid CmmNode O C
l CmmNode O C
l'


eqLastWith :: (BlockId -> BlockId -> Bool) -> CmmNode O C -> CmmNode O C -> Bool
eqLastWith :: (BlockId -> BlockId -> Bool) -> CmmNode O C -> CmmNode O C -> Bool
eqLastWith BlockId -> BlockId -> Bool
eqBid (CmmBranch BlockId
bid1) (CmmBranch BlockId
bid2) = BlockId -> BlockId -> Bool
eqBid BlockId
bid1 BlockId
bid2
eqLastWith BlockId -> BlockId -> Bool
eqBid (CmmCondBranch CmmExpr
c1 BlockId
t1 BlockId
f1 Maybe Bool
l1) (CmmCondBranch CmmExpr
c2 BlockId
t2 BlockId
f2 Maybe Bool
l2) =
  CmmExpr
c1 CmmExpr -> CmmExpr -> Bool
forall a. Eq a => a -> a -> Bool
== CmmExpr
c2 Bool -> Bool -> Bool
&& Maybe Bool
l1 Maybe Bool -> Maybe Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe Bool
l2 Bool -> Bool -> Bool
&& BlockId -> BlockId -> Bool
eqBid BlockId
t1 BlockId
t2 Bool -> Bool -> Bool
&& BlockId -> BlockId -> Bool
eqBid BlockId
f1 BlockId
f2
eqLastWith BlockId -> BlockId -> Bool
eqBid (CmmCall CmmExpr
t1 Maybe BlockId
c1 [GlobalReg]
g1 Int
a1 Int
r1 Int
u1) (CmmCall CmmExpr
t2 Maybe BlockId
c2 [GlobalReg]
g2 Int
a2 Int
r2 Int
u2) =
  CmmExpr
t1 CmmExpr -> CmmExpr -> Bool
forall a. Eq a => a -> a -> Bool
== CmmExpr
t2 Bool -> Bool -> Bool
&& (BlockId -> BlockId -> Bool)
-> Maybe BlockId -> Maybe BlockId -> Bool
forall a b. (a -> b -> Bool) -> Maybe a -> Maybe b -> Bool
forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq BlockId -> BlockId -> Bool
eqBid Maybe BlockId
c1 Maybe BlockId
c2 Bool -> Bool -> Bool
&& Int
a1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
a2 Bool -> Bool -> Bool
&& Int
r1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r2 Bool -> Bool -> Bool
&& Int
u1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
u2 Bool -> Bool -> Bool
&& [GlobalReg]
g1 [GlobalReg] -> [GlobalReg] -> Bool
forall a. Eq a => a -> a -> Bool
== [GlobalReg]
g2
eqLastWith BlockId -> BlockId -> Bool
eqBid (CmmSwitch CmmExpr
e1 SwitchTargets
ids1) (CmmSwitch CmmExpr
e2 SwitchTargets
ids2) =
  CmmExpr
e1 CmmExpr -> CmmExpr -> Bool
forall a. Eq a => a -> a -> Bool
== CmmExpr
e2 Bool -> Bool -> Bool
&& (BlockId -> BlockId -> Bool)
-> SwitchTargets -> SwitchTargets -> Bool
eqSwitchTargetWith BlockId -> BlockId -> Bool
eqBid SwitchTargets
ids1 SwitchTargets
ids2
eqLastWith BlockId -> BlockId -> Bool
_ CmmNode O C
_ CmmNode O C
_ = Bool
False

-- | 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.
copyTicks :: LabelMap BlockId -> CmmGraph -> CmmGraph
copyTicks :: LabelMap BlockId -> CmmGraph -> CmmGraph
copyTicks LabelMap BlockId
env CmmGraph
g
  | LabelMap BlockId -> Bool
forall a. LabelMap a -> Bool
forall (map :: * -> *) a. IsMap map => map a -> Bool
mapNull LabelMap BlockId
env = CmmGraph
g
  | Bool
otherwise   = BlockId -> LabelMap CmmBlock -> CmmGraph
ofBlockMap (CmmGraph -> BlockId
forall (n :: Extensibility -> Extensibility -> *).
GenCmmGraph n -> BlockId
g_entry CmmGraph
g) (LabelMap CmmBlock -> CmmGraph) -> LabelMap CmmBlock -> CmmGraph
forall a b. (a -> b) -> a -> b
$ (CmmBlock -> CmmBlock) -> LabelMap CmmBlock -> LabelMap CmmBlock
forall a b. (a -> b) -> LabelMap a -> LabelMap b
forall (map :: * -> *) a b. IsMap map => (a -> b) -> map a -> map b
mapMap CmmBlock -> CmmBlock
copyTo LabelMap CmmBlock
blockMap
  where -- Reverse block merge map
        blockMap :: LabelMap CmmBlock
blockMap = CmmGraph -> LabelMap CmmBlock
toBlockMap CmmGraph
g
        revEnv :: Map BlockId Key
revEnv = (Map BlockId Key -> KeyOf LabelMap -> BlockId -> Map BlockId Key)
-> Map BlockId Key -> LabelMap BlockId -> Map BlockId Key
forall b a. (b -> KeyOf LabelMap -> a -> b) -> b -> LabelMap a -> b
forall (map :: * -> *) b a.
IsMap map =>
(b -> KeyOf map -> a -> b) -> b -> map a -> b
mapFoldlWithKey Map BlockId Key -> KeyOf LabelMap -> BlockId -> Map BlockId Key
Map BlockId Key -> BlockId -> BlockId -> Map BlockId Key
forall {k} {a}. Ord k => Map k [a] -> a -> k -> Map k [a]
insertRev Map BlockId Key
forall k a. Map k a
M.empty LabelMap BlockId
env
        insertRev :: Map k [a] -> a -> k -> Map k [a]
insertRev Map k [a]
m a
k k
x = ([a] -> [a] -> [a]) -> k -> [a] -> Map k [a] -> Map k [a]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith (([a] -> [a]) -> [a] -> [a] -> [a]
forall a b. a -> b -> a
const (a
ka -> [a] -> [a]
forall a. a -> [a] -> [a]
:)) k
x [a
k] Map k [a]
m
        -- Copy ticks and scopes into the given block
        copyTo :: CmmBlock -> CmmBlock
copyTo CmmBlock
block = case BlockId -> Map BlockId Key -> Maybe Key
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (CmmBlock -> BlockId
forall (x :: Extensibility). Block CmmNode C x -> BlockId
forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> BlockId
entryLabel CmmBlock
block) Map BlockId Key
revEnv of
          Maybe Key
Nothing -> CmmBlock
block
          Just Key
ls -> (CmmBlock -> CmmBlock -> CmmBlock)
-> CmmBlock -> DistinctBlocks -> CmmBlock
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr CmmBlock -> CmmBlock -> CmmBlock
forall {x :: Extensibility}.
CmmBlock -> Block CmmNode C x -> Block CmmNode C x
copy CmmBlock
block (DistinctBlocks -> CmmBlock) -> DistinctBlocks -> CmmBlock
forall a b. (a -> b) -> a -> b
$ (BlockId -> Maybe CmmBlock) -> Key -> DistinctBlocks
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((BlockId -> LabelMap CmmBlock -> Maybe CmmBlock)
-> LabelMap CmmBlock -> BlockId -> Maybe CmmBlock
forall a b c. (a -> b -> c) -> b -> a -> c
flip KeyOf LabelMap -> LabelMap CmmBlock -> Maybe CmmBlock
BlockId -> LabelMap CmmBlock -> Maybe CmmBlock
forall a. KeyOf LabelMap -> LabelMap a -> Maybe a
forall (map :: * -> *) a.
IsMap map =>
KeyOf map -> map a -> Maybe a
mapLookup LabelMap CmmBlock
blockMap) Key
ls
        copy :: CmmBlock -> Block CmmNode C x -> Block CmmNode C x
copy CmmBlock
from Block CmmNode C x
to =
          let ticks :: [CmmTickish]
ticks = CmmBlock -> [CmmTickish]
blockTicks CmmBlock
from
              CmmEntry  BlockId
_   CmmTickScope
scp0        = CmmBlock -> CmmNode C O
forall (n :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
Block n C x -> n C O
firstNode CmmBlock
from
              (CmmEntry BlockId
lbl CmmTickScope
scp1, Block CmmNode O x
code) = Block CmmNode C x -> (CmmNode C O, Block CmmNode O x)
forall (n :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
Block n C x -> (n C O, Block n O x)
blockSplitHead Block CmmNode C x
to
          in BlockId -> CmmTickScope -> CmmNode C O
CmmEntry BlockId
lbl (CmmTickScope -> CmmTickScope -> CmmTickScope
combineTickScopes CmmTickScope
scp0 CmmTickScope
scp1) CmmNode C O -> Block CmmNode O x -> Block CmmNode C x
forall (n :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
n C O -> Block n O x -> Block n C x
`blockJoinHead`
             (CmmNode O O -> Block CmmNode O x -> Block CmmNode O x)
-> Block CmmNode O x -> [CmmNode O O] -> Block CmmNode O x
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr CmmNode O O -> Block CmmNode O x -> Block CmmNode O x
forall (n :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
n O O -> Block n O x -> Block n O x
blockCons Block CmmNode O x
code ((CmmTickish -> CmmNode O O) -> [CmmTickish] -> [CmmNode O O]
forall a b. (a -> b) -> [a] -> [b]
map CmmTickish -> CmmNode O O
CmmTick [CmmTickish]
ticks)

-- Group by [Label]
-- See Note [Compressed TrieMap] in GHC.Core.Map.Expr about the usage of GenMap.
groupByLabel :: [(Key, DistinctBlocks)] -> [(Key, NonEmpty DistinctBlocks)]
groupByLabel :: [(Key, DistinctBlocks)] -> [(Key, NonEmpty DistinctBlocks)]
groupByLabel =
  ListMap
  (GenMap LabelMap)
  (Key (ListMap (GenMap LabelMap)), NonEmpty DistinctBlocks)
-> [(Key (ListMap (GenMap LabelMap)), DistinctBlocks)]
-> [(Key (ListMap (GenMap LabelMap)), NonEmpty DistinctBlocks)]
forall {m :: * -> *} {a}.
TrieMap m =>
m (Key m, NonEmpty a) -> [(Key m, a)] -> [(Key m, NonEmpty a)]
go (ListMap (GenMap LabelMap) (Key, NonEmpty DistinctBlocks)
forall a. ListMap (GenMap LabelMap) a
forall (m :: * -> *) a. TrieMap m => m a
TM.emptyTM :: TM.ListMap (TM.GenMap LabelMap) (Key, NonEmpty DistinctBlocks))
    where
      go :: m (Key m, NonEmpty a) -> [(Key m, a)] -> [(Key m, NonEmpty a)]
go !m (Key m, NonEmpty a)
m [] = ((Key m, NonEmpty a)
 -> [(Key m, NonEmpty a)] -> [(Key m, NonEmpty a)])
-> m (Key m, NonEmpty a)
-> [(Key m, NonEmpty a)]
-> [(Key m, NonEmpty a)]
forall a b. (a -> b -> b) -> m a -> b -> b
forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
TM.foldTM (:) m (Key m, NonEmpty a)
m []
      go !m (Key m, NonEmpty a)
m ((Key m
k,a
v) : [(Key m, a)]
entries) = m (Key m, NonEmpty a) -> [(Key m, a)] -> [(Key m, NonEmpty a)]
go (Key m
-> XT (Key m, NonEmpty a)
-> m (Key m, NonEmpty a)
-> m (Key m, NonEmpty a)
forall b. Key m -> XT b -> m b -> m b
forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
TM.alterTM Key m
k XT (Key m, NonEmpty a)
adjust m (Key m, NonEmpty a)
m) [(Key m, a)]
entries
        where --k' = map (getKey . getUnique) k
              adjust :: XT (Key m, NonEmpty a)
adjust Maybe (Key m, NonEmpty a)
Nothing       = (Key m, NonEmpty a) -> Maybe (Key m, NonEmpty a)
forall a. a -> Maybe a
Just (Key m
k, a -> NonEmpty a
forall a. a -> NonEmpty a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
v)
              adjust (Just (Key m
_,NonEmpty a
vs)) = (Key m, NonEmpty a) -> Maybe (Key m, NonEmpty a)
forall a. a -> Maybe a
Just (Key m
k, a
v a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
NE.<| NonEmpty a
vs)

groupByInt :: (a -> Int) -> [a] -> [[a]]
groupByInt :: forall a. (a -> Int) -> [a] -> [[a]]
groupByInt a -> Int
f [a]
xs = UniqFM Int [a] -> [[a]]
forall key elt. UniqFM key elt -> [elt]
nonDetEltsUFM (UniqFM Int [a] -> [[a]]) -> UniqFM Int [a] -> [[a]]
forall a b. (a -> b) -> a -> b
$ (UniqFM Int [a] -> a -> UniqFM Int [a])
-> UniqFM Int [a] -> [a] -> UniqFM Int [a]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' UniqFM Int [a] -> a -> UniqFM Int [a]
go UniqFM Int [a]
forall key elt. UniqFM key elt
emptyUFM [a]
xs
   -- See Note [Unique Determinism and code generation]
  where
    go :: UniqFM Int [a] -> a -> UniqFM Int [a]
go UniqFM Int [a]
m a
x = (Maybe [a] -> Maybe [a]) -> UniqFM Int [a] -> Int -> UniqFM Int [a]
forall key elt.
Uniquable key =>
(Maybe elt -> Maybe elt) -> UniqFM key elt -> key -> UniqFM key elt
alterUFM Maybe [a] -> Maybe [a]
addEntry UniqFM Int [a]
m (a -> Int
f a
x)
      where
        addEntry :: Maybe [a] -> Maybe [a]
addEntry Maybe [a]
xs = [a] -> Maybe [a]
forall a. a -> Maybe a
Just ([a] -> Maybe [a]) -> [a] -> Maybe [a]
forall a b. (a -> b) -> a -> b
$! [a] -> ([a] -> [a]) -> Maybe [a] -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [a
x] (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) Maybe [a]
xs