More fixes for expiring map crossing generations

Emulating that with a simple dumb map data structure revealed a variety
of obscure bugs that would probably happen in production.

net-mqtt.cabal

@@ -148,6 +148,7 @@ test-suite mqtt-test
     , containers >=0.5.0 && <0.7
     , crypton-connection >=0.3.0
     , deepseq >=1.4.3.0 && <1.5
+    , lens
     , monad-loops >=0.4.3
     , mtl
     , net-mqtt

package.yaml

@@ -84,3 +84,4 @@ tests:
     - tasty-discover
     - checkers
     - mtl
+    - lens

src/Data/Map/Strict/Expiring.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE DeriveFunctor   #-}
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE TupleSections   #-}
 module Data.Map.Strict.Expiring (
     Map,
     new,
@@ -53,25 +54,26 @@ newGen g m
     | otherwise = m
 
 -- | 𝑂(log𝑛). Insert a new value into the map to expire after the given generation.
+-- alterF :: (Functor f, Ord k) => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
 insert :: (Ord k, Ord g) => g -> k -> a -> Map g k a -> Map g k a
 insert g _ _ m | g < generation m = m
-insert g k v m@Map{..} = m {
-  map = Map.insert k (Entry v g) map,
-  aging = Map.insertWith (<>) g (Set.singleton k) aging
-}
+insert g k v m@Map{..} = case Map.alterF (, Just (Entry v g)) k map of
+    (Just old, m') -> m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) (removeAging (gen old) k aging)}
+    (Nothing, m')  -> m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) aging}
 
 -- | 𝑂(log𝑛). Lookup and update.
 -- The function returns changed value, if it is updated. Returns the original key value if the map entry is deleted.
 updateLookupWithKey :: (Ord g, Ord k) => g -> (k -> a -> Maybe a) -> k -> Map g k a -> (Maybe a, Map g k a)
 updateLookupWithKey g _ _ m | g < generation m = (Nothing, m)
-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.insertWith (<>) g (Set.singleton k) (removeAging (gen e) k aging)
-  })
-  where
-    f' _ e = (`Entry` g) <$> f k (value e)
+updateLookupWithKey g f k m@Map{..} = case Map.alterF f' k map of
+    ((Nothing, _), m')   -> (Nothing, m)
+    ((Just old, Nothing), m')  -> (Just (value old), m{map=m', aging = removeAging (gen old) k aging})
+    ((Just old, Just new), m') -> (Just (value new), m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) (removeAging (gen old) k aging)})
+    where
+        f' Nothing = ((Nothing, Nothing), Nothing)
+        f' (Just e) = case f k (value e) of
+            Nothing -> ((Just e, Nothing), Nothing)
+            Just v  -> ((Just e, Just (Entry v g)), Just (Entry v g))
 
 removeAging :: (Ord g, Ord k) => g -> k -> Map.Map g (Set k) -> Map.Map g (Set k)
 removeAging g k = Map.update (nonNull . Set.delete k) g

test/ExpiringSpec.hs

@@ -1,10 +1,15 @@
-{-# LANGUAGE BlockArguments    #-}
-{-# LANGUAGE LambdaCase        #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE TupleSections     #-}
+{-# LANGUAGE BlockArguments             #-}
+{-# LANGUAGE DerivingStrategies         #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase                 #-}
+{-# LANGUAGE OverloadedStrings          #-}
+{-# LANGUAGE TemplateHaskell            #-}
+{-# LANGUAGE TupleSections              #-}
+
 
 module ExpiringSpec where
 
+import           Control.Lens
 import           Data.Bool                (bool)
 import           Data.Foldable            (foldl', toList, traverse_)
 import           Data.Function            ((&))
@@ -21,18 +26,28 @@ data SomeKey = Key1 | Key2 | Key3 | Key4 | Key5
 instance Arbitrary SomeKey where
   arbitrary = arbitraryBoundedEnum
 
-data Mutation = Insert SomeKey Int
+newtype GenOffset = GenOffset { getOffset :: Int }
+  deriving (Eq, Ord)
+  deriving newtype (Show, Num, Bounded)
+
+instance Arbitrary GenOffset where
+  arbitrary = GenOffset <$> choose (0, 5)
+  shrink = fmap GenOffset . shrink . getOffset
+
+data Mutation = Insert GenOffset SomeKey Int
               | Delete SomeKey
-              | Update SomeKey Int
-              | UpdateNothing SomeKey
-              | NewGeneration (Positive Int)
+              | Update GenOffset SomeKey Int
+              | UpdateNothing GenOffset SomeKey
+              | NewGeneration GenOffset
   deriving Show
 
+makePrisms ''Mutation
+
 instance Arbitrary Mutation where
-  arbitrary = oneof [Insert <$> arbitrary <*> arbitrary,
+  arbitrary = oneof [Insert <$> arbitrary <*> arbitrary <*> arbitrary,
                      Delete <$> arbitrary,
-                     Update <$> arbitrary <*> arbitrary,
-                     UpdateNothing <$> arbitrary,
+                     Update <$> arbitrary <*> arbitrary <*> arbitrary,
+                     UpdateNothing <$> arbitrary <*> arbitrary,
                      NewGeneration <$> arbitrary
                      ]
 
@@ -47,22 +62,26 @@ prop_doesMapStuff ops lookups =
   checkCoverage $
   ((`Map.lookup` massocs) <$> lookups) === ((`ExpiringMap.lookup` eassocs) <$> lookups)
   where
-    massocs = foldl' (flip applyOpM) (mempty :: Map.Map SomeKey Int) ops
+    massocs = degenerate $ foldl' applyOpM (0, mempty) 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
-      UpdateNothing k -> snd . Map.updateLookupWithKey (\_ _ -> Nothing) k
-      NewGeneration _ -> const mempty
+    -- The emulation stores the generation along with the value, so when we're done, we filter out anything old and fmap away the generation.
+    degenerate :: (GenOffset, Map.Map SomeKey (GenOffset, Int)) -> Map.Map SomeKey Int
+    degenerate (gen, m) = snd <$> Map.filter ((>= gen) . fst) m
+
+    applyOpM (gen, m) = \case
+      Insert g k v      -> (gen, Map.insert k (gen+g, v) m)
+      Delete k          -> (gen, Map.delete k m)
+      Update g k v      -> (gen, snd $ Map.updateLookupWithKey (\_ _ -> Just (gen+g, v)) k m)
+      UpdateNothing _ k -> (gen, snd $ Map.updateLookupWithKey (\_ _ -> Nothing) k m)
+      NewGeneration n   -> (gen + n, Map.filter ((>= gen + n) . fst) m)
 
     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 n -> ExpiringMap.newGen (gen + getPositive n) m
+      Insert g k v      -> ExpiringMap.insert (gen + g) k v m
+      Delete k          -> ExpiringMap.delete k m
+      Update g k v      -> snd $ ExpiringMap.updateLookupWithKey (gen + g) (\k' _ -> bool Nothing (Just v) (k == k')) k m
+      UpdateNothing g k -> snd $ ExpiringMap.updateLookupWithKey (gen + g) (\_ _ -> Nothing) k m
+      NewGeneration g   -> ExpiringMap.newGen (gen + g) m
       where gen = ExpiringMap.generation m
 
 prop_updateReturn :: Int -> Property