Skip to content

Instantly share code, notes, and snippets.

@michaelt michaelt/oneshot.hs
Created Oct 31, 2016

Embed
What would you like to do?
oneShot rule for `await`
{-#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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.