UIO works, and now includes a plugin to eliminate the calls to uniqueState

src/Test/UIO.hs

@@ -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 #-}
+{-# NOINLINE testUIOReplicate #-}
-testUIOUnique :: IO ()
+testUIOReplicate :: IO ()
-testUIOUnique = do
+testUIOReplicate = do
-  r <- runUIO2 $ do
+  rs <- runUIO2 $ do
-    -- The following two lines can be merged by common subexpression elimination (CSE), which is very bad
+    rs <- replicateM 10 $ timeless $ newIORef 2
-    r <- unordered $ newIORef 2
+    forM_ rs $ \r ->
-    r2 <- unordered $ newIORef 2
+      unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
-    -- 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!
+    pure rs
-    unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
+  vs <- mapM readIORef rs
-    unordered $ atomicModifyIORef' r2 $ \v -> (v + 10, ())
+  print $ sum vs
-    pure r
-  print =<< readIORef r
 
 {-# NOINLINE testBoth #-}
 testIO :: IO ()

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)

src/UIO/Plugin.hs

@@ -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')