Switch to using a patched GHC

.gitignore

@@ -0,0 +1,5 @@
+*.dump-*
+*.hi
+*.*_hi
+*.o
+*.*_o

main.hs

@@ -2,6 +2,7 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE RecursiveDo #-}
 import System.Mem
 import System.Mem.Weak
 import Data.IORef
@@ -19,6 +20,14 @@ import System.IO.Unsafe
 import Control.Monad
 import Control.Concurrent
 import Data.Foldable
+import Data.These
+import Unsafe.Coerce
+import Control.Monad.Fix
+import Data.List.NonEmpty (NonEmpty (..))
+import Data.Semigroup
+
+import UIO
+import Test.UIO
 
 main :: IO ()
 main = do
@@ -28,6 +37,208 @@ main = do
   testGraphX
   testGraphO
 
+type UIO = IO
+
+data WeakBagInput a
+data WeakBagOutput a
+data WeakMutVarInput a = WeakMutVarInput (WeakWithLifespan (IORef a))
+data WeakMutVarOutput a = WeakMutVarOutput (IORef a)
+
+data WeakWithLifespan a = WeakWithLifespan (Weak a) Lifespan
+
+newWeakMutVar :: a -> UIO (WeakMutVarInput a, WeakMutVarOutput a)
+newWeakMutVar a = do
+  r <- newIORef a
+  l <- lifespanOfIORef r
+  w <- newWeakWithLifespan l r
+  pure (WeakMutVarInput w, WeakMutVarOutput r)
+
+writeWeakMutVar :: WeakMutVarInput a -> UIO (Effect a)
+writeWeakMutVar target = mapEffect const =<< modifyWeakMutVar target
+
+modifyWeakMutVar :: WeakMutVarInput a -> UIO (Effect (a -> a))
+modifyWeakMutVar (WeakMutVarInput w) = Effect <$> do
+  mapWeakWithLifespan' (\r f -> atomicModifyIORef' r $ \v -> (f v, ())) w
+
+readWeakMutVar :: WeakMutVarOutput a -> Coeffect a
+readWeakMutVar (WeakMutVarOutput r) = Coeffect $ readIORef r
+
+testEffect1 :: IO ()
+testEffect1 = do
+  (i, o) <- newWeakMutVar "A"
+  e <- traceEffect id =<< writeWeakMutVar i
+  performGCUntilFinalizersQuiesce
+  putStrLn "X"
+  runEffect e "B"
+  touch o
+  _ <- runCoeffect (readWeakMutVar o)
+  performGCUntilFinalizersQuiesce
+  putStrLn "Y"
+  runEffect e "C"
+  putStrLn "Z"
+
+testEffect2 :: IO ()
+testEffect2 = do
+  (ia, oa) <- newWeakMutVar (1 :: Int)
+  (ib, ob) <- newWeakMutVar 1
+  modifyA <- traceEffect (const "modifyA") =<< modifyWeakMutVar ia
+  modifyB <- traceEffect (const "modifyB") =<< modifyWeakMutVar ib
+  let readA = readWeakMutVar oa
+  incrementA <- mapEffect (\() -> (+1)) modifyA
+  addToB <- mapEffect (\n -> (+n)) modifyB
+  addAToB <- withCoeffect readA addToB
+  e <- incrementA `andThen` addAToB
+  performGCUntilFinalizersQuiesce
+  runEffect e ((), ())
+  touch oa
+
+runEffect :: Effect a -> a -> IO ()
+runEffect (Effect e) a = deRefWeakWithLifespan e >>= \case
+  Nothing -> pure ()
+  Just f -> f a
+
+traceEffect :: (a -> String) -> Effect a -> UIO (Effect a)
+traceEffect f = weakEffect $ putStrLn . f
+
+-- Run the given IO action if the provided Effect is still runnable.  Note that the IO action may run even if the effect is *not* runnable, depending on garbage collection timing.
+weakEffect :: (a -> IO ()) -> Effect a -> UIO (Effect a)
+weakEffect f (Effect e) = fmap Effect $ forWeakWithLifespan' e $ \fe a -> do
+  f a
+  fe a
+
+runCoeffect :: Coeffect a -> IO a
+runCoeffect (Coeffect c) = c
+
+
+emptyWeak :: Weak a
+emptyWeak = unsafeCoerce emptyWeakUnit
+
+{-# NOINLINE emptyWeakUnit #-}
+emptyWeakUnit :: Weak ()
+emptyWeakUnit = unsafePerformIO $ do
+  w <- mkWeakPtr () Nothing
+  finalize w
+  pure w
+
+newWeakWithLifespan :: Lifespan -> a -> UIO (WeakWithLifespan a)
+newWeakWithLifespan (Lifespan l) a = do
+  deRefWeak l >>= \case
+    Nothing -> pure $ WeakWithLifespan emptyWeak (Lifespan l)
+    Just r -> do
+      w <- mkWeakWithIORefKey r a
+      pure $ WeakWithLifespan w (Lifespan l)
+
+underlyingLifespan :: WeakWithLifespan a -> Lifespan
+underlyingLifespan (WeakWithLifespan _ l) = l
+
+deRefWeakWithLifespan :: WeakWithLifespan a -> UIO (Maybe a)
+deRefWeakWithLifespan (WeakWithLifespan w _) = deRefWeak w
+
+-- Applies the function strictly; usually this is what you want, so that extra data is not retained
+forWeakWithLifespan' :: WeakWithLifespan a -> (a -> b) -> UIO (WeakWithLifespan b)
+forWeakWithLifespan' = flip mapWeakWithLifespan'
+
+-- Applies the function strictly; usually this is what you want, so that extra data is not retained
+mapWeakWithLifespan' :: (a -> b) -> WeakWithLifespan a -> UIO (WeakWithLifespan b)
+mapWeakWithLifespan' f (WeakWithLifespan w l) = do
+  deRefWeak w >>= \case
+    Nothing -> pure $ WeakWithLifespan emptyWeak l
+    Just v -> newWeakWithLifespan l $! f v
+
+bothAlive :: WeakWithLifespan a -> WeakWithLifespan b -> UIO (WeakWithLifespan (a, b))
+bothAlive = undefined
+
+-- This is impossible because it assumes that weaks *do actual work* when you retrieve them.  This isn't what we want.
+alignAlive :: WeakWithLifespan a -> WeakWithLifespan b -> UIO (WeakWithLifespan (These a b))
+alignAlive = undefined
+
+data Effect a = Effect (WeakWithLifespan (a -> UIO ())) -- Take an `a` and do a side effect with it
+
+data Coeffect a = Coeffect (UIO a)  -- Read an `a` without doing any side effects
+
+-- Like `also`, but also guarantees order.  But I'm not sure what the semantics should really be here, since we could want coeffects ordered separately from effects
+andThen :: Effect a -> Effect b -> UIO (Effect (a, b))
+andThen = also
+
+also :: Effect a -> Effect b -> UIO (Effect (a, b))
+also (Effect ea) (Effect eb) = Effect <$> do
+  myLifespan <- unionLifespan (underlyingLifespan ea) (underlyingLifespan eb)
+  newWeakWithLifespan myLifespan $ \(a, b) -> do
+    deRefWeakWithLifespan ea >>= \case
+      Nothing -> pure ()
+      Just fa -> fa a
+    deRefWeakWithLifespan eb >>= \case
+      Nothing -> pure ()
+      Just fb -> fb b
+
+mapEffect :: (b -> a) -> Effect a -> UIO (Effect b)
+mapEffect f (Effect e) = Effect <$> mapWeakWithLifespan' (\fe -> fe . f) e
+
+withCoeffect :: Coeffect a -> Effect a -> UIO (Effect ())
+withCoeffect (Coeffect c) (Effect e) = Effect <$> do
+  let f fe () = do
+        v <- c
+        fe v
+  mapWeakWithLifespan' f e
+
+coAlso :: Coeffect a -> Coeffect b -> Coeffect (a, b)
+coAlso = undefined
+
+--TODO: This should use a WeakBag
+newtype Lifespan = Lifespan (Weak (IORef [LifespanBacklink]))
+newtype LifespanBacklink = LifespanBacklink (IORef [LifespanBacklink])
+
+modifyWeakIORef :: Weak (IORef a) -> (a -> a) -> IO ()
+modifyWeakIORef w f = deRefWeak w >>= \case
+  Nothing -> pure ()
+  Just r -> atomicModifyIORef' r $ \v -> (f v, ())
+
+lifespanOfIORef :: IORef a -> UIO Lifespan
+lifespanOfIORef basis = do
+  mine <- newIORef []
+  w <- mkWeakWithIORefKey basis mine -- This exploits the fact that System.Mem.Weak references keep the value alive even when the weak reference itself dies.
+  pure $ Lifespan w
+
+-- Return a lifespan
+unionLifespan :: Lifespan -> Lifespan -> UIO Lifespan
+unionLifespan (Lifespan a) (Lifespan b) = do
+  r <- newIORef []
+  w <- mkWeakWithIORefKey r r
+  modifyWeakIORef a (LifespanBacklink r :)
+  modifyWeakIORef b (LifespanBacklink r :)
+  pure $ Lifespan w
+
+intersectionLifespan :: Lifespan -> Lifespan -> UIO Lifespan
+intersectionLifespan a b = undefined
+
+-- If we do readWeakMutVar v `bind` writeWeakMutVar v', we should only keep v alive if `v'`'s output side is alive
+bind :: IO a -> (a -> IO b) -> IO b
+bind = undefined
+
+data Event a = Event
+  { _event_items :: WeakBagInput (a -> IO ()) -- Allows adding items to the weak bag.  Does not keep the weak bag alive; if the bag is gone, adding an item does nothing.
+  , _event_currentValue :: WeakMutVarOutput (Maybe a) -- Allows seeing the current state of the event: if Nothing, it either isn't firing or hasn't fired yet this frame; if Just, it has fired this frame.
+  --TODO: Position in topological ordering; only needed when we introduce Merge
+  }
+
+data Trigger a = Trigger
+  { _trigger_items :: WeakBagOutput (a -> IO ()) -- Allows retrieving the contents of the weak bag
+  , _trigger_currentValue :: WeakMutVarInput (Maybe a)
+  }
+
+newEvent :: IO (Trigger a, Event a)
+newEvent = undefined
+
+fireTrigger :: Trigger a -> a -> IO ()
+fireTrigger = undefined
+
+data Hold a = Hold
+  {
+  }
+
+newHold :: a -> Event a -> IO (Hold a)
+newHold = undefined
+
 -- Demonstrate that a Weak's key keeps its value alive, even if the Weak is dead
 testWeakChain :: IO ()
 testWeakChain = do

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

src/Test/UIO.hs

@@ -0,0 +1,134 @@
+{-# LANGUAGE RecursiveDo #-}
+{-# OPTIONS_GHC -fplugin UIO #-}
+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"
+  testUIOUnique
+  putStrLn "testUIOReplicate"
+  testUIOReplicate
+  putStrLn "testUIOBadTimeless"
+  testUIOBadTimeless
+  putStrLn "testUIOMany"
+  testUIOMany
+  putStrLn "testUIOCycle"
+  testUIOCycle
+  putStrLn "testIO"
+  testIO
+  putStrLn "testMultiModule"
+  testMultiModule
+  putStrLn "testFloatIn"
+  testFloatIn
+  putStrLn "testUIOPrintLots"
+  testUIOPrintLots
+-}
+
+{-# NOINLINE testUIOUnique #-}
+testUIOUnique :: IO ()
+testUIOUnique = 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
+    -- 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
+  7 <- readIORef r
+  pure ()
+
+{-# NOINLINE testUIOFix #-}
+testUIOFix :: IO ()
+testUIOFix = do
+  r <- runUIO $ mdo
+    unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
+    r <- timeless $ newIORef 2
+    pure r
+  7 <- readIORef r
+  pure ()
+
+{-# NOINLINE testUIOBadTimeless #-}
+testUIOBadTimeless :: IO ()
+testUIOBadTimeless = do
+  r <- runUIO $ mdo
+    timeless $ atomicModifyIORef' r $ \v -> (v + 5, ())
+    r <- timeless $ newIORef 2
+    pure r
+  2 <- readIORef r
+  pure ()
+
+{-# NOINLINE testUIOMany #-}
+testUIOMany :: IO ()
+testUIOMany = do
+  r <- runUIO $ do
+    r <- timeless $ newIORef 0
+    unordered $ atomicModifyIORef' r $ \v -> (v + 1, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 2, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 3, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 4, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 6, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 7, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 8, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 9, ())
+    unordered $ atomicModifyIORef' r $ \v -> (v + 10, ())
+    pure r
+  55 <- readIORef r
+  pure ()
+
+{-# NOINLINE testUIOReplicate #-}
+testUIOReplicate :: IO ()
+testUIOReplicate = do
+  rs <- runUIO $ do
+    rs <- replicateM 10 $ timeless $ newIORef 2
+    forM_ rs $ \r ->
+      unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
+    pure rs
+  70 <- fmap sum $ mapM readIORef rs
+  pure ()
+
+{-# NOINLINE testUIOPrintLots #-}
+testUIOPrintLots :: IO ()
+testUIOPrintLots = runUIO $ do
+  replicateM_ 1000000 $ unordered $ putStrLn "Task"
+
+newtype CycleRef = CycleRef (IORef CycleRef)
+
+{-# NOINLINE testUIOCycle #-}
+testUIOCycle :: IO ()
+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
+  r <- newIORef 2
+  atomicModifyIORef' r $ \v -> (v + 5, ())
+  7 <- readIORef r
+  pure ()

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

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

src/UIO.hs

@@ -0,0 +1,91 @@
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE UnboxedTuples #-}
+module UIO
+  ( module UIO
+  , module X
+  ) where
+
+import Control.Monad.Fix
+import Data.Foldable
+import Data.Semigroup
+import GHC.IO
+import GHC.Prim
+import GHC.Magic
+
+import UIO.Plugin as X
+
+data RemainingWork = RemainingWork
+
+instance Semigroup RemainingWork where
+  {-# INLINE (<>) #-}
+  (<>) = seq
+  sconcat = mconcat . toList
+  stimes _ d = d
+
+instance Monoid RemainingWork where
+  {-# INLINE mempty #-}
+  mempty = RemainingWork
+  mconcat = \case
+    [] -> 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 UIO a = UIO { unUIO :: State# RealWorld -> (RemainingWork, a) }
+
+instance Functor UIO where
+  {-# INLINE fmap #-}
+  fmap f x = x >>= (pure . f)
+
+instance Applicative UIO where
+  {-# INLINE pure #-}
+  {-# INLINE (*>) #-}
+  {-# INLINE (<*>) #-}
+  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))
+  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 UIO where
+  {-# INLINE (>>)   #-}
+  {-# INLINE (>>=)  #-}
+  (>>)      = (*>)
+  UIO m >>= k = UIO (\s ->
+    let (ms, a) = m (uniqueState# 1# s)
+        (ks, b) = unUIO (k a) (uniqueState# 2# s)
+    in (ms <> ks, b))
+
+instance MonadFix UIO where
+  {-# INLINE mfix #-}
+  mfix k = UIO (\s ->
+    let (ks, result) = unUIO (k result) s
+    in (ks, result))
+
+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 -> 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 -> UIO a
+timeless (IO m) = UIO (\s -> (RemainingWork, case m s of (# _, x #) -> x))
+
+{-# INLINE listen #-}
+listen :: UIO a -> UIO (RemainingWork, a)
+listen (UIO m) = UIO (\s -> let (done, a) = m s in (done, (done, a)))
+
+{-# INLINE after #-}
+after :: RemainingWork -> UIO a -> UIO a
+after w (UIO m) = UIO (\s -> m (w `seq` 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)
+
+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
+
+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)
+          (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 (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')
+-}