From 25cf31c7dd93fd79ca264075a20390e384b038dc Mon Sep 17 00:00:00 2001 From: Ryan Trinkle Date: Mon, 10 Apr 2023 15:37:08 -0400 Subject: [PATCH] UIO works and produces decent core, except for the noinline'd `uniqueState` calls --- .gitignore | 7 ++++ main.hs | 73 ++----------------------------------- src/Test/UIO.hs | 60 +++++++++++++++++++++++++++++++ src/UIO.hs | 96 +++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 166 insertions(+), 70 deletions(-) create mode 100644 .gitignore create mode 100644 src/Test/UIO.hs create mode 100644 src/UIO.hs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..75a909a --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +# Ignore all files without a dot in them (usually binaries) +* +!*.* + +*.dump-* +*.hi +*.o diff --git a/main.hs b/main.hs index fbc6102..d9643ac 100644 --- a/main.hs +++ b/main.hs @@ -26,6 +26,9 @@ import Control.Monad.Fix import Data.List.NonEmpty (NonEmpty (..)) import Data.Semigroup +import UIO +import Test.UIO + main :: IO () main = do testDualWeak @@ -34,76 +37,6 @@ main = do testGraphX testGraphO -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 s - (ks, b) = k s - in (ms `seq` ks, b)) - UIO2 m <*> UIO2 k = UIO2 (\s -> - let (ms, f) = m s - (ks, x) = k s - in (ms `seq` ks, f x)) - -instance Monad UIO2 where - {-# INLINE (>>) #-} - {-# INLINE (>>=) #-} - (>>) = (*>) - UIO2 m >>= k = UIO2 (\s -> - let (ms, a) = m s - (ks, b) = unUIO2 (k a) 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 - (done, result) <- IO (\s -> (# s, m s #)) - evaluate done - pure result - --- The following is marked NOINLINE because unsafeDupablePerformIO is marked NOINLINE. I don't really understand it. -{-# NOINLINE unordered #-} -unordered :: IO a -> UIO2 a -unordered (IO m) = UIO2 (\s -> case m s of (# _, x #) -> (Done, x)) - -testUIO2 :: IO () -testUIO2 = do - r <- runUIO2 $ mdo - unordered $ atomicModifyIORef' r $ \v -> (v + 5, ()) - r <- unordered $ newIORef 2 - unordered $ atomicModifyIORef' r $ \v -> (v + 10, ()) - pure r - print =<< readIORef r - type UIO = IO data WeakBagInput a diff --git a/src/Test/UIO.hs b/src/Test/UIO.hs new file mode 100644 index 0000000..838a916 --- /dev/null +++ b/src/Test/UIO.hs @@ -0,0 +1,60 @@ +{-# LANGUAGE RecursiveDo #-} +module Test.UIO where + +import Data.IORef + +import UIO + +testBoth :: IO () +testBoth = do + testUIOFix + testUIOUnique + testUIOMany + testIO + +{-# NOINLINE testUIOFix #-} +testUIOFix :: IO () +testUIOFix = do + r <- runUIO2 $ mdo + unordered $ atomicModifyIORef' r $ \v -> (v + 5, ()) + r <- unordered $ newIORef 2 + pure r + print =<< readIORef r + +{-# NOINLINE testUIOMany #-} +testUIOMany :: IO () +testUIOMany = do + r <- runUIO2 $ do + r <- unordered $ newIORef 0 + unordered $ atomicModifyIORef' r $ \v -> (v + 1, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 2, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 3, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 4, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 5, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 6, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 7, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 8, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 9, ()) + unordered $ atomicModifyIORef' r $ \v -> (v + 10, ()) + pure r + print =<< readIORef r + +{-# NOINLINE testUIOUnique #-} +testUIOUnique :: IO () +testUIOUnique = do + r <- runUIO2 $ do + -- The following two lines can be merged by common subexpression elimination (CSE), which is very bad + r <- unordered $ newIORef 2 + r2 <- unordered $ newIORef 2 + -- Note that the following two lines must be different, otherwise they will *also* be merged by CSE, which will make the test appear to succeed! + unordered $ atomicModifyIORef' r $ \v -> (v + 5, ()) + unordered $ atomicModifyIORef' r2 $ \v -> (v + 10, ()) + pure r + print =<< readIORef r + +{-# NOINLINE testBoth #-} +testIO :: IO () +testIO = do + r <- newIORef 2 + atomicModifyIORef' r $ \v -> (v + 5, ()) + print =<< readIORef r diff --git a/src/UIO.hs b/src/UIO.hs new file mode 100644 index 0000000..67920ff --- /dev/null +++ b/src/UIO.hs @@ -0,0 +1,96 @@ +{-# 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)