tweaks

src/TimerWheel.hs

@@ -20,7 +20,6 @@ module TimerWheel
 where
 
 import qualified Data.Atomics as Atomics
-import Data.Foldable (traverse_)
 import Data.Primitive.Array (MutableArray)
 import qualified Data.Primitive.Array as Array
 import GHC.Base (RealWorld)
@@ -337,7 +336,7 @@ timerBucketInsert timestamp timer =
 -- Fire every timer in a bucket, in order
 timerBucketFire :: Counter -> TimerBucket -> IO ()
 timerBucketFire numTimers =
-  traverse_ \case
+  TimestampMap.foreach \_ -> \case
     Timers1 timer -> timerActionThenDecr timer
     TimersN timers -> foldr f (pure ()) timers
   where
@@ -346,6 +345,7 @@ timerBucketFire numTimers =
       acc
       timerActionThenDecr timer
 
+    timerActionThenDecr :: Timer -> IO ()
     timerActionThenDecr timer = do
       timerAction timer
       decrCounter_ numTimers

src/TimerWheel/Internal/Prelude.hs

@@ -10,7 +10,7 @@ where
 import Control.Monad as X (when)
 import Data.Coerce as X (coerce)
 import Data.Fixed (E6, Fixed)
-import Data.IORef as X (newIORef, readIORef, writeIORef)
+import Data.IORef as X (IORef, newIORef, readIORef, writeIORef)
 import Data.Map as X (Map)
 import Data.Word as X (Word64)
 import GHC.Generics as X (Generic)

src/TimerWheel/Internal/TimestampMap.hs

@@ -4,6 +4,7 @@ module TimerWheel.Internal.TimestampMap
   ( TimestampMap,
     delete,
     empty,
+    foreach,
     insert,
     lookup,
     null,
@@ -31,18 +32,6 @@ type Prefix = Word64
 
 type Mask = Word64
 
-instance Foldable TimestampMap where
-  -- Use lambda t to be inlinable with two arguments only.
-  foldr f z = \case
-    Bin _ _ l r -> go (go z r) l
-    t -> go z t
-    where
-      go acc = \case
-        Nil -> acc
-        Tip _ x -> f x acc
-        Bin _ _ l r -> go (go acc r) l
-  {-# INLINE foldr #-}
-
 bin :: Prefix -> Mask -> TimestampMap a -> TimestampMap a -> TimestampMap a
 bin _ _ l Nil = l
 bin _ _ Nil r = r
@@ -79,6 +68,15 @@ empty =
   Nil
 {-# INLINE empty #-}
 
+foreach :: (Timestamp -> a -> IO ()) -> TimestampMap a -> IO ()
+foreach f =
+  go
+  where
+    go = \case
+      Nil -> pure ()
+      Tip k x -> f (coerce @Word64 @Timestamp k) x
+      Bin _ _ l r -> go l >> go r
+
 insert :: forall a. Timestamp -> a -> TimestampMap a -> TimestampMap a
 insert = coerce @(Word64 -> a -> TimestampMap a -> TimestampMap a) insert_
 {-# INLINE insert #-}