Review suggestions: prefer toplevel functions to internal workers

haskell · Apr 3, 2024 · 4a15a7a · 4a15a7a
1 parent 2c6d4a0
commit 4a15a7a
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 42 deletions.
diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
@@ -1956,38 +1956,41 @@ withoutKeys m (Set.Bin _ k ls rs) = case splitMember k m of
 -- @
 partitionKeys :: forall k a. Ord k => Map k a -> Set k -> (Map k a, Map k a)
 partitionKeys xs ys =
-  case go xs ys of
+  case partitionKeysWorker xs ys of
     xs' :*: ys' -> (xs', ys')
-  where
-    go :: Map k a -> Set k -> StrictPair (Map k a) (Map k a)
-    go Tip                 _           = Tip :*: Tip
-    go m                   Set.Tip     = Tip :*: m
-    go m@(Bin _ k x lm rm) s@Set.Bin{} =
-      case b of
-        True  -> with :*: without
-          where
-            with    =
-              if lmWith `ptrEq` lm && rmWith `ptrEq` rm
-              then m
-              else link k x lmWith rmWith
-            without =
-              link2 lmWithout rmWithout
-        False -> with :*: without
-          where
-            with    = link2 lmWith rmWith
-            without =
-              if lmWithout `ptrEq` lm && rmWithout `ptrEq` rm
-              then m
-              else link k x lmWithout rmWithout
-        where
-          !(lmWith :*: lmWithout) = go lm ls'
-          !(rmWith :*: rmWithout) = go rm rs'
-
-          !(!ls', b, !rs') = Set.splitMember k s
 #if __GLASGOW_HASKELL__
 {-# INLINABLE partitionKeys #-}
 #endif
 
+partitionKeysWorker :: Ord k => Map k a -> Set k -> StrictPair (Map k a) (Map k a)
+partitionKeysWorker Tip                 _           = Tip :*: Tip
+partitionKeysWorker m                   Set.Tip     = Tip :*: m
+partitionKeysWorker m@(Bin _ k x lm rm) s@Set.Bin{} =
+  case b of
+    True  -> with :*: without
+      where
+        with    =
+          if lmWith `ptrEq` lm && rmWith `ptrEq` rm
+          then m
+          else link k x lmWith rmWith
+        without =
+          link2 lmWithout rmWithout
+    False -> with :*: without
+      where
+        with    = link2 lmWith rmWith
+        without =
+          if lmWithout `ptrEq` lm && rmWithout `ptrEq` rm
+          then m
+          else link k x lmWithout rmWithout
+    where
+      !(lmWith :*: lmWithout) = partitionKeysWorker lm ls'
+      !(rmWith :*: rmWithout) = partitionKeysWorker rm rs'
+
+      !(!ls', b, !rs') = Set.splitMember k s
+#if __GLASGOW_HASKELL__
+{-# INLINABLE partitionKeysWorker #-}
+#endif
+
 -- | \(O(n+m)\). Difference with a combining function.
 -- When two equal keys are
 -- encountered, the combining function is applied to the values of these keys.

diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs
@@ -1316,23 +1316,25 @@ splitS x (Bin _ y l r)
           EQ -> (l :*: r)
 {-# INLINABLE splitS #-}
 
+splitMemberS :: Ord a => a -> Set a -> StrictTriple (Set a) Bool (Set a)
+splitMemberS x = go
+  where
+    go Tip           = StrictTriple Tip False Tip
+    go (Bin _ y l r) = case compare x y of
+      LT -> let StrictTriple lt found gt = splitMemberS x l
+            in StrictTriple lt found (link y gt r)
+      GT -> let StrictTriple lt found gt = splitMemberS x r
+            in StrictTriple (link y l lt) found gt
+      EQ -> StrictTriple l True r
+#if __GLASGOW_HASKELL__
+{-# INLINABLE splitMemberS #-}
+#endif
+
 -- | \(O(\log n)\). Performs a 'split' but also returns whether the pivot
 -- element was found in the original set.
 splitMember :: Ord a => a -> Set a -> (Set a,Bool,Set a)
-splitMember k0 s = case go k0 s of
+splitMember k0 s = case splitMemberS k0 s of
   StrictTriple l b r -> (l, b, r)
-  where
-    go :: Ord a => a -> Set a -> StrictTriple (Set a) Bool (Set a)
-    go _ Tip = StrictTriple Tip False Tip
-    go x (Bin _ y l r)
-       = case compare x y of
-           LT -> let StrictTriple lt found gt = go x l
-                     !gt' = link y gt r
-                 in StrictTriple lt found gt'
-           GT -> let StrictTriple lt found gt = go x r
-                     !lt' = link y l lt
-                 in StrictTriple lt' found gt
-           EQ -> StrictTriple l True r
 #if __GLASGOW_HASKELL__
 {-# INLINABLE splitMember #-}
 #endif

diff --git a/containers/src/Utils/Containers/Internal/StrictTriple.hs b/containers/src/Utils/Containers/Internal/StrictTriple.hs
@@ -3,8 +3,6 @@
 {-# LANGUAGE Safe #-}
 #endif
 
-#include "containers.h"
-
 -- | A strict triple
 
 module Utils.Containers.Internal.StrictTriple (StrictTriple(..)) where