Added a test for new generation, found a bug

dustin · Aug 24, 2023 · 3d93334 · 3d93334
1 parent b7b0277
commit 3d93334
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 15 deletions.
diff --git a/src/Data/Map/Strict/Expiring.hs b/src/Data/Map/Strict/Expiring.hs
@@ -68,7 +68,7 @@ updateLookupWithKey g f k m@Map{..} = case Map.updateLookupWithKey f' k map of
   (Nothing, _) -> (Nothing, m)
   (Just e, m') -> (Just (value e), m {
     map = m',
-    aging = Map.insert g (Set.singleton k) (removeAging (gen e) k aging)
+    aging = Map.insertWith (<>) g (Set.singleton k) (removeAging (gen e) k aging)
   })
   where
     f' _ e = (`Entry` g) <$> f k (value e)

diff --git a/test/ExpiringSpec.hs b/test/ExpiringSpec.hs
@@ -5,12 +5,13 @@
 
 module ExpiringSpec where
 
-import           Data.Foldable            (toList, traverse_)
+import           Data.Bool                (bool)
+import           Data.Foldable            (foldl', toList, traverse_)
+import           Data.Function            ((&))
 import qualified Data.Map.Strict          as Map
 import qualified Data.Map.Strict.Expiring as ExpiringMap
 import           Data.Set                 (Set)
 import qualified Data.Set                 as Set
-import Data.Bool (bool)
 
 import           Test.QuickCheck
 
@@ -24,13 +25,15 @@ data Mutation = Insert SomeKey Int
               | Delete SomeKey
               | Update SomeKey Int
               | UpdateNothing SomeKey
+              | NewGeneration
   deriving Show
 
 instance Arbitrary Mutation where
   arbitrary = oneof [Insert <$> arbitrary <*> arbitrary,
                      Delete <$> arbitrary,
                      Update <$> arbitrary <*> arbitrary,
-                     UpdateNothing <$> arbitrary
+                     UpdateNothing <$> arbitrary,
+                     pure NewGeneration
                      ]
 
 allOpTypes :: [String]
@@ -44,20 +47,23 @@ prop_expMapDoesMapStuff ops lookups =
   checkCoverage $
   ((`Map.lookup` massocs) <$> lookups) === ((`ExpiringMap.lookup` eassocs) <$> lookups)
   where
-    massocs = foldr applyOpM (mempty :: Map.Map SomeKey Int) ops
-    eassocs = foldr applyOpE (ExpiringMap.new 0) ops
+    massocs = foldl' (flip applyOpM) (mempty :: Map.Map SomeKey Int) ops
+    eassocs = foldl' applyOpE (ExpiringMap.new 0) ops
 
     applyOpM = \case
-      Insert k v -> Map.insert k v
-      Delete k -> Map.delete k
-      Update k v -> snd . Map.updateLookupWithKey (\_ _ -> Just v) k
+      Insert k v      -> Map.insert k v
+      Delete k        -> Map.delete k
+      Update k v      -> snd . Map.updateLookupWithKey (\_ _ -> Just v) k
       UpdateNothing k -> snd . Map.updateLookupWithKey (\_ _ -> Nothing) k
-
-    applyOpE = \case
-      Insert k v -> ExpiringMap.insert 1 k v
-      Delete k -> ExpiringMap.delete k
-      Update k v -> snd . ExpiringMap.updateLookupWithKey 1 (\k' _ -> bool Nothing (Just v) (k == k')) k
-      UpdateNothing k -> snd . ExpiringMap.updateLookupWithKey 1 (\_ _ -> Nothing) k
+      NewGeneration   -> const mempty
+
+    applyOpE m = \case
+      Insert k v      -> ExpiringMap.insert gen k v m
+      Delete k        -> ExpiringMap.delete k m
+      Update k v      -> snd $ ExpiringMap.updateLookupWithKey gen (\k' _ -> bool Nothing (Just v) (k == k')) k m
+      UpdateNothing k -> snd $ ExpiringMap.updateLookupWithKey gen (\_ _ -> Nothing) k m
+      NewGeneration   -> ExpiringMap.newGen (gen + 1) m
+      where gen = ExpiringMap.generation m
 
 prop_expiringMapWorks :: Int -> [Int] -> Property
 prop_expiringMapWorks baseGen keys = Just keys === traverse (`ExpiringMap.lookup` m) keys