From ebdd46e76e491fccf9a1f43135a2b8bf322ff302 Mon Sep 17 00:00:00 2001
From: Ryan Trinkle <ryan@trinkle.org>
Date: Wed, 12 Apr 2023 21:44:43 -0400
Subject: [PATCH] Add a failing test for the plugin

---
 runTest                            |  5 +++++
 src/Test/UIO.hs                    | 23 ++++++++++---------
 src/Test/UIO/MultiModule.hs        | 19 ++++++++++++++++
 src/Test/UIO/MultiModule/Callee.hs | 14 ++++++++++++
 src/UIO.hs                         | 36 +++++++++++++++---------------
 src/UIO/Plugin.hs                  |  8 +++----
 6 files changed, 72 insertions(+), 33 deletions(-)
 create mode 100644 runTest
 create mode 100644 src/Test/UIO/MultiModule.hs
 create mode 100644 src/Test/UIO/MultiModule/Callee.hs

diff --git a/runTest b/runTest
new file mode 100644
index 0000000..cfc006d
--- /dev/null
+++ b/runTest
@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ghc -O3 -isrc -ddump-to-file -ddump-simpl -ddump-prep -ddump-cse -fforce-recomp -package ghc -dynamic-too -Wall -main-is Test.UIO.test -o ./test src/Test/UIO.hs
+./test
diff --git a/src/Test/UIO.hs b/src/Test/UIO.hs
index b9f2d3d..7cc52ca 100644
--- a/src/Test/UIO.hs
+++ b/src/Test/UIO.hs
@@ -6,9 +6,10 @@ import Data.IORef
 import Control.Monad
 
 import UIO
+import Test.UIO.MultiModule
 
-testBoth :: IO ()
-testBoth = do
+test :: IO ()
+test = do
   putStrLn "testUIOFix"
   testUIOFix
   putStrLn "testUIOUnique"
@@ -23,13 +24,15 @@ testBoth = do
   testUIOCycle
   putStrLn "testIO"
   testIO
+  putStrLn "testMultiModule"
+  testMultiModule
 --  putStrLn "testUIOPrintLots"
 --  testUIOPrintLots
 
 {-# NOINLINE testUIOUnique #-}
 testUIOUnique :: IO ()
 testUIOUnique = do
-  r <- runUIO2 $ do
+  r <- runUIO $ 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
@@ -43,7 +46,7 @@ testUIOUnique = do
 {-# NOINLINE testUIOFix #-}
 testUIOFix :: IO ()
 testUIOFix = do
-  r <- runUIO2 $ mdo
+  r <- runUIO $ mdo
     unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
     r <- timeless $ newIORef 2
     pure r
@@ -53,7 +56,7 @@ testUIOFix = do
 {-# NOINLINE testUIOBadTimeless #-}
 testUIOBadTimeless :: IO ()
 testUIOBadTimeless = do
-  r <- runUIO2 $ mdo
+  r <- runUIO $ mdo
     timeless $ atomicModifyIORef' r $ \v -> (v + 5, ())
     r <- timeless $ newIORef 2
     pure r
@@ -63,7 +66,7 @@ testUIOBadTimeless = do
 {-# NOINLINE testUIOMany #-}
 testUIOMany :: IO ()
 testUIOMany = do
-  r <- runUIO2 $ do
+  r <- runUIO $ do
     r <- timeless $ newIORef 0
     unordered $ atomicModifyIORef' r $ \v -> (v + 1, ())
     unordered $ atomicModifyIORef' r $ \v -> (v + 2, ())
@@ -82,7 +85,7 @@ testUIOMany = do
 {-# NOINLINE testUIOReplicate #-}
 testUIOReplicate :: IO ()
 testUIOReplicate = do
-  rs <- runUIO2 $ do
+  rs <- runUIO $ do
     rs <- replicateM 10 $ timeless $ newIORef 2
     forM_ rs $ \r ->
       unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
@@ -92,18 +95,18 @@ testUIOReplicate = do
 
 {-# NOINLINE testUIOPrintLots #-}
 testUIOPrintLots :: IO ()
-testUIOPrintLots = runUIO2 $ do
+testUIOPrintLots = runUIO $ do
   replicateM_ 1000000 $ unordered $ putStrLn "Task"
 
 newtype CycleRef = CycleRef (IORef CycleRef)
 
 {-# NOINLINE testUIOCycle #-}
 testUIOCycle :: IO ()
-testUIOCycle = runUIO2 $ mdo
+testUIOCycle = runUIO $ mdo
   r <- timeless $ newIORef $ CycleRef r
   pure ()
 
-{-# NOINLINE testBoth #-}
+{-# NOINLINE testIO #-}
 testIO :: IO ()
 testIO = do
   r <- newIORef 2
diff --git a/src/Test/UIO/MultiModule.hs b/src/Test/UIO/MultiModule.hs
new file mode 100644
index 0000000..828aabe
--- /dev/null
+++ b/src/Test/UIO/MultiModule.hs
@@ -0,0 +1,19 @@
+{-# OPTIONS_GHC -fplugin UIO #-}
+module Test.UIO.MultiModule where
+
+import Data.IORef
+
+import UIO
+import Test.UIO.MultiModule.Callee
+
+testMultiModule :: IO ()
+testMultiModule = do
+  r <- runUIO caller
+  12 <- readIORef r
+  pure ()
+
+--TODO: I think this should fail because callee should get inlined here, and then CSE should take place, and it should find multiple newMutVar# operations taking the same state token as input, and they should be CSE'd away.  However, that doesn't seem to happen.  The callee might need a different internal structure to be susceptible to CSE.
+caller :: UIO (IORef Int)
+caller = do
+  callee
+  callee
diff --git a/src/Test/UIO/MultiModule/Callee.hs b/src/Test/UIO/MultiModule/Callee.hs
new file mode 100644
index 0000000..e9fc018
--- /dev/null
+++ b/src/Test/UIO/MultiModule/Callee.hs
@@ -0,0 +1,14 @@
+{-# OPTIONS_GHC -fplugin UIO #-}
+module Test.UIO.MultiModule.Callee where
+
+import Data.IORef
+
+import UIO
+
+--NOTE: It's important that this NOT be marked INLINE; if it is, then the unfolding will be the *unoptimized* code, which will still contain `uniqueState` invocations.  These will be eliminated by the UIO plugin in the caller, which will work correctly, despite the plugin's behavior being unreliable.
+callee :: UIO (IORef Int)
+callee = do
+  r <- timeless $ newIORef 2
+  unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
+  unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
+  pure r
diff --git a/src/UIO.hs b/src/UIO.hs
index d74faca..af0e140 100644
--- a/src/UIO.hs
+++ b/src/UIO.hs
@@ -30,56 +30,56 @@ instance Monoid RemainingWork where
     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 -> (RemainingWork, a) }
+newtype UIO a = UIO { unUIO :: State# RealWorld -> (RemainingWork, a) }
 
-instance Functor UIO2 where
+instance Functor UIO where
   {-# INLINE fmap #-}
   fmap f x = x >>= (pure . f)
 
-instance Applicative UIO2 where
+instance Applicative UIO where
   {-# INLINE pure #-}
   {-# INLINE (*>) #-}
   {-# INLINE (<*>) #-}
-  pure x = UIO2 (\_ -> (RemainingWork, x))
-  UIO2 m *> UIO2 k = UIO2 (\s ->
+  pure x = UIO (\_ -> (RemainingWork, x))
+  UIO m *> UIO k = UIO (\s ->
     let (ms, _) = m (uniqueState 1# s)
         (ks, b) = k (uniqueState 2# s)
     in (ms <> ks, b))
-  UIO2 m <*> UIO2 k = UIO2 (\s ->
+  UIO m <*> UIO k = UIO (\s ->
     let (ms, f) = m (uniqueState 1# s)
         (ks, x) = k (uniqueState 2# s)
     in (ms <> ks, f x))
 
-instance Monad UIO2 where
+instance Monad UIO where
   {-# INLINE (>>)   #-}
   {-# INLINE (>>=)  #-}
   (>>)      = (*>)
-  UIO2 m >>= k = UIO2 (\s ->
+  UIO m >>= k = UIO (\s ->
     let (ms, a) = m (uniqueState 1# s)
-        (ks, b) = unUIO2 (k a) (uniqueState 2# s)
+        (ks, b) = unUIO (k a) (uniqueState 2# s)
     in (ms <> ks, b))
 
-instance MonadFix UIO2 where
+instance MonadFix UIO where
   {-# INLINE mfix #-}
-  mfix k = UIO2 (\s ->
-    let (ks, result) = unUIO2 (k result) s
+  mfix k = UIO (\s ->
+    let (ks, result) = unUIO (k result) s
     in (ks, result))
 
-runUIO2 :: UIO2 a -> IO a
-runUIO2 (UIO2 m) = do
+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?
   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 -> let x = case m s of (# _, x #) -> x in (x `seq` RemainingWork, x))
+unordered :: IO a -> UIO a
+unordered (IO m) = UIO (\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))
+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 #-}
diff --git a/src/UIO/Plugin.hs b/src/UIO/Plugin.hs
index 3e22048..a340001 100644
--- a/src/UIO/Plugin.hs
+++ b/src/UIO/Plugin.hs
@@ -6,12 +6,10 @@ 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"
+  Just m -> moduleNameString (moduleName m) == "UIO" && getOccString f == "uniqueState"
   Nothing  -> False
 
 plugin :: Plugin
@@ -32,7 +30,7 @@ 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
+          (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
@@ -44,7 +42,7 @@ inlineSpecificFunction guts = do
     return guts { mg_binds = inlinedBinds }
 
 inlineSpecificFunctionTransform :: CoreExpr -> CoreM CoreExpr
-inlineSpecificFunctionTransform e@(App (App (Var f) _) s)
+inlineSpecificFunctionTransform (App (App (Var f) _) s)
   | isSpecificFunction f = return s
 inlineSpecificFunctionTransform e = return e