97 lines
3.0 KiB
Haskell
97 lines
3.0 KiB
Haskell
|
{-# LANGUAGE MagicHash #-}
|
||
|
|
{-# LANGUAGE LambdaCase #-}
|
||
|
|
{-# LANGUAGE UnboxedTuples #-}
|
||
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
||
|
|
{-# LANGUAGE RecursiveDo #-}
|
||
|
|
module UIO where
|
||
|
|
|
||
|
|
import System.Mem
|
||
|
|
import System.Mem.Weak
|
||
|
|
import Data.IORef
|
||
|
|
import GHC.IORef
|
||
|
|
import GHC.STRef
|
||
|
|
import GHC.IO
|
||
|
|
import GHC.Weak
|
||
|
|
import GHC.Prim
|
||
|
|
import Control.Monad.Primitive
|
||
|
|
import Data.IntMap.Strict (IntMap)
|
||
|
|
import qualified Data.IntMap.Strict as IntMap
|
||
|
|
import Data.Set (Set)
|
||
|
|
import qualified Data.Set as Set
|
||
|
|
import System.IO.Unsafe
|
||
|
|
import Control.Monad
|
||
|
|
import Control.Concurrent
|
||
|
|
import Data.Foldable
|
||
|
|
import Data.These
|
||
|
|
import Unsafe.Coerce
|
||
|
|
import Control.Monad.Fix
|
||
|
|
import Data.List.NonEmpty (NonEmpty (..))
|
||
|
|
import Data.Semigroup
|
||
|
|
import GHC.Magic
|
||
|
|
|
||
|
|
data Done = Done
|
||
|
|
|
||
|
|
instance Semigroup Done where
|
||
|
|
{-# INLINE (<>) #-}
|
||
|
|
(<>) = seq
|
||
|
|
sconcat = mconcat . toList
|
||
|
|
stimes _ d = d
|
||
|
|
|
||
|
|
instance Monoid Done where
|
||
|
|
{-# INLINE mempty #-}
|
||
|
|
mempty = Done
|
||
|
|
mconcat = \case
|
||
|
|
[] -> Done
|
||
|
|
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 UIO2 a = UIO2 { unUIO2 :: State# RealWorld -> (Done, a) }
|
||
|
|
|
||
|
|
instance Functor UIO2 where
|
||
|
|
{-# INLINE fmap #-}
|
||
|
|
fmap f x = x >>= (pure . f)
|
||
|
|
|
||
|
|
instance Applicative UIO2 where
|
||
|
|
{-# INLINE pure #-}
|
||
|
|
{-# INLINE (*>) #-}
|
||
|
|
{-# INLINE (<*>) #-}
|
||
|
|
pure x = UIO2 (\s -> (Done, x))
|
||
|
|
UIO2 m *> UIO2 k = UIO2 (\s ->
|
||
|
|
let (ms, _) = m (uniqueState 1# s)
|
||
|
|
(ks, b) = k (uniqueState 2# s)
|
||
|
|
in (ms `seq` ks, b))
|
||
|
|
UIO2 m <*> UIO2 k = UIO2 (\s ->
|
||
|
|
let (ms, f) = m (uniqueState 1# s)
|
||
|
|
(ks, x) = k (uniqueState 2# s)
|
||
|
|
in (ms `seq` ks, f x))
|
||
|
|
|
||
|
|
instance Monad UIO2 where
|
||
|
|
{-# INLINE (>>) #-}
|
||
|
|
{-# INLINE (>>=) #-}
|
||
|
|
(>>) = (*>)
|
||
|
|
UIO2 m >>= k = UIO2 (\s ->
|
||
|
|
let (ms, a) = m (uniqueState 1# s)
|
||
|
|
(ks, b) = unUIO2 (k a) (uniqueState 2# s)
|
||
|
|
in (ms `seq` ks, b))
|
||
|
|
|
||
|
|
instance MonadFix UIO2 where
|
||
|
|
mfix k = UIO2 (\s ->
|
||
|
|
let (ks, result) = unUIO2 (k result) s
|
||
|
|
in (ks, result))
|
||
|
|
|
||
|
|
runUIO2 :: UIO2 a -> IO a
|
||
|
|
runUIO2 (UIO2 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
|
||
|
|
(!Done, 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 -> UIO2 a
|
||
|
|
unordered (IO m) = UIO2 (\s -> case m s of (# _, x #) -> (Done, x))
|
||
|
|
|
||
|
|
-- Force GHC to treat each of these state tokens as unique. This way, multiple identical calls, e.g. to newIORef are not treated as identical, because they have different state tokens. Ideally, we would inline this after common sub-expression elimination finishes, so that it is costless.
|
||
|
|
{-# INLINE uniqueState #-}
|
||
|
|
uniqueState :: Int# -> State# RealWorld -> State# RealWorld
|
||
|
|
uniqueState = noinline (\_ s -> s)
|