From 450e6ab697be2cd2263d75dc5625e85910c4dd9b Mon Sep 17 00:00:00 2001
From: Mitchell Rosen <mitchellwrosen@gmail.com>
Date: Sun, 1 Oct 2023 15:22:59 -0400
Subject: [PATCH] move some code around

---
 src/TimerWheel.hs                             |  36 ++---
 src/TimerWheel/Internal/Prelude.hs            |  17 +-
 .../TimerWheel/Internal/WordMap.hs}           | 149 ++++++++----------
 test/{ => timer-wheel}/Main.hs                |   0
 timer-wheel.cabal                             |   8 +-
 5 files changed, 89 insertions(+), 121 deletions(-)
 rename src/{TimerWheel/Internal/TimestampMap.hs => wordmap/TimerWheel/Internal/WordMap.hs} (50%)
 rename test/{ => timer-wheel}/Main.hs (100%)

diff --git a/src/TimerWheel.hs b/src/TimerWheel.hs
index b294b4b..f68d294 100644
--- a/src/TimerWheel.hs
+++ b/src/TimerWheel.hs
@@ -44,8 +44,8 @@ import qualified TimerWheel.Internal.Nanoseconds as Nanoseconds
 import TimerWheel.Internal.Prelude
 import TimerWheel.Internal.Timestamp (Timestamp)
 import qualified TimerWheel.Internal.Timestamp as Timestamp
-import TimerWheel.Internal.TimestampMap (TimestampMap)
-import qualified TimerWheel.Internal.TimestampMap as TimestampMap
+import TimerWheel.Internal.WordMap (WordMap)
+import qualified TimerWheel.Internal.WordMap as WordMap
 
 -- | A timer wheel is a vector-of-collections-of timers to fire. Timers may be one-shot or recurring, and may be
 -- scheduled arbitrarily far in the future.
@@ -126,7 +126,7 @@ create ::
   -- | ​
   IO TimerWheel
 create scope (Config spokes0 resolution0) = do
-  buckets <- Array.newArray spokes TimestampMap.empty
+  buckets <- Array.newArray spokes WordMap.empty
   numTimers <- newCounter
   timerIdSupply <- newCounter
   Ki.fork_ scope (runTimerReaperThread buckets numTimers resolution)
@@ -290,14 +290,14 @@ timestampToIndex buckets resolution timestamp =
 -- Timer bucket operations
 
 type TimerBucket =
-  TimestampMap Timers
+  WordMap Timers
 
 timerBucketDelete :: Timestamp -> TimerId -> TimerBucket -> Maybe TimerBucket
-timerBucketDelete timestamp timerId bucket =
-  case TimestampMap.lookup timestamp bucket of
+timerBucketDelete (coerce @Timestamp @Word64 -> timestamp) timerId bucket =
+  case WordMap.lookup timestamp bucket of
     Nothing -> Nothing
     Just (Timers1 timer)
-      | timerId == getTimerId timer -> Just $! TimestampMap.delete timestamp bucket
+      | timerId == getTimerId timer -> Just $! WordMap.delete timestamp bucket
       | otherwise -> Nothing
     Just (TimersN timers0) ->
       case timersDelete timerId timers0 of
@@ -307,11 +307,11 @@ timerBucketDelete timestamp timerId bucket =
                 case timers1 of
                   [timer] -> Timers1 timer
                   _ -> TimersN timers1
-           in Just $! TimestampMap.insert timestamp timers2 bucket
+           in Just $! WordMap.insert timestamp timers2 bucket
 
 timerBucketInsert :: Timestamp -> Timer0 -> TimerBucket -> TimerBucket
 timerBucketInsert timestamp timer =
-  TimestampMap.upsert timestamp (Timers1 timer) \case
+  WordMap.upsert (coerce @Timestamp @Word64 timestamp) (Timers1 timer) \case
     Timers1 old -> TimersN [timer, old]
     TimersN old -> TimersN (timer : old)
 
@@ -372,15 +372,15 @@ atomicMaybeModifyArray buckets index doDelete = do
           if success then pure True else loop ticket1
 
 atomicExtractExpiredTimersFromBucket :: MutableArray RealWorld TimerBucket -> Int -> Timestamp -> IO TimerBucket
-atomicExtractExpiredTimersFromBucket buckets index now = do
+atomicExtractExpiredTimersFromBucket buckets index (coerce @Timestamp @Word64 -> now) = do
   ticket0 <- Atomics.readArrayElem buckets index
   loop ticket0
   where
     loop :: Atomics.Ticket TimerBucket -> IO TimerBucket
     loop ticket = do
-      let Pair expired bucket1 = TimestampMap.splitL now (Atomics.peekTicket ticket)
-      if TimestampMap.null expired
-        then pure TimestampMap.empty
+      let WordMap.Pair expired bucket1 = WordMap.splitL now (Atomics.peekTicket ticket)
+      if WordMap.null expired
+        then pure WordMap.empty
         else do
           (success, ticket1) <- Atomics.casArrayElem buckets index ticket bucket1
           if success then pure expired else loop ticket1
@@ -515,15 +515,15 @@ runTimerReaperThread buckets numTimers resolution = do
       where
         fireTimerBucket :: TimerBucket -> IO ()
         fireTimerBucket bucket0 =
-          case TimestampMap.pop bucket0 of
-            TimestampMap.PopNada -> pure ()
-            TimestampMap.PopAlgo timestamp timers bucket1 ->
+          case WordMap.pop bucket0 of
+            WordMap.PopNada -> pure ()
+            WordMap.PopAlgo timestamp timers bucket1 ->
               case timers of
                 Timers1 timer -> do
-                  expired2 <- fireTimer bucket1 timestamp timer
+                  expired2 <- fireTimer bucket1 (coerce @Word64 @Timestamp timestamp) timer
                   fireTimerBucket expired2
                 TimersN timers1 -> do
-                  bucket2 <- fireTimers bucket1 timestamp timers1
+                  bucket2 <- fireTimers bucket1 (coerce @Word64 @Timestamp timestamp) timers1
                   fireTimerBucket bucket2
 
         fireTimers :: TimerBucket -> Timestamp -> [Timer0] -> IO TimerBucket
diff --git a/src/TimerWheel/Internal/Prelude.hs b/src/TimerWheel/Internal/Prelude.hs
index 5f73a4b..9c12a14 100644
--- a/src/TimerWheel/Internal/Prelude.hs
+++ b/src/TimerWheel/Internal/Prelude.hs
@@ -1,8 +1,5 @@
 module TimerWheel.Internal.Prelude
-  ( Pair (..),
-    mapPairL,
-    mapPairR,
-    Seconds,
+  ( Seconds,
     module X,
   )
 where
@@ -15,18 +12,6 @@ import Data.Word as X (Word64)
 import GHC.Generics as X (Generic)
 import Prelude as X hiding (lookup, null)
 
--- A strict pair
-data Pair a b
-  = Pair !a !b
-
-mapPairL :: (a -> b) -> Pair a x -> Pair b x
-mapPairL f (Pair x y) = Pair (f x) y
-{-# INLINE mapPairL #-}
-
-mapPairR :: (a -> b) -> Pair x a -> Pair x b
-mapPairR f (Pair x y) = Pair x (f y)
-{-# INLINE mapPairR #-}
-
 -- | A number of seconds, with nanosecond precision.
 --
 -- You can use numeric literals to construct a value of this type, e.g. @0.5@.
diff --git a/src/TimerWheel/Internal/TimestampMap.hs b/src/wordmap/TimerWheel/Internal/WordMap.hs
similarity index 50%
rename from src/TimerWheel/Internal/TimestampMap.hs
rename to src/wordmap/TimerWheel/Internal/WordMap.hs
index 485d8fd..34a5b93 100644
--- a/src/TimerWheel/Internal/TimestampMap.hs
+++ b/src/wordmap/TimerWheel/Internal/WordMap.hs
@@ -1,10 +1,10 @@
--- A Word64Map with code cribbed from containers and GHC
+-- An IntMap-like container with code cribbed from containers and GHC
+--
 -- TODO: proper code attribution
-module TimerWheel.Internal.TimestampMap
-  ( TimestampMap,
+module TimerWheel.Internal.WordMap
+  ( WordMap (..),
     delete,
     empty,
-    foreach,
     insert,
     lookup,
     null,
@@ -12,95 +12,79 @@ module TimerWheel.Internal.TimestampMap
     pop,
     splitL,
     upsert,
+
+    -- * Strict pair
+    Pair (..),
   )
 where
 
 import Data.Bits
 import Data.Word
-import TimerWheel.Internal.Nanoseconds (Nanoseconds (..))
 import TimerWheel.Internal.Prelude
-import TimerWheel.Internal.Timestamp (Timestamp (..))
-
-data TimestampMap a
-  = Bin
-      {-# UNPACK #-} !Prefix
-      {-# UNPACK #-} !Mask
-      !(TimestampMap a)
-      !(TimestampMap a)
+
+data WordMap a
+  = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(WordMap1 a) !(WordMap1 a)
   | Tip {-# UNPACK #-} !Word64 a
   | Nil
 
+-- A non-empty word map.
+type WordMap1 a = WordMap a
+
 type Prefix = Word64
 
 type Mask = Word64
 
-bin :: Prefix -> Mask -> TimestampMap a -> TimestampMap a -> TimestampMap a
+bin :: Prefix -> Mask -> WordMap a -> WordMap a -> WordMap a
 bin _ _ l Nil = l
 bin _ _ Nil r = r
 bin p m l r = Bin p m l r
 {-# INLINE bin #-}
 
-binCheckLeft :: Prefix -> Mask -> TimestampMap a -> TimestampMap a -> TimestampMap a
+binCheckLeft :: Prefix -> Mask -> WordMap a -> WordMap1 a -> WordMap1 a
 binCheckLeft _ _ Nil r = r
 binCheckLeft p m l r = Bin p m l r
 {-# INLINE binCheckLeft #-}
 
-binCheckRight :: Prefix -> Mask -> TimestampMap a -> TimestampMap a -> TimestampMap a
+binCheckRight :: Prefix -> Mask -> WordMap1 a -> WordMap a -> WordMap1 a
 binCheckRight _ _ l Nil = l
 binCheckRight p m l r = Bin p m l r
 {-# INLINE binCheckRight #-}
 
-delete :: forall a. Timestamp -> TimestampMap a -> TimestampMap a
-delete = coerce @(Word64 -> TimestampMap a -> TimestampMap a) delete_
-{-# INLINE delete #-}
-
-delete_ :: Word64 -> TimestampMap a -> TimestampMap a
-delete_ !k = \case
-  t@(Bin p m l r)
-    | nomatch k p m -> t
-    | zero k m -> binCheckLeft p m (delete_ k l) r
-    | otherwise -> binCheckRight p m l (delete_ k r)
-  t@(Tip kx _)
-    | k == kx -> Nil
-    | otherwise -> t
-  Nil -> Nil
-
-empty :: TimestampMap a
+delete :: Word64 -> WordMap a -> WordMap a
+delete !k t =
+  case t of
+    Bin p m l r
+      | nomatch k p m -> t
+      | zero k m -> binCheckLeft p m (delete k l) r
+      | otherwise -> binCheckRight p m l (delete k r)
+    Tip kx _
+      | k == kx -> Nil
+      | otherwise -> t
+    Nil -> Nil
+
+empty :: WordMap a
 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 #-}
-
-insert_ :: Word64 -> a -> TimestampMap a -> TimestampMap a
-insert_ !k !x t =
+insert :: Word64 -> a -> WordMap a -> WordMap a
+insert !k !x t =
   case t of
     Bin p m l r
       | nomatch k p m -> link k (Tip k x) p t
-      | zero k m -> Bin p m (insert_ k x l) r
-      | otherwise -> Bin p m l (insert_ k x r)
+      | zero k m -> Bin p m (insert k x l) r
+      | otherwise -> Bin p m l (insert k x r)
     Tip ky _
       | k == ky -> Tip k x
       | otherwise -> link k (Tip k x) ky t
     Nil -> Tip k x
 
-link :: Prefix -> TimestampMap a -> Prefix -> TimestampMap a -> TimestampMap a
+link :: Prefix -> WordMap a -> Prefix -> WordMap a -> WordMap a
 link p1 t1 p2 = linkWithMask (branchMask p1 p2) p1 t1
 {-# INLINE link #-}
 
 -- `linkWithMask` is useful when the `branchMask` has already been computed
-linkWithMask :: Mask -> Prefix -> TimestampMap a -> TimestampMap a -> TimestampMap a
+linkWithMask :: Mask -> Prefix -> WordMap a -> WordMap a -> WordMap a
 linkWithMask m p0 t1 t2
   | zero p0 m = Bin p m t1 t2
   | otherwise = Bin p m t2 t1
@@ -108,42 +92,33 @@ linkWithMask m p0 t1 t2
     p = mask p0 m
 {-# INLINE linkWithMask #-}
 
-lookup :: forall a. Timestamp -> TimestampMap a -> Maybe a
-lookup = coerce @(Word64 -> TimestampMap a -> Maybe a) lookup_
-{-# INLINE lookup #-}
-
-lookup_ :: forall a. Word64 -> TimestampMap a -> Maybe a
-lookup_ !k =
+lookup :: forall a. Word64 -> WordMap a -> Maybe a
+lookup !k =
   go
   where
-    go :: TimestampMap a -> Maybe a
+    go :: WordMap a -> Maybe a
     go = \case
       Bin _ m l r -> go (if zero k m then l else r)
       Tip kx x -> if k == kx then Just x else Nothing
       Nil -> Nothing
 
-null :: TimestampMap a -> Bool
+null :: WordMap a -> Bool
 null = \case
   Nil -> True
   _ -> False
 {-# INLINE null #-}
 
-singleton :: Word64 -> a -> TimestampMap a
+singleton :: Word64 -> a -> WordMap a
 singleton k !x = Tip k x
 {-# INLINE singleton #-}
 
--- @splitL k xs@ splits @xs@ into @ys@ and @zs@, where every timestamp in @ys@ is less than or equal to @k@, and every
--- timestamp in @zs@ is greater than @k@.
-splitL :: forall a. Timestamp -> TimestampMap a -> Pair (TimestampMap a) (TimestampMap a)
-splitL = coerce @(Word64 -> TimestampMap a -> Pair (TimestampMap a) (TimestampMap a)) splitL_
-{-# INLINE splitL #-}
-
-splitL_ :: forall a. Word64 -> TimestampMap a -> Pair (TimestampMap a) (TimestampMap a)
-splitL_ k = \case
-  Bin p m l r | m < 0 -> mapPairL (\ll -> bin p m ll r) (go l)
-  t -> go t
+-- @splitL k xs@ splits @xs@ into @ys@ and @zs@, where every value in @ys@ is less than or equal to @k@, and every value
+-- in @zs@ is greater than @k@.
+splitL :: forall a. Word64 -> WordMap a -> Pair (WordMap a) (WordMap a)
+splitL k =
+  go
   where
-    go :: TimestampMap a -> Pair (TimestampMap a) (TimestampMap a)
+    go :: WordMap a -> Pair (WordMap a) (WordMap a)
     go = \case
       t@(Bin p m l r)
         | nomatch k p m ->
@@ -159,9 +134,9 @@ splitL_ k = \case
 
 data Pop a
   = PopNada
-  | PopAlgo !Timestamp !a !(TimestampMap a)
+  | PopAlgo !Word64 !a !(WordMap a)
 
-pop :: TimestampMap a -> Pop a
+pop :: WordMap a -> Pop a
 pop = \case
   Nil -> PopNada
   Bin p m l r0 | m < 0 ->
@@ -171,25 +146,21 @@ pop = \case
   t -> pop1 t
 {-# INLINE pop #-}
 
-pop1 :: TimestampMap a -> Pop a
+pop1 :: WordMap1 a -> Pop a
 pop1 = \case
   Bin p m l0 r ->
     case pop1 l0 of
       PopAlgo k x l1 -> PopAlgo k x (binCheckLeft p m l1 r)
       PopNada -> undefined
-  Tip k x -> PopAlgo (coerce @Word64 @Timestamp k) x Nil
+  Tip k x -> PopAlgo k x Nil
   Nil -> undefined
 
-upsert :: forall a. Timestamp -> a -> (a -> a) -> TimestampMap a -> TimestampMap a
-upsert = coerce @(Word64 -> a -> (a -> a) -> TimestampMap a -> TimestampMap a) upsert_
-{-# INLINE upsert #-}
-
-upsert_ :: Word64 -> a -> (a -> a) -> TimestampMap a -> TimestampMap a
-upsert_ !k x f = \case
+upsert :: Word64 -> a -> (a -> a) -> WordMap a -> WordMap a
+upsert !k x f = \case
   t@(Bin p m l r)
     | nomatch k p m -> link k (singleton k x) p t
-    | zero k m -> Bin p m (upsert_ k x f l) r
-    | otherwise -> Bin p m l (upsert_ k x f r)
+    | zero k m -> Bin p m (upsert k x f l) r
+    | otherwise -> Bin p m l (upsert k x f r)
   t@(Tip ky y)
     | k == ky -> Tip k $! f y
     | otherwise -> link k (singleton k x) ky t
@@ -227,3 +198,15 @@ branchMask p1 p2 =
 highestBitMask :: Word64 -> Word64
 highestBitMask w = unsafeShiftL 1 (63 - countLeadingZeros w)
 {-# INLINE highestBitMask #-}
+
+-- A strict pair
+data Pair a b
+  = Pair !a !b
+
+mapPairL :: (a -> b) -> Pair a x -> Pair b x
+mapPairL f (Pair x y) = Pair (f x) y
+{-# INLINE mapPairL #-}
+
+mapPairR :: (a -> b) -> Pair x a -> Pair x b
+mapPairR f (Pair x y) = Pair x (f y)
+{-# INLINE mapPairR #-}
diff --git a/test/Main.hs b/test/timer-wheel/Main.hs
similarity index 100%
rename from test/Main.hs
rename to test/timer-wheel/Main.hs
diff --git a/timer-wheel.cabal b/timer-wheel.cabal
index 1f0f48a..1ee9be4 100644
--- a/timer-wheel.cabal
+++ b/timer-wheel.cabal
@@ -72,22 +72,22 @@ library
     primitive ^>= 0.7 || ^>= 0.8,
   exposed-modules:
     TimerWheel
-  hs-source-dirs: src
+  hs-source-dirs: src, src/wordmap
   other-modules:
     TimerWheel.Internal.Counter
     TimerWheel.Internal.Nanoseconds
     TimerWheel.Internal.Prelude
     TimerWheel.Internal.Timestamp
-    TimerWheel.Internal.TimestampMap
+    TimerWheel.Internal.WordMap
 
-test-suite tests
+test-suite timer-wheel-tests
   import: component
   build-depends:
     ki,
     random ^>= 1.2,
     timer-wheel,
   ghc-options: -threaded -with-rtsopts=-N2
-  hs-source-dirs: test
+  hs-source-dirs: test/timer-wheel
   main-is: Main.hs
   type: exitcode-stdio-1.0