From 6b27e947ee391469a9326af96f1a3383d65555f6 Mon Sep 17 00:00:00 2001
From: Ryan Trinkle <ryan.trinkle@obsidian.systems>
Date: Sun, 9 Jul 2023 12:34:47 -0400
Subject: [PATCH] Switch to using a patched GHC

---
 src/Test/UIO.hs   | 25 ++++++++++++++++++++++---
 src/UIO.hs        | 28 ++++++++++++++--------------
 src/UIO/Plugin.hs | 18 ++++++++++--------
 3 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/src/Test/UIO.hs b/src/Test/UIO.hs
index 7cc52ca..5086f5d 100644
--- a/src/Test/UIO.hs
+++ b/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"
---  testUIOPrintLots
+  putStrLn "testFloatIn"
+  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
diff --git a/src/UIO.hs b/src/UIO.hs
index af0e140..3eaa92d 100644
--- a/src/UIO.hs
+++ b/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)
-        (ks, b) = k (uniqueState 2# s)
+    let (ms, _) = m (uniqueState# 1# s)
+        (ks, b) = k (uniqueState# 2# s)
     in (ms <> ks, b))
   UIO m <*> UIO k = UIO (\s ->
-    let (ms, f) = m (uniqueState 1# s)
-        (ks, x) = k (uniqueState 2# s)
+    let (ms, f) = m (uniqueState# 1# 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)
-        (ks, b) = unUIO (k a) (uniqueState 2# s)
+    let (ms, a) = m (uniqueState# 1# 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.
-{-# NOINLINE uniqueState #-}
-uniqueState :: Int# -> State# RealWorld -> State# RealWorld
-uniqueState = uniqueState'
+{-# INLINE listen #-}
+listen :: UIO a -> UIO (RemainingWork, a)
+listen (UIO m) = UIO (\s -> let (done, a) = m s in (done, (done, a)))
 
-{-# NOINLINE uniqueState' #-}
-uniqueState' :: Int# -> State# RealWorld -> State# RealWorld
-uniqueState' _ s = s
+{-# INLINE after #-}
+after :: RemainingWork -> UIO a -> UIO a
+after w (UIO m) = UIO (\s -> m (w `seq` s))
diff --git a/src/UIO/Plugin.hs b/src/UIO/Plugin.hs
index a340001..6784d4b 100644
--- a/src/UIO/Plugin.hs
+++ b/src/UIO/Plugin.hs
@@ -4,20 +4,21 @@
 module UIO.Plugin (plugin) where
 
 import GHC.Plugins
-import Data.Generics.Schemes (everywhereM)
-import Data.Generics.Aliases (mkM)
+--import Data.Generics.Schemes (everywhereM)
+--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')
+-}