From 5b9de04692450f4922152199fecc6c220d279e37 Mon Sep 17 00:00:00 2001
From: Harendra Kumar <harendra@composewell.com>
Date: Tue, 20 May 2025 10:06:29 +0530
Subject: [PATCH] Change the signature of pollIntIORef

---
 core/src/Streamly/Internal/Data/IORef.hs      | 39 +++++++++++++------
 core/src/Streamly/Internal/Data/MutArray.hs   |  2 +-
 .../Internal/Data/Stream/Concurrent.hs        |  2 +-
 3 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/core/src/Streamly/Internal/Data/IORef.hs b/core/src/Streamly/Internal/Data/IORef.hs
index 7105a49fc8..32126b568a 100644
--- a/core/src/Streamly/Internal/Data/IORef.hs
+++ b/core/src/Streamly/Internal/Data/IORef.hs
@@ -32,11 +32,16 @@ module Streamly.Internal.Data.IORef
 
     -- Read
     , readIORef
+    , pollGenericIORef
+    , pollIORefInt
+
+    -- Deprecated
     , pollIntIORef
     )
 where
 
 #include "inline.hs"
+#include "deprecation.h"
 
 import Control.Monad.IO.Class (MonadIO(..))
 #if __GLASGOW_HASKELL__ >= 810
@@ -88,20 +93,30 @@ modifyIORef' var g = do
   x <- readIORef var
   writeIORef var (g x)
 
--- | Generate a stream by continuously reading the IORef.
---
--- This operation reads the IORef without any synchronization. It can be
--- assumed to be atomic because the IORef (MutableByteArray) is always aligned
--- to Int boundaries, we are assuming that compiler uses single instructions to
--- access the memory. It may read stale values though until caches are
--- synchronised in a multiprocessor architecture.
---
--- /Pre-release/
-{-# INLINE_NORMAL pollIntIORef #-}
-pollIntIORef :: (MonadIO m, Unbox a) => IORef a -> D.Stream m a
-pollIntIORef var = D.Stream step ()
+-- | Internal, do not use.
+{-# INLINE_NORMAL pollGenericIORef #-}
+pollGenericIORef :: (MonadIO m, Unbox a) => IORef a -> D.Stream m a
+pollGenericIORef var = D.Stream step ()
 
     where
 
     {-# INLINE_LATE step #-}
     step _ () = liftIO (readIORef var) >>= \x -> return $ D.Yield x ()
+
+{-# DEPRECATED pollIntIORef "Use pollIORefInt instead." #-}
+{-# INLINE_NORMAL pollIntIORef #-}
+pollIntIORef :: (MonadIO m, Unbox a) => IORef a -> D.Stream m a
+pollIntIORef = pollGenericIORef
+
+-- | Generate a stream by continuously reading the IORef.
+--
+-- This operation reads the IORef without any synchronization. It can be
+-- assumed to be atomic because the size fits into machine register size. We
+-- are assuming that compiler uses single instructions to access the memory. It
+-- may read stale values though until caches are synchronised in a
+-- multiprocessor architecture.
+--
+-- /Pre-release/
+{-# INLINE_NORMAL pollIORefInt #-}
+pollIORefInt :: MonadIO m => IORef Int -> D.Stream m Int
+pollIORefInt = pollGenericIORef
diff --git a/core/src/Streamly/Internal/Data/MutArray.hs b/core/src/Streamly/Internal/Data/MutArray.hs
index 371d397c76..f58532e4d6 100644
--- a/core/src/Streamly/Internal/Data/MutArray.hs
+++ b/core/src/Streamly/Internal/Data/MutArray.hs
@@ -402,7 +402,7 @@ type IORef = IORef.IORef
 
 {-# DEPRECATED pollIntIORef "Use pollIntIORef from MutByteArray module." #-}
 pollIntIORef :: (MonadIO m, Unbox a) => IORef a -> Stream m a
-pollIntIORef = IORef.pollIntIORef
+pollIntIORef = IORef.pollGenericIORef
 
 {-# DEPRECATED newIORef "Use newIORef from MutByteArray module." #-}
 newIORef :: forall a. Unbox a => a -> IO (IORef a)
diff --git a/src/Streamly/Internal/Data/Stream/Concurrent.hs b/src/Streamly/Internal/Data/Stream/Concurrent.hs
index 3212b4e330..1b3d36d398 100644
--- a/src/Streamly/Internal/Data/Stream/Concurrent.hs
+++ b/src/Streamly/Internal/Data/Stream/Concurrent.hs
@@ -746,7 +746,7 @@ parTapCount predicate fld (D.Stream step state) = D.Stream step' Nothing
         countVar <- liftIO $ Unboxed.newIORef (0 :: Int)
         tid <- forkManaged
             $ void $ fld
-            $ Unboxed.pollIntIORef countVar
+            $ Unboxed.pollIORefInt countVar
         return $ Skip (Just (countVar, tid, state))
 
     step' gst (Just (countVar, tid, st)) = do