From 5da5ab9045320deac4cf0e2907350366bd9404bc Mon Sep 17 00:00:00 2001
From: Ryan Trinkle <ryan@trinkle.org>
Date: Mon, 10 Apr 2023 11:15:20 -0400
Subject: [PATCH] Initial implementation of UIO2

---
 main.hs | 65 +++++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 63 insertions(+), 2 deletions(-)

diff --git a/main.hs b/main.hs
index 119aed9..c754f27 100644
--- a/main.hs
+++ b/main.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE RecursiveDo #-}
 import System.Mem
 import System.Mem.Weak
 import Data.IORef
@@ -21,6 +22,7 @@ import Control.Concurrent
 import Data.Foldable
 import Data.These
 import Unsafe.Coerce
+import Control.Monad.Fix
 
 main :: IO ()
 main = do
@@ -30,6 +32,67 @@ main = do
   testGraphX
   testGraphO
 
+-- 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 -> ((), 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 -> ((), 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 = do
+    m <- unordered newEmptyMVar
+    UIO2 (\s ->
+            let (rs, ans) = unUIO2 (unordered $ readMVar m) s
+                (ks, result) = unUIO2 (k ans) s
+                (ps, _) = unUIO2 (unordered $ putMVar m result) s
+            in (ps `seq` ks `seq` rs, result))
+
+runUIO2 :: UIO2 a -> IO a
+runUIO2 (UIO2 m) = do
+  (done, result) <- IO (\s -> (# s, m s #))
+  evaluate result
+  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 #) -> ((), x))
+
+testUIO2 :: IO ()
+testUIO2 = do
+  r <- runUIO2 $ mdo
+    unordered $ writeIORef r 5
+    r <- unordered $ newIORef 2
+    pure r
+  print =<< readIORef r
+
+type UIO = IO
+
 data WeakBagInput a
 data WeakBagOutput a
 data WeakMutVarInput a = WeakMutVarInput (WeakWithLifespan (IORef a))
@@ -100,8 +163,6 @@ weakEffect f (Effect e) = fmap Effect $ forWeakWithLifespan' e $ \fe a -> do
 runCoeffect :: Coeffect a -> IO a
 runCoeffect (Coeffect c) = c
 
--- Unordered IO - we want to allocate things, strictly evaluate things, etc., but we don't actually care what order it is done in
-type UIO = IO
 
 emptyWeak :: Weak a
 emptyWeak = unsafeCoerce emptyWeakUnit