oneShot rule for 'await >>= f'

pipes_oneshot.hs

{-#LANGUAGE BangPatterns #-}
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

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

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)
   #-}

large = 10000000 :: Int

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

receive :: IORef Int -> Consumer Int IO ()
receive ref = do
  n <- await
  liftIO $ writeIORef ref $! (n+1::Int)
  receive ref 
  
receive' :: IORef Int -> Consumer Int IO ()
receive' ref = do
  n <- await'
  liftIO $ writeIORef ref $! (n+1::Int)
  receive' ref 

main = do 
  ref <- newIORef 0
  let b = receive' ref
      large = 10000000 :: Int
      larger = 2*large
  runEffect $ emit large 0 ref >-> b
  readIORef ref >>= print
  runEffect $ emit larger 0 ref >-> b
  readIORef ref >>= print