Simon Jakobi pushed to branch wip/sjakobi/udfm-placement at Glasgow Haskell Compiler / GHC Commits: e4fa09ab by Simon Jakobi at 2026-07-03T12:47:23+02:00 Add Word64Map.compareSize; use it in UniqDFM's sorting guards The guards in eltsUDFM/udfmToList need the map's size only to compare it against a small bound, but computed it with a full O(n) traversal. Add compareSize, a verbatim port of Data.IntMap.compareSize from containers (haskell/containers#1139), which stops traversing as soon as the outcome is decided, and use it for both guards: the singleton check becomes O(1) and the placement/mergesort dispatch O(min(n, i)). Assisted-by: Claude Fable 5 - - - - - 1fdedb00 by Simon Jakobi at 2026-07-03T12:53:57+02:00 Fill placementSort's array directly from the Word64Map eltsUDFM and udfmToList materialized an intermediate list (for udfmToList even one of freshly allocated TaggedVals) only for placementSort to consume it in a single strict fill pass. Instead pass placementSort the map itself plus a per-element projection, and fill the array with the new Word64Map.traverseWithKey_, which visits each entry with an applicative action without building any intermediate structure. Instantiated at ST, this compiles to a fill loop that allocates nothing beyond the array (verified in the Core dump). Assisted-by: Claude Fable 5 - - - - - b4022ce8 by Simon Jakobi at 2026-07-06T01:17:53+02:00 Fixes and wibbles - - - - - 5 changed files: - compiler/GHC/Data/Word64Map/Internal.hs - compiler/GHC/Data/Word64Map/Lazy.hs - compiler/GHC/Data/Word64Map/Strict.hs - compiler/GHC/Data/Word64Map/Strict/Internal.hs - compiler/GHC/Types/Unique/DFM.hs Changes: ===================================== compiler/GHC/Data/Word64Map/Internal.hs ===================================== @@ -72,6 +72,7 @@ module GHC.Data.Word64Map.Internal ( -- * Query , null , size + , compareSize , member , notMember , lookup @@ -169,6 +170,7 @@ module GHC.Data.Word64Map.Internal ( , map , mapWithKey , traverseWithKey + , traverseWithKey_ , traverseMaybeWithKey , mapAccum , mapAccumWithKey @@ -522,6 +524,8 @@ null _ = False -- > size empty == 0 -- > size (singleton 1 'a') == 1 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 +-- +-- See also: 'compareSize' size :: Word64Map a -> Int size = go 0 where @@ -529,6 +533,23 @@ size = go 0 go acc (Tip _ _) = 1 + acc go acc Nil = acc +-- | \(O(\min(n,c))\). Compare the number of entries in the map to an @Int@. +-- +-- @compareSize m c@ returns the same result as @compare ('size' m) c@ but is +-- more efficient when @c@ is smaller than the size of the map. +compareSize :: Word64Map a -> Int -> Ordering +compareSize Nil c0 = compare 0 c0 +compareSize _ c0 | c0 <= 0 = GT +compareSize t c0 = compare 0 (go t (c0 - 1)) + where + go (Bin _ _ _ _) 0 = -1 + go (Bin _ _ l r) c + | c' < 0 = c' + | otherwise = go r c' + where + c' = go l (c - 1) + go _ c = c -- Must be Tip (Nil is never a child of Bin) + -- | \(O(\min(n,W))\). Is the key a member of the map? -- -- > member 5 (fromList [(5,'a'), (3,'b')]) == True @@ -2500,6 +2521,18 @@ traverseWithKey f = go | otherwise = liftA2 (Bin p m) (go l) (go r) {-# INLINE traverseWithKey #-} +-- | \(O(n)\). Visit each key\/value pair in ascending key order, discarding +-- the results. +-- +-- Unlike @traverse_ . toList@, this builds no intermediate structure. +traverseWithKey_ :: Applicative t => (Key -> a -> t ()) -> Word64Map a -> t () +traverseWithKey_ f = go + where + go Nil = pure () + go (Tip k v) = f k v + go (Bin _ _ l r) = go l *> go r +{-# INLINE traverseWithKey_ #-} + -- | \(O(n)\). The function @'mapAccum'@ threads an accumulating -- argument through the map in ascending order of keys. -- ===================================== compiler/GHC/Data/Word64Map/Lazy.hs ===================================== @@ -113,6 +113,7 @@ module GHC.Data.Word64Map.Lazy ( -- ** Size , WM.null , size + , compareSize -- * Combine @@ -148,6 +149,7 @@ module GHC.Data.Word64Map.Lazy ( , WM.map , mapWithKey , traverseWithKey + , traverseWithKey_ , traverseMaybeWithKey , mapAccum , mapAccumWithKey ===================================== compiler/GHC/Data/Word64Map/Strict.hs ===================================== @@ -130,6 +130,7 @@ module GHC.Data.Word64Map.Strict ( -- ** Size , null , size + , compareSize -- * Combine @@ -165,6 +166,7 @@ module GHC.Data.Word64Map.Strict ( , map , mapWithKey , traverseWithKey + , traverseWithKey_ , traverseMaybeWithKey , mapAccum , mapAccumWithKey ===================================== compiler/GHC/Data/Word64Map/Strict/Internal.hs ===================================== @@ -132,6 +132,7 @@ module GHC.Data.Word64Map.Strict.Internal ( -- ** Size , null , size + , compareSize -- * Combine @@ -167,6 +168,7 @@ module GHC.Data.Word64Map.Strict.Internal ( , map , mapWithKey , traverseWithKey + , traverseWithKey_ , traverseMaybeWithKey , mapAccum , mapAccumWithKey @@ -322,12 +324,14 @@ import GHC.Data.Word64Map.Internal , spanAntitone , restrictKeys , size + , compareSize , split , splitLookup , splitRoot , toAscList , toDescList , toList + , traverseWithKey_ , union , unions , withoutKeys ===================================== compiler/GHC/Types/Unique/DFM.hs ===================================== @@ -97,9 +97,9 @@ import qualified GHC.Data.Word64Set as W -- If the client of the map performs operations on the map in deterministic -- order then `udfmToList` returns them in deterministic order. -- --- There is an implementation cost: each element is given a serial number +-- There is an implementation cost: each element is given an insertion tag -- as it is added, and functions like `udfmToList` or `eltsUDFM` order their --- results by this serial number (see +-- results by this tag (see -- Note [Cost of deterministic iteration]). So you should only use `UniqDFM` -- if you need the deterministic property. -- @@ -114,7 +114,7 @@ import qualified GHC.Data.Word64Set as W -- -- -- There's more than one way to implement this. The implementation here tags --- every value with the insertion time that can later be used to sort the +-- every value with its insertion tag that can later be used to sort the -- values when asked to convert to a list. -- -- An alternative would be to have @@ -133,11 +133,11 @@ import qualified GHC.Data.Word64Set as W -- -- may also be worth considering. Compare Dhall.Map in the dhall package. --- | A type of values tagged with insertion time +-- | A type of values carrying an insertion tag data TaggedVal val = TaggedVal !val - {-# UNPACK #-} !Int -- ^ insertion time + {-# UNPACK #-} !Int -- ^ insertion tag deriving stock (Data, Functor, Foldable, Traversable) taggedFst :: TaggedVal val -> val @@ -159,13 +159,25 @@ instance Eq val => Eq (TaggedVal val) where data UniqDFM key ele = UDFM !(M.Word64Map (TaggedVal ele)) -- A map where keys are Unique's values and - -- values are tagged with insertion time. - -- The invariant is that all the tags will - -- be distinct within a single map + -- values carry an insertion tag. {-# UNPACK #-} !Int -- Upper bound on the values' insertion - -- time. See Note [Overflow on plusUDFM] + -- tags. See Note [Overflow on plusUDFM] + -- See Note [UDFM invariants] deriving (Data, Functor) +{- Note [UDFM invariants] +~~~~~~~~~~~~~~~~~~~~~~~~~ +In a map (UDFM m ub): + + (a) The insertion tags of the elements of m are distinct. + (b) Every tag lies in [0, ub). + +Consequently ub >= size m. + +The tags determine the order of deterministic iteration (eltsUDFM, +udfmToList); see Note [Sorting a UDFM]. +-} + -- | Deterministic. -- -- See Note [Cost of deterministic iteration]. @@ -266,8 +278,8 @@ plusUDFM_CK f udfml@(UDFM _ i) udfmr@(UDFM _ j) -- Note [Overflow on plusUDFM] -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~ -- There are multiple ways of implementing plusUDFM. --- The main problem that needs to be solved is overlap on times of --- insertion between different keys in two maps. +-- The main problem that needs to be solved is overlap on insertion +-- tags between different keys in two maps. -- Consider: -- -- A = fromList [(a, (x, 1))] @@ -355,84 +367,108 @@ nonDetStrictFoldUDFM k z (UDFM m _i) = foldl' k' z m where k' acc (TaggedVal v _) = k v acc --- Note [Cost of deterministic iteration] --- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ --- Deterministic iteration orders elements by insertion tag, and any such --- ordering must inspect every element's tag before it can emit the first --- element. So even the head of the result costs a full traversal of the map --- plus -- on the main path -- the allocation of an O(n)-sized array (see --- Note [Sorting a UDFM]). Laziness in the result list only avoids allocating --- for elements that are never demanded; it does not make the iteration --- incremental. #27459 shows this cost biting in consumers that demanded only --- the head. --- --- So: to test for emptiness, use isNullUDFM rather than null on eltsUDFM; --- for order-oblivious queries, prefer short-circuiting anyUDFM/allUDFM; and --- if you don't need the deterministic order at all, use the nonDet functions. +{- Note [Cost of deterministic iteration] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Deterministic iteration -- eltsUDFM, udfmToList, and everything built on +them, such as foldUDFM and UniqDFM's Foldable instance -- orders elements +by insertion tag, and any such ordering must inspect every element's tag +before it can emit the first element. So even the head of the result costs +a full traversal of the map plus -- on the main path -- the allocation of +an O(n)-sized array (see Note [Sorting a UDFM]). Laziness in the result +list only avoids allocating for elements that are never demanded; it does +not make the iteration incremental. #27459 shows this cost biting in +consumers that demanded only the head. + +So: to test for emptiness, use isNullUDFM rather than null on eltsUDFM; +for order-oblivious queries, prefer short-circuiting anyUDFM/allUDFM; and +if you don't need the deterministic order at all, use the nonDet functions. +-} -- | Deterministic, in order of insertion. -- -- See Note [Sorting a UDFM] and Note [Cost of deterministic iteration]. eltsUDFM :: UniqDFM key elt -> [elt] -eltsUDFM (UDFM m i) - | n <= 1 = map taggedFst (M.elems m) - | usePlacement n i = placementSort i (M.elems m) - | otherwise = map taggedFst (sort_it m) - where n = M.size m +eltsUDFM (UDFM m ub) + | M.compareSize m 1 /= GT = map taggedFst (M.elems m) + | usePlacement m ub = placementSort ub (\_ tv -> tv) m + | otherwise = map taggedFst (sort_it m) sort_it :: M.Word64Map (TaggedVal elt) -> [TaggedVal elt] sort_it m = sortBy (compare `on` taggedSnd) (M.elems m) --- Note [Sorting a UDFM] --- ~~~~~~~~~~~~~~~~~~~~~ --- Deterministic iteration must order elements by insertion tag. Instead of a --- comparison sort -- the list mergesort behind sortBy allocates ~n*log n cons --- cells -- we exploit the invariant that in (UDFM m i) all tags are distinct --- Ints in [0, i): allocate an array of size i, write each element at --- @index = tag@, freeze, and read out in index order. That's O(i) work (which --- subsumes the O(n) fill, since distinct tags force n <= i), no comparisons, --- and the readout is lazy, so consumers that demand only a prefix pay almost --- nothing beyond the fill (but the fill itself is unavoidable; see --- Note [Cost of deterministic iteration]). --- --- Holes: slots whose tag never occurs keep the initial sentinel, a TaggedVal --- with tag -1. Real tags are non-negative, so the readout skips on tag < 0; --- the sentinel's value field is never touched (it is unsafeCoerced ()). --- --- This sorting method loses when i is much larger than n: i never shrinks --- (overwrites keep bumping it, delete/filter shrink n but not i). We compute --- n = M.size m (O(n), cheap next to either sort) and fall back to the --- mergesort when i > 4 * n. Maps built by plain insertion -- the common --- case -- have i == n. The guard also caps the fast path's O(i) at O(n). - -usePlacement :: Int -> Int -> Bool -usePlacement n i = i <= 4 * n - --- | Order a list of 'TaggedVal's by tag, by placing each at array index = --- its tag. --- --- The tags must be distinct and in @[0, i)@. --- See Note [Sorting a UDFM]. -placementSort :: forall r. Int -> [TaggedVal r] -> [r] -placementSort i tvs = runST (ST (\s0 -> - case newSmallArray i hole s0 of - (# s1, marr #) -> case fill marr tvs s1 of - s2 -> case unsafeFreezeSmallArray marr s2 of +{- Note [Sorting a UDFM] +~~~~~~~~~~~~~~~~~~~~~~~~ +Deterministic iteration must order elements by insertion tag. A comparison +sort is needlessly expensive: the list mergesort behind sortBy allocates +~n*log n cons cells, and in #27459 those allocations made up a large share +of the compiler's total allocation on the InstanceMatching perf tests. + +Instead we exploit the UDFM invariants (all tags in (UDFM m ub) are distinct +Ints in [0, ub); see Note [UDFM invariants]): allocate an array of size ub, +write each element at +@index = tag@, freeze, and read out in index order. The fill is O(n) work, +where n = size m; the readout is O(ub); and since ub >= n, the whole sort is +O(ub) -- with no comparisons. The readout is lazy, so consumers that demand +only a prefix pay almost nothing beyond the fill (but the fill itself is +unavoidable; see Note [Cost of deterministic iteration]). + +Holes: slots whose tag never occurs keep the initial sentinel, a TaggedVal +with tag -1. Real tags are non-negative, so the readout skips on tag < 0. +The sentinel's value field is never read, but something must fill it: not a +panic thunk -- TaggedVal's strict value field would force it as soon as the +readout inspects a hole's tag -- so we borrow the value of an arbitrary map +element. (Safe: the callers' guards send maps of size < 2 down a different +path, so the map is never empty here.) + +This sorting method loses when ub is much larger than n = M.size m: ub never +shrinks (overwrites keep bumping it, delete/filter shrink n but not ub). We +fall back to the mergesort when ub > 4 * n; the decision about which sort to +use is made by usePlacement. It checks the threshold with M.compareSize, +which stops traversing the map as soon as the outcome is decided (so the +guard costs O(min(n, ub)), not a full O(n) size computation). Maps built by +plain insertion -- the common case -- have ub == n. The guard also caps the +fast path's O(ub) at O(n). +-} + +-- | @ub <= 4 * size m@, computed without a full 'M.size' traversal. +usePlacement :: M.Word64Map a -> Int -> Bool +usePlacement m ub = M.compareSize m ceil_ub_div_4 /= LT + where + ceil_ub_div_4 = (ub + 3) `div` 4 -- ceil(ub/4): ub <= 4*n iff n >= ceil(ub/4) + +-- | Order the map's elements by tag, by placing @mk key elt@ at array index = +-- the tag of @elt@. +-- +-- The map must be non-empty, its tags distinct and in @[0, ub)@; @mk@ must +-- preserve the tag. See Note [Sorting a UDFM]. +placementSort :: forall e r. Int + -> (M.Key -> TaggedVal e -> TaggedVal r) + -> M.Word64Map (TaggedVal e) + -> [r] +{-# INLINE placementSort #-} -- specializes mk into the fill loop +placementSort ub mk m = runST (ST (\s0 -> + case newSmallArray ub hole s0 of + (# s1, marr #) -> case fill marr s1 of + (# s2, () #) -> case unsafeFreezeSmallArray marr s2 of (# s3, arr #) -> (# s3, readout arr 0 #))) where + -- The tag -1 marks unfilled slots; the value field is never read, but + -- it is strict, so it needs a real value of type r -- borrow one from + -- the map. See Note [Sorting a UDFM]. hole :: TaggedVal r - hole = TaggedVal (unsafeCoerce ()) (-1) + hole = case M.findMin m of + (k, tv) -> TaggedVal (taggedFst (mk k tv)) (-1) - fill :: SmallMutableArray s (TaggedVal r) -> [TaggedVal r] -> State# s -> State# s - fill _ [] s = s - fill marr (tv : tvs') s = - case writeSmallArray marr (taggedSnd tv) tv s of - s' -> fill marr tvs' s' + fill :: SmallMutableArray s (TaggedVal r) -> State# s -> (# State# s, () #) + fill marr s = case M.traverseWithKey_ write m of ST st -> st s + where + write k tv = ST (\s' -> + (# writeSmallArray marr (taggedSnd tv) (mk k tv) s', () #)) readout :: SmallArray (TaggedVal r) -> Int -> [r] readout arr j - | j >= i = [] + | j >= ub = [] | t < 0 = readout arr (j + 1) | otherwise = v : readout arr (j + 1) where TaggedVal v t = indexSmallArray arr j @@ -461,19 +497,17 @@ udfmRestrictKeysSet (UDFM val_set i) set = -- as this already incurs most of the cost of returning the full list. -- See Note [Cost of deterministic iteration]. udfmToList :: UniqDFM key elt -> [(Unique, elt)] -udfmToList (UDFM m i) - | n <= 1 = [ (mkUniqueGrimily k, taggedFst v) | (k, v) <- M.toList m ] +udfmToList (UDFM m ub) + | M.compareSize m 1 /= GT = + [ (mkUniqueGrimily k, taggedFst v) | (k, v) <- M.toList m ] -- Unlike eltsUDFM, this allocates a fresh TaggedVal + pair per element - -- before the sort. If it ever matters, a parallel Word64 array of - -- keys filled in the same pass would avoid the eager boxes. - | usePlacement n i = placementSort i - (M.foldrWithKey (\k tv rest -> - TaggedVal (mkUniqueGrimily k, taggedFst tv) (taggedSnd tv) : rest) [] m) + -- (they make up the result). + | usePlacement m ub = placementSort ub + (\k tv -> TaggedVal (mkUniqueGrimily k, taggedFst tv) (taggedSnd tv)) m | otherwise = [ (mkUniqueGrimily k, taggedFst v) | (k, v) <- sortBy (compare `on` (taggedSnd . snd)) $ M.toList m ] - where n = M.size m -- Determines whether two 'UniqDFM's contain the same keys. equalKeysUDFM :: UniqDFM key a -> UniqDFM key b -> Bool View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/1c645b9c8d5318d36dd37c23c4abb6a... -- View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/1c645b9c8d5318d36dd37c23c4abb6a... You're receiving this email because of your account on gitlab.haskell.org.