Use unboxed Maybe for map lookups on GHC

containers-tests/benchmarks/IntMap.hs

@@ -25,25 +25,25 @@ import Utils.Random (shuffle)
 
 main = do
     let m     = M.fromAscList elems_hits   :: M.IntMap Int
-    let m'    = M.fromAscList elems_mid    :: M.IntMap Int
-    let m''   = M.fromAscList elems_most   :: M.IntMap Int
-    let m'''  = M.fromAscList elems_misses :: M.IntMap Int
-    let m'''' = M.fromAscList elems_mixed  :: M.IntMap Int
-        m_random = M.fromList elems_random
+        m_mid = M.fromAscList elems_mid :: M.IntMap Int
+        m_most = M.fromAscList elems_most :: M.IntMap Int
+        m_misses = M.fromAscList elems_misses :: M.IntMap Int
+        m_mixed = M.fromList elems_mixed :: M.IntMap Int
+        m_random = M.fromList elems_random :: M.IntMap ()
         s = S.fromList keys
         s_random2 = S.fromList keys_random2
     evaluate $
       rnf [elems_asc, elems_random, elems_randomDups, elems_fromListWorstCase]
-    evaluate $ rnf [m, m', m'', m''', m'''']
+    evaluate $ rnf [m, m_mid, m_most, m_misses, m_mixed]
     evaluate $ rnf m_random
     evaluate $ rnf [s, s_random2]
     evaluate $ rnf evens
     defaultMain
         [ bench "lookup_hits" $ whnf (lookup keys) m
-        , bench "lookup_half" $ whnf (lookup keys) m'
-        , bench "lookup_most" $ whnf (lookup keys) m''
-        , bench "lookup_misses" $ whnf (lookup keys'') m'''
-        , bench "lookup_mixed" $ whnf (lookup keys) m''''
+        , bench "lookup_half" $ whnf (lookup keys) m_mid
+        , bench "lookup_most" $ whnf (lookup keys) m_most
+        , bench "lookup_misses" $ whnf (lookup keys'') m_misses
+        , bench "lookup_mixed" $ whnf (lookup keys) m_mixed
         , bench "index" $ whnf (indexMany evens) m
         , bench "insert" $ whnf (ins elems) M.empty
         , bench "insertWith empty" $ whnf (insWith elems) M.empty

containers-tests/containers-tests.cabal

@@ -135,6 +135,7 @@ library
     Utils.Containers.Internal.State
     Utils.Containers.Internal.Strict
     Utils.Containers.Internal.EqOrdUtil
+    Utils.Containers.Internal.UnboxedMaybe
 
   if impl(ghc >= 8.6)
     ghc-options: -Werror

containers/containers.cabal

@@ -92,5 +92,6 @@ Library
         Utils.Containers.Internal.EqOrdUtil
         Utils.Containers.Internal.BitUtil
         Utils.Containers.Internal.BitQueue
+        Utils.Containers.Internal.UnboxedMaybe
 
     include-dirs: include

containers/src/Data/IntMap/Internal.hs

@@ -354,6 +354,8 @@ import Data.IntSet.Internal.IntTreeCommons
 import Utils.Containers.Internal.BitUtil (shiftLL, shiftRL, iShiftRL, wordSize)
 import Utils.Containers.Internal.Strict
   (StrictPair(..), StrictTriple(..), toPair)
+import Utils.Containers.Internal.UnboxedMaybe (UMaybe)
+import qualified Utils.Containers.Internal.UnboxedMaybe as UMaybe
 
 #ifdef __GLASGOW_HASKELL__
 import Data.Coerce
@@ -649,26 +651,25 @@ notMember k m = not $ member k m
 
 -- | \(O(\min(n,W))\). Look up the value at a key in the map. See also 'Data.Map.lookup'.
 
--- See Note: Local 'go' functions and capturing
 lookup :: Key -> IntMap a -> Maybe a
-lookup !k = go
-  where
-    go (Bin p l r) | left k p  = go l
-                   | otherwise = go r
-    go (Tip kx x) | k == kx   = Just x
-                  | otherwise = Nothing
-    go Nil = Nothing
+lookup k m = UMaybe.toMaybe (lookupGo k m)
+-- Inline so that GHC can eliminate the Maybe if it matched on at the call-site.
+{-# INLINE lookup #-}
+
+lookupGo :: Key -> IntMap a -> UMaybe a
+lookupGo !k (Bin p l r)
+  | left k p = lookupGo k l
+  | otherwise = lookupGo k r
+lookupGo !k (Tip kx x)
+  | k == kx = UMaybe.just x
+  | otherwise = UMaybe.nothing ()
+lookupGo !_ Nil = UMaybe.nothing ()
 
--- See Note: Local 'go' functions and capturing]
 find :: Key -> IntMap a -> a
-find !k = go
+find k m = case lookup k m of
+  Nothing -> not_found
+  Just x -> x
   where
-    go (Bin p l r) | left k p  = go l
-                   | otherwise = go r
-    go (Tip kx x) | k == kx   = x
-                  | otherwise = not_found
-    go Nil = not_found
-
     not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map")
 
 -- | \(O(\min(n,W))\). The expression @('findWithDefault' def k map)@

containers/src/Data/Map/Internal.hs

@@ -401,6 +401,8 @@ import Utils.Containers.Internal.Strict
   (StrictPair(..), StrictTriple(..), toPair)
 import Utils.Containers.Internal.BitQueue
 import Utils.Containers.Internal.EqOrdUtil (EqM(..), OrdM(..))
+import Utils.Containers.Internal.UnboxedMaybe (UMaybe)
+import qualified Utils.Containers.Internal.UnboxedMaybe as UMaybe
 #ifdef DEFINE_ALTERF_FALLBACK
 import Utils.Containers.Internal.BitUtil (wordSize)
 #endif
@@ -572,14 +574,17 @@ size (Bin sz _ _ _ _) = sz
 -- >   John's currency: Just "Euro"
 -- >   Pete's currency: Nothing
 lookup :: Ord k => k -> Map k a -> Maybe a
-lookup = go
-  where
-    go !_ Tip = Nothing
-    go k (Bin _ kx x l r) = case compare k kx of
-      LT -> go k l
-      GT -> go k r
-      EQ -> Just x
-{-# INLINABLE lookup #-}
+lookup k m = UMaybe.toMaybe (lookupGo k m)
+-- Inline so that GHC can eliminate the Maybe if it matched on at the call-site.
+{-# INLINE lookup #-}
+
+lookupGo :: Ord k => k -> Map k a -> UMaybe a
+lookupGo !_ Tip = UMaybe.nothing ()
+lookupGo !k (Bin _ kx x l r) = case compare k kx of
+  LT -> lookupGo k l
+  GT -> lookupGo k r
+  EQ -> UMaybe.just x
+{-# INLINABLE lookupGo #-}
 
 -- | \(O(\log n)\). Is the key a member of the map? See also 'notMember'.
 --
@@ -605,13 +610,9 @@ notMember k m = not $ member k m
 {-# INLINABLE notMember #-}
 
 find :: Ord k => k -> Map k a -> a
-find = go
-  where
-    go !_ Tip = error "Map.!: given key is not an element in the map"
-    go k (Bin _ kx x l r) = case compare k kx of
-      LT -> go k l
-      GT -> go k r
-      EQ -> x
+find k m = case lookup k m of
+  Nothing -> error "Map.!: given key is not an element in the map"
+  Just x -> x
 {-# INLINABLE find #-}
 
 -- | \(O(\log n)\). The expression @('findWithDefault' def k map)@ returns

containers/src/Utils/Containers/Internal/UnboxedMaybe.hs

@@ -0,0 +1,43 @@
+{-# LANGUAGE CPP #-}
+#if __GLASGOW_HASKELL__ >= 810
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE UnliftedNewtypes #-}
+#endif
+
+-- | Unboxed Maybe type, for internal use.
+module Utils.Containers.Internal.UnboxedMaybe
+  ( UMaybe
+  , nothing
+  , just
+  , toMaybe
+  ) where
+
+#if __GLASGOW_HASKELL__ >= 810
+newtype UMaybe a = UMaybe (# (# #) | a #)
+
+-- Takes an argument because
+-- Top-level bindings for unlifted types aren't allowed.
+nothing :: b -> UMaybe a
+nothing _ = UMaybe (# (# #) | #)
+
+just :: a -> UMaybe a
+just x = UMaybe (# | x #)
+
+toMaybe :: UMaybe a -> Maybe a
+toMaybe (UMaybe m) = case m of
+  (# _ | #) -> Nothing
+  (# | x #) -> Just x
+{-# INLINE toMaybe #-}
+#else
+type UMaybe = Maybe
+
+nothing :: b -> UMaybe a
+nothing _ = Nothing
+
+just :: a -> UMaybe a
+just = Just
+
+toMaybe :: UMaybe a -> Maybe a
+toMaybe m = m
+#endif