Add upsert

In comparison to alter this avoids DeleteEntry constraints
and, when the function does not inline, eliminates a redundant Maybe box.

bench/Main.hs

@@ -241,6 +241,7 @@ utilities n = do
     hSize <- vh n
     hMember <- vh n
     hFindWithDefault <- vh n
+    hUpsert <- vh n
     hAlter <- vh n
     hAlterM <- vh n
     hUnion1 <- vh n
@@ -264,6 +265,7 @@ utilities n = do
         , bench "size" $ nfIO (bhusize hSize)
         , bench "member" $ nfIO (bhumember n hMember)
         , bench "findWithDefault" $ nfIO (bhufindWithDefault n hFindWithDefault)
+        , bench "upsert" $ nfIO (bhuupsert n hUpsert)
         , bench "alter" $ nfIO (bhualter n hAlter)
         , bench "alterM" $ nfIO (bhualterM n hAlterM)
         , bench "union" $ nfIO (bhuunion hUnion1 hUnion2)
@@ -318,6 +320,11 @@ bhufindWithDefault n ht = do
               | otherwise = return ()
     go 0
 
+bhuupsert n ht = do
+    let go !i | i <= n = VH.upsert ht (maybe minBound succ) i >> go (i + 1)
+              | otherwise = return ()
+    go 0
+
 bhualter n ht = do
     let go !i | i <= n = VH.alter ht (fmap succ) i >> go (i + 1)
               | otherwise = return ()

src/Data/Vector/Hashtables.hs

@@ -31,6 +31,7 @@ module Data.Vector.Hashtables
     , lookup'
     , insert
     , delete
+    , upsert
     , alter
     , alterM
 

src/Data/Vector/Hashtables/Internal.hs

@@ -541,6 +541,46 @@ findWithDefault
 findWithDefault ht v k = return . fromMaybe v =<< at' ht k
 {-# INLINE findWithDefault #-}
 
+-- | /O(1)/ in the best case, /O(n)/ in the worst case.
+-- The expression (@'upsert' ht f k@) updates or inserts the value @x@ at @k@.
+--
+-- > let f _ = "c"
+-- > ht <- fromList [(5,"a"), (3,"b")]
+-- > upsert ht f 7
+-- > toList ht
+-- > [(3, "b"), (5, "a"), (7, "c")]
+--
+-- > ht <- fromList [(5,"a"), (3,"b")]
+-- > upsert ht f 5
+-- > toList ht
+-- > [(3, "b"), (5, "c")]
+--
+upsert
+  :: ( MVector ks k, MVector vs v
+     , PrimMonad m, Hashable k, Eq k
+     )
+  => Dictionary (PrimState m) ks k vs v -> (Maybe v -> v) -> k -> m ()
+upsert ht f k = do
+  d@Dictionary{..} <- readMutVar . getDRef $ ht
+  let
+      hashCode' = hash k .&. mask
+      targetBucket = hashCode' `rem` A.length buckets
+
+      onFound' value' dict i = insertWithIndex targetBucket hashCode' k value' (getDRef ht) dict i
+
+      onFound dict i = do
+        d'@Dictionary{..} <- readMutVar . getDRef $ dict
+        v <- value !~ i
+        onFound' (f (Just v)) d' i
+
+      onNothing dict = do
+        d' <- readMutVar . getDRef $ dict
+        onFound' (f Nothing) d' (-1)
+
+  void $ atWithOrElse ht k onFound onNothing
+
+{-# INLINE upsert #-}
+
 -- | /O(1)/ in the best case, /O(n)/ in the worst case.
 -- The expression (@'alter' ht f k@) alters the value @x@ at @k@, or absence thereof.
 -- 'alter' can be used to insert, delete, or update a value in a 'Dictionary'.

test/Data/Vector/HashTablesSpec.hs

@@ -117,6 +117,8 @@ class HashTableTest ks vs where
 
   testAlter :: VH.Dictionary (PrimState IO) ks Int vs Int -> (Maybe Int -> Maybe Int) -> Int -> IO ()
 
+  testUpsert :: VH.Dictionary (PrimState IO) ks Int vs Int -> (Maybe Int -> Int) -> Int -> IO ()
+
   testUnion
     :: VH.Dictionary (PrimState IO) ks Int vs Int
     -> VH.Dictionary (PrimState IO) ks Int vs Int
@@ -177,6 +179,11 @@ mkSpec ksp vsp = describe (specDescription ksp vsp) $
       it "when altering is nothing - key deleted from table" $ property prop_alterDelete
 
       it "when altering is just a result - key updated with result" $ property prop_alterUpdate
+
+      it "when upserting a new key - key is set to value" $ property prop_upsertInsert
+
+      it "when upserting an existing key - key updated with result" $ property prop_upsertUpdate
+
       it "intersection + symmetric difference of two tables is equal to union of two tables" $ property prop_union
 
   where
@@ -302,6 +309,21 @@ mkSpec ksp vsp = describe (specDescription ksp vsp) $
       v <- testAt ht x
       v `shouldBe` (negate y)
 
+    prop_upsertInsert :: HashTableTest ks vs => (Int, Int) -> IO ()
+    prop_upsertInsert (x, y) = do
+      ht <- testInit (Proxy @ks) (Proxy @vs) 10
+      testUpsert ht (maybe 0 negate) x
+      v <- testAt ht x
+      v `shouldBe` 0
+
+    prop_upsertUpdate :: HashTableTest ks vs => (Int, Int) -> IO ()
+    prop_upsertUpdate (x, y) = do
+      ht <- testInit (Proxy @ks) (Proxy @vs) 10
+      testInsert ht x y
+      testUpsert ht (maybe 0 negate) x
+      v <- testAt ht x
+      v `shouldBe` (negate y)
+
     prop_union :: Positive Int -> Property
     prop_union (Positive n) = forAll (twoListsN n) $ \(xs, ys) -> do
       ht1 <- testFromList (Proxy @ks) (Proxy @vs) xs
@@ -337,6 +359,7 @@ instance HashTableTest M.MVector M.MVector where
   testNull = VH.null
   testMember = VH.member
   testAlter = VH.alter
+  testUpsert = VH.upsert
   testUnion = VH.union
   testDifference = VH.difference
   testIntersection = VH.intersection
@@ -361,6 +384,7 @@ instance HashTableTest SM.MVector SM.MVector where
   testNull = VH.null
   testMember = VH.member
   testAlter = VH.alter
+  testUpsert = VH.upsert
   testUnion = VH.union
   testDifference = VH.difference
   testIntersection = VH.intersection
@@ -385,6 +409,7 @@ instance HashTableTest SM.MVector M.MVector where
   testNull = VH.null
   testMember = VH.member
   testAlter = VH.alter
+  testUpsert = VH.upsert
   testUnion = VH.union
   testDifference = VH.difference
   testIntersection = VH.intersection
@@ -409,6 +434,7 @@ instance HashTableTest M.MVector UM.MVector where
   testNull = VH.null
   testMember = VH.member
   testAlter = VH.alter
+  testUpsert = VH.upsert
   testUnion = VH.union
   testDifference = VH.difference
   testIntersection = VH.intersection