oneShot rule for `await`

oneshot.hs

{-#LANGUAGE BangPatterns, MagicHash, UnboxedTuples #-}
import Pipes
import qualified Pipes.Prelude as P
import Control.Monad.Trans.State.Strict
import GHC.Magic
import qualified Pipes.Internal as I
import Data.IORef
import Control.Monad
import GHC.Types
import GHC.Prim

liftio :: IO r -> Proxy a b c d IO r
liftio = liftIO
{-#INLINE [0] liftio #-}

iobind :: IO t -> (t -> Proxy a' a b' b IO r) -> Proxy a' a b' b IO r
iobind (IO f) g = I.M (IO (oneShot (\s -> case f s of (# a, b #) -> (# a, g b #)))) 
{-#INLINE iobind #-}

-- bindio io f = I.M (IO (oneShot (fmap f io)))

await' :: Monad m => Proxy () a y' y m a
await' = await
{-#INLINE [1] await' #-}

-- a similar rule for 'yield' has no effect
yield' :: Monad m => a -> Producer a m ()
yield' = yield  
{-#INLINE [1] yield' #-}

{-#RULES
   "await' hack" forall f . await' >>= f = I.Request () (oneShot f);

   "yield' hack" forall f x . yield' x >>= f = I.Respond x (oneShot f);
   
   "liftio hack" forall m f . liftio m >>= f = iobind m f
   #-}

sender :: Int -> Int -> IORef Int -> Producer Int IO ()
sender !m limit ref = when (m < limit) $ do 
    n <- liftIO (readIORef ref) 
    yield n
    sender (m+1) limit ref

receiver :: IORef Int -> Consumer Int IO ()
receiver ref = do
  n <- await   -- leaks with standard await
--  n <- await'     -- does not leak with the patched await
  liftIO (writeIORef ref $! n+1)
  receiver ref

main = do 
  ref <- newIORef 0
  let receiver_ref = receiver ref
      large = 10000000 :: Int
      larger = 2*large
  runEffect $ sender 0 large ref >-> receiver_ref  
  readIORef ref >>= print
  runEffect $ sender 1 large ref >-> receiver_ref 
  readIORef ref >>= print

-- -----------------------------
-- assorted other things to try
-- -----------------------------



-- this is a pipe not a sink and produces the same problem
receiver' :: IORef Int -> Pipe Int Int IO ()
receiver' ref = do
  -- liftIO $ getCurrentTime -- it leaks the same if the first constructor is M
  n <- await'
  -- n <- await'  -- does not leak with the patched `await`
  liftIO (writeIORef ref $! n+1)
  yield n
  receiver' ref


main' = do 
  ref <- newIORef 0
  let receiver_ref = receiver ref
      large = 10000000 :: Int
      larger = 2*large
  runEffect $ sender 0     large ref >-> receiver_ref >-> P.drain
  readIORef ref >>= print
  runEffect $ sender 1     large ref >-> receiver_ref >-> P.drain
  readIORef ref >>= print
 


-- it is typical that this works fine with full laziness and standard await
receiver_ :: IORef Int -> Consumer Int IO ()
receiver_ ref = loop where
  loop = do
    n <- await
    liftIO (writeIORef ref $! n+1)
    loop