diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index 5dd19ed3c..3c1c57f87 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/containers/src/Data/Map/Internal.hs @@ -4125,25 +4125,24 @@ deleteFindMax t = case maxViewWithKey t of -- See Note [Iterator] in Data.Set.Internal -newtype Iterator k a = Iterator (Stack k a) - iterDown :: Map k a -> Stack k a -> Stack k a -iterDown Tip stk = stk iterDown (Bin _ kx x l r) stk = iterDown l (Push kx x r stk) +iterDown Tip stk = stk -iterator :: Map k a -> Iterator k a -iterator m = Iterator (iterDown m Nada) +-- Create an iterator from a Map. +iterator :: Map k a -> Stack k a +iterator m = iterDown m Nada -iterNext :: Iterator k a -> Maybe (StrictPair (KeyValue k a) (Iterator k a)) -iterNext (Iterator stk) = case stk of - Nada -> Nothing - Push kx x r stk' -> Just $! KeyValue kx x :*: Iterator (iterDown r stk') +-- Get the next key-value and the remaining iterator. +iterNext :: Stack k a -> Maybe (StrictPair (KeyValue k a) (Stack k a)) +iterNext (Push kx x r stk) = Just $! KeyValue kx x :*: iterDown r stk +iterNext Nada = Nothing {-# INLINE iterNext #-} -iterNull :: Iterator k a -> Bool -iterNull (Iterator stk) = case stk of - Nada -> True - Push _ _ _ _ -> False +-- Whether there are no more key-values in the iterator. +iterNull :: Stack k a -> Bool +iterNull (Push _ _ _ _) = False +iterNull Nada = True {-------------------------------------------------------------------- [balance l x r] balances two trees with value x. diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs index f1db71c76..2ce1a7683 100644 --- a/containers/src/Data/Set/Internal.hs +++ b/containers/src/Data/Set/Internal.hs @@ -1278,35 +1278,40 @@ foldl'Stack f = go -- Note [Iterator] -- ~~~~~~~~~~~~~~~ --- This is an efficient way to consume a Set one element at a time. --- Alternately, this may be done by toAscList. toAscList when consumed via --- List.foldr will rewrite to Set.foldr (thanks to rewrite rules), which is --- quite efficient. However, sometimes that is not possible, such as in the --- second arg of '==' or 'compare', where manifesting the list cons cells --- is unavoidable and makes things slower. +-- Iteration, using a Stack as an iterator, is an efficient way to consume a Set +-- one element at a time. Alternately, this may be done by toAscList. toAscList +-- when consumed via List.foldr will rewrite to Set.foldr (thanks to rewrite +-- rules), which is quite efficient. However, sometimes that is not possible, +-- such as in the second arg of '==' or 'compare', where manifesting the list +-- cons cells is unavoidable and makes things slower. -- -- Concretely, compare on Set Int using toAscList takes ~21% more time compared -- to using Iterator, on GHC 9.6.3. - -newtype Iterator a = Iterator (Stack a) +-- +-- The heart of this implementation is the `iterDown` function. It walks down +-- the left spine of the tree, pushing the value and right child on the stack, +-- until a Tip is reached. The next value is now at the top of the stack. To get +-- to the value after that, `iterDown` is called again with the right child and +-- the remaining stack. iterDown :: Set a -> Stack a -> Stack a -iterDown Tip stk = stk iterDown (Bin _ x l r) stk = iterDown l (Push x r stk) +iterDown Tip stk = stk -iterator :: Set a -> Iterator a -iterator s = Iterator (iterDown s Nada) +-- Create an iterator from a Set. +iterator :: Set a -> Stack a +iterator s = iterDown s Nada -iterNext :: Iterator a -> Maybe (StrictPair a (Iterator a)) -iterNext (Iterator stk) = case stk of - Nada -> Nothing - Push x r stk' -> Just $! x :*: Iterator (iterDown r stk') +-- Get the next element and the remaining iterator. +iterNext :: Stack a -> Maybe (StrictPair a (Stack a)) +iterNext (Push x r stk) = Just $! x :*: iterDown r stk +iterNext Nada = Nothing {-# INLINE iterNext #-} -iterNull :: Iterator a -> Bool -iterNull (Iterator stk) = case stk of - Nada -> True - Push _ _ _ -> False +-- Whether there are no more elements in the iterator. +iterNull :: Stack a -> Bool +iterNull (Push _ _ _) = False +iterNull Nada = True {-------------------------------------------------------------------- Eq