More UIO2 work

main.hs

@@ -23,6 +23,8 @@ import Data.Foldable
 import Data.These
 import Unsafe.Coerce
 import Control.Monad.Fix
+import Data.List.NonEmpty (NonEmpty (..))
+import Data.Semigroup
 
 main :: IO ()
 main = do
@@ -32,8 +34,23 @@ main = do
   testGraphX
   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 -> ((), a) }
+newtype UIO2 a = UIO2 { unUIO2 :: State# RealWorld -> (Done, a) }
 
 instance Functor UIO2 where
   {-# INLINE fmap #-}
@@ -43,7 +60,7 @@ instance Applicative UIO2 where
   {-# INLINE pure #-}
   {-# INLINE (*>) #-}
   {-# INLINE (<*>) #-}
-  pure x = UIO2 (\s -> ((), x))
+  pure x = UIO2 (\s -> (Done, x))
   UIO2 m *> UIO2 k = UIO2 (\s ->
     let (ms, _) = m s
         (ks, b) = k s
@@ -63,31 +80,27 @@ instance Monad UIO2 where
     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))
+  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 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))
+unordered (IO m) = UIO2 (\s -> case m s of (# _, x #) -> (Done, x))
 
 testUIO2 :: IO ()
 testUIO2 = do
   r <- runUIO2 $ mdo
-    unordered $ writeIORef r 5
+    unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
     r <- unordered $ newIORef 2
+    unordered $ atomicModifyIORef' r $ \v -> (v + 10, ())
     pure r
   print =<< readIORef r