UIO works, and now includes a plugin to eliminate the calls to uniqueState
This commit is contained in:
Ryan Trinkle
parent
25cf31c7dd
commit
f2ff3dbb09
@ -1,23 +1,55 @@
|
||||
{-# LANGUAGE RecursiveDo #-}
|
||||
{-# OPTIONS_GHC -fplugin UIO.Plugin #-}
|
||||
module Test.UIO where
|
||||
|
||||
import Data.IORef
|
||||
import Control.Monad
|
||||
|
||||
import UIO
|
||||
|
||||
testBoth :: IO ()
|
||||
testBoth = do
|
||||
putStrLn "testUIOFix"
|
||||
testUIOFix
|
||||
putStrLn "testUIOUnique"
|
||||
testUIOUnique
|
||||
putStrLn "testUIOReplicate"
|
||||
testUIOReplicate
|
||||
putStrLn "testUIOBadTimeless"
|
||||
testUIOBadTimeless
|
||||
putStrLn "testUIOMany"
|
||||
testUIOMany
|
||||
putStrLn "testIO"
|
||||
testIO
|
||||
|
||||
{-# 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 <- timeless $ newIORef 2
|
||||
r2 <- timeless $ 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 testUIOFix #-}
|
||||
testUIOFix :: IO ()
|
||||
testUIOFix = do
|
||||
r <- runUIO2 $ mdo
|
||||
unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
|
||||
r <- unordered $ newIORef 2
|
||||
r <- timeless $ newIORef 2
|
||||
pure r
|
||||
print =<< readIORef r
|
||||
|
||||
{-# NOINLINE testUIOBadTimeless #-}
|
||||
testUIOBadTimeless :: IO ()
|
||||
testUIOBadTimeless = do
|
||||
r <- runUIO2 $ mdo
|
||||
timeless $ atomicModifyIORef' r $ \v -> (v + 5, ())
|
||||
r <- timeless $ newIORef 2
|
||||
pure r
|
||||
print =<< readIORef r
|
||||
|
||||
@ -25,7 +57,7 @@ testUIOFix = do
|
||||
testUIOMany :: IO ()
|
||||
testUIOMany = do
|
||||
r <- runUIO2 $ do
|
||||
r <- unordered $ newIORef 0
|
||||
r <- timeless $ newIORef 0
|
||||
unordered $ atomicModifyIORef' r $ \v -> (v + 1, ())
|
||||
unordered $ atomicModifyIORef' r $ \v -> (v + 2, ())
|
||||
unordered $ atomicModifyIORef' r $ \v -> (v + 3, ())
|
||||
@ -39,18 +71,16 @@ testUIOMany = do
|
||||
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 testUIOReplicate #-}
|
||||
testUIOReplicate :: IO ()
|
||||
testUIOReplicate = do
|
||||
rs <- runUIO2 $ do
|
||||
rs <- replicateM 10 $ timeless $ newIORef 2
|
||||
forM_ rs $ \r ->
|
||||
unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
|
||||
pure rs
|
||||
vs <- mapM readIORef rs
|
||||
print $ sum vs
|
||||
|
||||
{-# NOINLINE testBoth #-}
|
||||
testIO :: IO ()
|
||||
|
||||
39
src/UIO.hs
39
src/UIO.hs
@ -29,23 +29,23 @@ import Data.List.NonEmpty (NonEmpty (..))
|
||||
import Data.Semigroup
|
||||
import GHC.Magic
|
||||
|
||||
data Done = Done
|
||||
data RemainingWork = RemainingWork
|
||||
|
||||
instance Semigroup Done where
|
||||
instance Semigroup RemainingWork where
|
||||
{-# INLINE (<>) #-}
|
||||
(<>) = seq
|
||||
sconcat = mconcat . toList
|
||||
stimes _ d = d
|
||||
|
||||
instance Monoid Done where
|
||||
instance Monoid RemainingWork where
|
||||
{-# INLINE mempty #-}
|
||||
mempty = Done
|
||||
mempty = RemainingWork
|
||||
mconcat = \case
|
||||
[] -> Done
|
||||
[] -> 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 UIO2 a = UIO2 { unUIO2 :: State# RealWorld -> (Done, a) }
|
||||
newtype UIO2 a = UIO2 { unUIO2 :: State# RealWorld -> (RemainingWork, a) }
|
||||
|
||||
instance Functor UIO2 where
|
||||
{-# INLINE fmap #-}
|
||||
@ -55,15 +55,15 @@ instance Applicative UIO2 where
|
||||
{-# INLINE pure #-}
|
||||
{-# INLINE (*>) #-}
|
||||
{-# INLINE (<*>) #-}
|
||||
pure x = UIO2 (\s -> (Done, x))
|
||||
pure x = UIO2 (\s -> (RemainingWork, x))
|
||||
UIO2 m *> UIO2 k = UIO2 (\s ->
|
||||
let (ms, _) = m (uniqueState 1# s)
|
||||
(ks, b) = k (uniqueState 2# s)
|
||||
in (ms `seq` ks, b))
|
||||
in (ms <> 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))
|
||||
in (ms <> ks, f x))
|
||||
|
||||
instance Monad UIO2 where
|
||||
{-# INLINE (>>) #-}
|
||||
@ -72,7 +72,7 @@ instance Monad UIO2 where
|
||||
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))
|
||||
in (ms <> ks, b))
|
||||
|
||||
instance MonadFix UIO2 where
|
||||
mfix k = UIO2 (\s ->
|
||||
@ -82,15 +82,26 @@ instance MonadFix UIO2 where
|
||||
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?
|
||||
(!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 -> UIO2 a
|
||||
unordered (IO m) = UIO2 (\s -> case m s of (# _, x #) -> (Done, x))
|
||||
unordered (IO m) = UIO2 (\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 -> UIO2 a
|
||||
timeless (IO m) = UIO2 (\s -> (RemainingWork, case m s of (# _, x #) -> 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 #-}
|
||||
{-# NOINLINE uniqueState #-}
|
||||
uniqueState :: Int# -> State# RealWorld -> State# RealWorld
|
||||
uniqueState = noinline (\_ s -> s)
|
||||
uniqueState = uniqueState'
|
||||
|
||||
{-# NOINLINE uniqueState' #-}
|
||||
uniqueState' :: Int# -> State# RealWorld -> State# RealWorld
|
||||
uniqueState' _ s = s
|
||||
-- This implementation seems to work sometimes, but I don't understand why, and it seems highly dependent on other aspects of the implementation.
|
||||
-- uniqueState = runRW# (\s _ _ -> s)
|
||||
|
||||
57
src/UIO/Plugin.hs
Normal file
57
src/UIO/Plugin.hs
Normal file
@ -0,0 +1,57 @@
|
||||
{-# LANGUAGE BangPatterns #-}
|
||||
{-# LANGUAGE LambdaCase #-}
|
||||
|
||||
module UIO.Plugin (plugin) where
|
||||
|
||||
import GHC.Plugins
|
||||
import Data.Generics.Schemes (everywhereM)
|
||||
import Data.Generics.Aliases (mkM)
|
||||
import GHC.Types.Name
|
||||
import Debug.Trace
|
||||
|
||||
isSpecificFunction :: Id -> Bool
|
||||
isSpecificFunction f = case nameModule_maybe (idName f) of
|
||||
Just mod -> moduleNameString (moduleName mod) == "UIO" && getOccString f == "uniqueState"
|
||||
Nothing -> False
|
||||
|
||||
plugin :: Plugin
|
||||
plugin = defaultPlugin
|
||||
{ installCoreToDos = install
|
||||
, pluginRecompile = purePlugin
|
||||
}
|
||||
|
||||
install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
|
||||
install _ todos = do
|
||||
let myPass = CoreDoPluginPass ("Inline UIO.uniqueState") inlineSpecificFunction
|
||||
return $ case insertAfterLastCSE myPass todos of
|
||||
Nothing -> myPass : todos
|
||||
Just newTodos -> newTodos
|
||||
|
||||
insertAfterLastCSE :: CoreToDo -> [CoreToDo] -> Maybe [CoreToDo]
|
||||
insertAfterLastCSE myPass = fmap reverse . go . reverse
|
||||
where go = \case
|
||||
[] -> Nothing
|
||||
passes@(CoreCSE {} : _) -> Just (myPass : passes)
|
||||
passes@(h@(CoreDoPasses subPasses) : t) -> case insertAfterLastCSE myPass subPasses of
|
||||
Nothing -> fmap (h :) $ insertAfterLastCSE myPass t
|
||||
Just newSubPasses -> Just $ CoreDoPasses newSubPasses : t
|
||||
h : t -> fmap (h :) $ insertAfterLastCSE myPass t
|
||||
|
||||
inlineSpecificFunction :: ModGuts -> CoreM ModGuts
|
||||
inlineSpecificFunction guts = do
|
||||
let !binds = mg_binds guts
|
||||
inlinedBinds <- mapM (inlineSpecificFunctionBindM inlineSpecificFunctionTransform) binds
|
||||
return guts { mg_binds = inlinedBinds }
|
||||
|
||||
inlineSpecificFunctionTransform :: CoreExpr -> CoreM CoreExpr
|
||||
inlineSpecificFunctionTransform e@(App (App (Var f) _) s)
|
||||
| isSpecificFunction f = return s
|
||||
inlineSpecificFunctionTransform e = return e
|
||||
|
||||
inlineSpecificFunctionBindM :: (CoreExpr -> CoreM CoreExpr) -> CoreBind -> CoreM CoreBind
|
||||
inlineSpecificFunctionBindM transform (NonRec b e) = do
|
||||
e' <- everywhereM (mkM transform) e
|
||||
return (NonRec b e')
|
||||
inlineSpecificFunctionBindM transform (Rec pairs) = do
|
||||
pairs' <- mapM (\(b, e) -> everywhereM (mkM transform) e >>= \e' -> return (b, e')) pairs
|
||||
return (Rec pairs')
|
||||
Loading…
Reference in New Issue
Block a user