Introduce recursive worker to ‘splitMember’ to increase inlining chances

haskell · Oct 15, 2023 · 0e6e28d · 0e6e28d
1 parent 43c15f0
commit 0e6e28d
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 12 deletions.
diff --git a/containers-tests/containers-tests.cabal b/containers-tests/containers-tests.cabal
@@ -115,6 +115,7 @@ library
     Utils.Containers.Internal.PtrEquality
     Utils.Containers.Internal.State
     Utils.Containers.Internal.StrictMaybe
+    Utils.Containers.Internal.StrictTriple
 
   if impl(ghc)
     other-modules:

diff --git a/containers/containers.cabal b/containers/containers.cabal
@@ -79,6 +79,7 @@ Library
         Utils.Containers.Internal.StrictMaybe
         Utils.Containers.Internal.PtrEquality
         Utils.Containers.Internal.Coercions
+        Utils.Containers.Internal.StrictTriple
     if impl(ghc)
       other-modules:
         Utils.Containers.Internal.TypeError

diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
@@ -398,6 +398,7 @@ import qualified Data.Set.Internal as Set
 import Data.Set.Internal (Set)
 import Utils.Containers.Internal.PtrEquality (ptrEq)
 import Utils.Containers.Internal.StrictPair
+import Utils.Containers.Internal.StrictTriple
 import Utils.Containers.Internal.StrictMaybe
 import Utils.Containers.Internal.BitQueue
 #ifdef DEFINE_ALTERF_FALLBACK
@@ -3983,8 +3984,6 @@ splitMember k0 m = case go k0 m of
 {-# INLINABLE splitMember #-}
 #endif
 
-data StrictTriple a b c = StrictTriple !a !b !c
-
 {--------------------------------------------------------------------
   Utility functions that maintain the balance properties of the tree.
   All constructors assume that all values in [l] < [k] and all values

diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs
@@ -250,6 +250,7 @@ import qualified Data.Foldable as Foldable
 import Control.DeepSeq (NFData(rnf))
 
 import Utils.Containers.Internal.StrictPair
+import Utils.Containers.Internal.StrictTriple
 import Utils.Containers.Internal.PtrEquality
 
 #if __GLASGOW_HASKELL__
@@ -1318,16 +1319,20 @@ splitS x (Bin _ y l r)
 -- | \(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 _ Tip = (Tip, False, Tip)
-splitMember x (Bin _ y l r)
-   = case compare x y of
-       LT -> let (lt, found, gt) = splitMember x l
-                 !gt' = link y gt r
-             in (lt, found, gt')
-       GT -> let (lt, found, gt) = splitMember x r
-                 !lt' = link y l lt
-             in (lt', found, gt)
-       EQ -> (l, True, r)
+splitMember k0 s = case go 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
@@ -0,0 +1,24 @@
+-- |
+-- Module:     Utils.Containers.Internal.StrictTriple
+-- Copyright:  (c) Sergey Vinokurov 2023
+-- License:    Apache-2.0 (see LICENSE)
+-- Maintainer: [email protected]
+
+
+{-# LANGUAGE CPP #-}
+#if !defined(TESTING) && defined(__GLASGOW_HASKELL__)
+{-# LANGUAGE Safe #-}
+#endif
+
+#include "containers.h"
+
+-- | A strict triple
+
+module Utils.Containers.Internal.StrictTriple (StrictTriple(..)) where
+
+-- | The same as a regular Haskell tuple, but
+--
+-- @
+-- StrictTriple x y _|_ = StrictTriple x _|_ z = StrictTriple _|_ y z = _|_
+-- @
+data StrictTriple a b c = StrictTriple !a !b !c