Switch to using a patched GHC

src/Test/UIO.hs

@@ -1,15 +1,21 @@
 {-# LANGUAGE RecursiveDo #-}
 {-# OPTIONS_GHC -fplugin UIO #-}
-module Test.UIO where
+module Test.UIO (test) where
 
+import GHC.IO
+import GHC.Magic
 import Data.IORef
 import Control.Monad
 
 import UIO
 import Test.UIO.MultiModule
+import Test.UIO.Builder
 
 test :: IO ()
 test = do
+  putStrLn "testUIOBuilder"
+  testUIOBuilder
+{-
   putStrLn "testUIOFix"
   testUIOFix
   putStrLn "testUIOUnique"
@@ -26,8 +32,11 @@ test = do
   testIO
   putStrLn "testMultiModule"
   testMultiModule
---  putStrLn "testUIOPrintLots"
+  putStrLn "testFloatIn"
---  testUIOPrintLots
+  testFloatIn
+  putStrLn "testUIOPrintLots"
+  testUIOPrintLots
+-}
 
 {-# NOINLINE testUIOUnique #-}
 testUIOUnique :: IO ()
@@ -106,6 +115,16 @@ testUIOCycle = runUIO $ mdo
   r <- timeless $ newIORef $ CycleRef r
   pure ()
 
+{-# NOINLINE testFloatIn #-}
+testFloatIn :: IO ()
+testFloatIn = runUIO $ do
+  r <- timeless $ newIORef 0
+  x <- unordered $ atomicModifyIORef r $ \x -> (succ x, succ x)
+  let {-# NOINLINE blah #-}
+      blah :: IO ()
+      blah = IO $ oneShot (\s -> unIO (print x) s)
+  unordered $ replicateM_ 10 blah
+
 {-# NOINLINE testIO #-}
 testIO :: IO ()
 testIO = do

src/UIO.hs

@@ -11,6 +11,7 @@ import Data.Foldable
 import Data.Semigroup
 import GHC.IO
 import GHC.Prim
+import GHC.Magic
 
 import UIO.Plugin as X
 
@@ -42,12 +43,12 @@ instance Applicative UIO where
   {-# INLINE (<*>) #-}
   pure x = UIO (\_ -> (RemainingWork, x))
   UIO m *> UIO k = UIO (\s ->
-    let (ms, _) = m (uniqueState 1# s)
+    let (ms, _) = m (uniqueState# 1# s)
-        (ks, b) = k (uniqueState 2# s)
+        (ks, b) = k (uniqueState# 2# s)
     in (ms <> ks, b))
   UIO m <*> UIO k = UIO (\s ->
-    let (ms, f) = m (uniqueState 1# s)
+    let (ms, f) = m (uniqueState# 1# s)
-        (ks, x) = k (uniqueState 2# s)
+        (ks, x) = k (uniqueState# 2# s)
     in (ms <> ks, f x))
 
 instance Monad UIO where
@@ -55,8 +56,8 @@ instance Monad UIO where
   {-# INLINE (>>=)  #-}
   (>>)      = (*>)
   UIO m >>= k = UIO (\s ->
-    let (ms, a) = m (uniqueState 1# s)
+    let (ms, a) = m (uniqueState# 1# s)
-        (ks, b) = unUIO (k a) (uniqueState 2# s)
+        (ks, b) = unUIO (k a) (uniqueState# 2# s)
     in (ms <> ks, b))
 
 instance MonadFix UIO where
@@ -68,7 +69,7 @@ instance MonadFix UIO where
 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?
+  (!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.
@@ -81,11 +82,10 @@ unordered (IO m) = UIO (\s -> let x = case m s of (# _, x #) -> x in (x `seq` Re
 timeless :: IO a -> UIO a
 timeless (IO m) = UIO (\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 listen #-}
-{-# NOINLINE uniqueState #-}
+listen :: UIO a -> UIO (RemainingWork, a)
-uniqueState :: Int# -> State# RealWorld -> State# RealWorld
+listen (UIO m) = UIO (\s -> let (done, a) = m s in (done, (done, a)))
-uniqueState = uniqueState'
 
-{-# NOINLINE uniqueState' #-}
+{-# INLINE after #-}
-uniqueState' :: Int# -> State# RealWorld -> State# RealWorld
+after :: RemainingWork -> UIO a -> UIO a
-uniqueState' _ s = s
+after w (UIO m) = UIO (\s -> m (w `seq` s))

src/UIO/Plugin.hs

@@ -4,20 +4,21 @@
 module UIO.Plugin (plugin) where
 
 import GHC.Plugins
-import Data.Generics.Schemes (everywhereM)
+--import Data.Generics.Schemes (everywhereM)
-import Data.Generics.Aliases (mkM)
+--import Data.Generics.Aliases (mkM)
 
+plugin :: Plugin
+plugin = defaultPlugin
+    { {- installCoreToDos = install
+    , -} pluginRecompile = purePlugin
+    }
+
+{-
 isSpecificFunction :: Id -> Bool
 isSpecificFunction f = case nameModule_maybe (idName f) of
   Just m -> moduleNameString (moduleName m) == "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
@@ -53,3 +54,4 @@ inlineSpecificFunctionBindM transform (NonRec b e) = do
 inlineSpecificFunctionBindM transform (Rec pairs) = do
     pairs' <- mapM (\(b, e) -> everywhereM (mkM transform) e >>= \e' -> return (b, e')) pairs
     return (Rec pairs')
+-}