-- Unboxed counterparts to data structures

{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE UnliftedNewtypes #-}

{-# OPTIONS_GHC -fno-omit-interface-pragmas #-}
  -- If you use -fomit-interface-pragmas for your build, we won't
  -- inline the matcher for JustUB, and that turns out to have a
  -- catastropic effect on Lint, which uses unboxed Maybes.
  -- Simple fix: switch off -fomit-interface-pragmas for this tiny
  -- and very stable module.

module GHC.Data.Unboxed (
  MaybeUB(JustUB, NothingUB),
  fmapMaybeUB, fromMaybeUB, apMaybeUB, maybeUB,
  traverseMaybeUB
  ) where

import GHC.Prelude

-- | Like Maybe, but using unboxed sums.
--
-- Use with care. Using a unboxed maybe is not always a win
-- in execution *time* even when allocations go down. So make
-- sure to benchmark for execution time as well. If the difference
-- in *runtime* for the compiler is too small to measure it's likely
-- better to use a regular Maybe instead.
--
-- This is since it causes more function arguments to be passed, and
-- potentially more variables to be captured by closures increasing
-- closure size.
newtype MaybeUB a = MaybeUB (# (# #) | a #)

pattern JustUB :: a -> MaybeUB a
pattern $mJustUB :: forall {r} {a}. MaybeUB a -> (a -> r) -> ((# #) -> r) -> r
$bJustUB :: forall a. a -> MaybeUB a
JustUB x = MaybeUB (# | x #)

pattern NothingUB :: MaybeUB a
pattern $mNothingUB :: forall {r} {a}. MaybeUB a -> ((# #) -> r) -> ((# #) -> r) -> r
$bNothingUB :: (# #) -> forall a. MaybeUB a
NothingUB = MaybeUB (# (# #) | #)

{-# COMPLETE NothingUB, JustUB #-}

fromMaybeUB :: a -> MaybeUB a -> a
fromMaybeUB :: forall a. a -> MaybeUB a -> a
fromMaybeUB a
d MaybeUB a
NothingUB = a
d
fromMaybeUB a
_ (JustUB a
x) = a
x

apMaybeUB :: MaybeUB (a -> b) -> MaybeUB a -> MaybeUB b
apMaybeUB :: forall a b. MaybeUB (a -> b) -> MaybeUB a -> MaybeUB b
apMaybeUB (JustUB a -> b
f) (JustUB a
x) = b -> MaybeUB b
forall a. a -> MaybeUB a
JustUB (a -> b
f a
x)
apMaybeUB MaybeUB (a -> b)
_ MaybeUB a
_ = (# #) -> forall a. MaybeUB a
forall a. MaybeUB a
NothingUB

fmapMaybeUB :: (a -> b) -> MaybeUB a -> MaybeUB b
fmapMaybeUB :: forall a b. (a -> b) -> MaybeUB a -> MaybeUB b
fmapMaybeUB a -> b
_f MaybeUB a
NothingUB = (# #) -> forall a. MaybeUB a
forall a. MaybeUB a
NothingUB
fmapMaybeUB a -> b
f (JustUB a
x) = b -> MaybeUB b
forall a. a -> MaybeUB a
JustUB (b -> MaybeUB b) -> b -> MaybeUB b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
x

maybeUB :: b -> (a -> b) -> MaybeUB a -> b
maybeUB :: forall b a. b -> (a -> b) -> MaybeUB a -> b
maybeUB b
_def a -> b
f (JustUB a
x) = a -> b
f a
x
maybeUB b
def a -> b
_f MaybeUB a
NothingUB = b
def

toMaybe :: MaybeUB a -> Maybe a
toMaybe :: forall a. MaybeUB a -> Maybe a
toMaybe MaybeUB a
NothingUB = Maybe a
forall a. Maybe a
Nothing
toMaybe (JustUB a
a) = a -> Maybe a
forall a. a -> Maybe a
Just a
a

-- | Like 'traverse' for the 'Maybe' applicative, but avoiding intermediate
-- allocations.
--
-- The passed-in function must inline for this to help at all.
traverseMaybeUB
  :: forall a b
  .  (a -> Maybe b) -- ^ @INLINE@ function to map over the list
  -> [a] -> Maybe [b]
{-# INLINE traverseMaybeUB #-}
  -- 'traverseMaybeUB' must inline so that the function it is passed can inline
traverseMaybeUB :: forall a b. (a -> Maybe b) -> [a] -> Maybe [b]
traverseMaybeUB a -> Maybe b
f = \ [a]
xs -> MaybeUB [b] -> Maybe [b]
forall a. MaybeUB a -> Maybe a
toMaybe (MaybeUB [b] -> Maybe [b]) -> MaybeUB [b] -> Maybe [b]
forall a b. (a -> b) -> a -> b
$ [a] -> MaybeUB [b]
go [a]
xs
  where
    go :: [a] -> MaybeUB [b]
    go :: [a] -> MaybeUB [b]
go [] = [b] -> MaybeUB [b]
forall a. a -> MaybeUB a
JustUB []
    go (a
a:[a]
as)
      | Just b
b <- a -> Maybe b
f a
a -- f is assumed to inline
      , JustUB [b]
bs <- [a] -> MaybeUB [b]
go [a]
as
      = [b] -> MaybeUB [b]
forall a. a -> MaybeUB a
JustUB (b
bb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
bs)
      | Bool
otherwise
      = (# #) -> forall a. MaybeUB a
forall a. MaybeUB a
NothingUB