92 lines
2.8 KiB
Haskell
92 lines
2.8 KiB
Haskell
{-# LANGUAGE MagicHash #-}
|
|
{-# LANGUAGE LambdaCase #-}
|
|
{-# LANGUAGE UnboxedTuples #-}
|
|
module UIO
|
|
( module UIO
|
|
, module X
|
|
) where
|
|
|
|
import Control.Monad.Fix
|
|
import Data.Foldable
|
|
import Data.Semigroup
|
|
import GHC.IO
|
|
import GHC.Prim
|
|
import GHC.Magic
|
|
|
|
import UIO.Plugin as X
|
|
|
|
data RemainingWork = RemainingWork
|
|
|
|
instance Semigroup RemainingWork where
|
|
{-# INLINE (<>) #-}
|
|
(<>) = seq
|
|
sconcat = mconcat . toList
|
|
stimes _ d = d
|
|
|
|
instance Monoid RemainingWork where
|
|
{-# INLINE mempty #-}
|
|
mempty = RemainingWork
|
|
mconcat = \case
|
|
[] -> RemainingWork
|
|
x : xs -> x `seq` mconcat xs
|
|
|
|
-- Unordered IO - we want to allocate things, strictly evaluate things, etc., but we don't actually care what order it is done in
|
|
newtype UIO a = UIO { unUIO :: State# RealWorld -> (RemainingWork, a) }
|
|
|
|
instance Functor UIO where
|
|
{-# INLINE fmap #-}
|
|
fmap f x = x >>= (pure . f)
|
|
|
|
instance Applicative UIO where
|
|
{-# INLINE pure #-}
|
|
{-# INLINE (*>) #-}
|
|
{-# INLINE (<*>) #-}
|
|
pure x = UIO (\_ -> (RemainingWork, x))
|
|
UIO m *> UIO k = UIO (\s ->
|
|
let (ms, _) = m (uniqueState# 1# s)
|
|
(ks, b) = k (uniqueState# 2# s)
|
|
in (ms <> ks, b))
|
|
UIO m <*> UIO k = UIO (\s ->
|
|
let (ms, f) = m (uniqueState# 1# s)
|
|
(ks, x) = k (uniqueState# 2# s)
|
|
in (ms <> ks, f x))
|
|
|
|
instance Monad UIO where
|
|
{-# INLINE (>>) #-}
|
|
{-# INLINE (>>=) #-}
|
|
(>>) = (*>)
|
|
UIO m >>= k = UIO (\s ->
|
|
let (ms, a) = m (uniqueState# 1# s)
|
|
(ks, b) = unUIO (k a) (uniqueState# 2# s)
|
|
in (ms <> ks, b))
|
|
|
|
instance MonadFix UIO where
|
|
{-# INLINE mfix #-}
|
|
mfix k = UIO (\s ->
|
|
let (ks, result) = unUIO (k result) s
|
|
in (ks, result))
|
|
|
|
runUIO :: UIO a -> IO a
|
|
runUIO (UIO m) = do
|
|
-- We use a bang pattern here instead of "evaluate", because "evaluate" leaves a "seq#" clutting up our core, but the bang pattern does not
|
|
(!RemainingWork, result) <- IO (\s -> (# s, m s #)) --TODO: This returns the same state we were given; should we call uniqueState# 1# or something on it?
|
|
pure result
|
|
|
|
-- The following is marked NOINLINE because unsafeDupablePerformIO is marked NOINLINE. I don't really understand it.
|
|
{-# INLINE unordered #-}
|
|
unordered :: IO a -> UIO a
|
|
unordered (IO m) = UIO (\s -> let x = case m s of (# _, x #) -> x in (x `seq` RemainingWork, x))
|
|
|
|
-- | Perform an action only when its result is needed. This action will be unique, but the computation will be considered finished regardless of whether this action has run. This is appropriate for functions like `newIORef`.
|
|
{-# INLINE timeless #-}
|
|
timeless :: IO a -> UIO a
|
|
timeless (IO m) = UIO (\s -> (RemainingWork, case m s of (# _, x #) -> x))
|
|
|
|
{-# INLINE listen #-}
|
|
listen :: UIO a -> UIO (RemainingWork, a)
|
|
listen (UIO m) = UIO (\s -> let (done, a) = m s in (done, (done, a)))
|
|
|
|
{-# INLINE after #-}
|
|
after :: RemainingWork -> UIO a -> UIO a
|
|
after w (UIO m) = UIO (\s -> m (w `seq` s))
|