ryantrinkle/weak-refs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

UIO works and produces decent core, except for the noinline'd uniqueState calls

Browse Source
This commit is contained in:
Ryan Trinkle 2023-04-10 15:37:08 -04:00
parent b6c2099b02
commit 25cf31c7dd
4 changed files with 166 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
.gitignore vendored Normal file
View File

@ -0,0 +1,7 @@
# Ignore all files without a dot in them (usually binaries)
*
!*.*
*.dump-*
*.hi
*.o

73
main.hs
View File

@ -26,6 +26,9 @@ import Control.Monad.Fix
import Data.List.NonEmpty (NonEmpty (..)) import Data.List.NonEmpty (NonEmpty (..))
import Data.Semigroup import Data.Semigroup
import UIO
import Test.UIO
main :: IO () main :: IO ()
main = do main = do
testDualWeak testDualWeak
@ -34,76 +37,6 @@ main = do
testGraphX testGraphX
testGraphO 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 type UIO = IO
data WeakBagInput a data WeakBagInput a

60
src/Test/UIO.hs Normal file
View File

@ -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

96
src/UIO.hs Normal file
View File

@ -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)