[#48] Add adjustWithKey and adjust for TMap (#50)

* [#48] Add adjustWithKey and adjust for TMap

* Fix space

* Improve adjust function

Rename to adjust for TypeRepMap

CHANGELOG.md

@@ -6,9 +6,11 @@ The change log is available [on GitHub][2].
 
 # 0.3.0
 
-* [#43](https://github.com/kowainik/typerep-map/issues/43):
+* [#46](https://github.com/kowainik/typerep-map/issues/46):
   Make `Show` instance for `TypeRepMap` show keys.
   Add `keys` function.
+* [#48](https://github.com/kowainik/typerep-map/issues/48):
+  Add `adjust` function for `TypeRepMap` and  `TMap`.
 
 # 0.2.0
 

src/Data/TMap.hs

@@ -34,6 +34,7 @@ module Data.TMap
        , unionWith
        , union
        , map
+       , adjust
 
          -- * Query
        , lookup
@@ -157,8 +158,13 @@ keys = F.keys
 
 -- | Map a function over the values.
 map :: (forall a. Typeable a => a -> a) -> TMap -> TMap
-map f = F.hoistWithKey (fIdentity f)
-  where
-    fIdentity :: forall a. Typeable a => (a -> a) -> Identity a -> Identity a
-    fIdentity = coerce
+map f = F.hoistWithKey (liftToIdentity f)
 {-# INLINE map #-}
+
+-- | Update a value with the result of the provided function.
+adjust :: Typeable a => (a -> a) -> TMap -> TMap
+adjust f = F.adjust (liftToIdentity f)
+{-# INLINE adjust #-}
+
+liftToIdentity :: forall a. Typeable a => (a -> a) -> Identity a -> Identity a
+liftToIdentity = coerce

src/Data/TypeRepMap.hs

@@ -51,6 +51,7 @@ module Data.TypeRepMap
          -- * Modification
        , insert
        , delete
+       , adjust
        , hoist
        , hoistA
        , hoistWithKey

src/Data/TypeRepMap/Internal.hs

@@ -23,13 +23,15 @@ module Data.TypeRepMap.Internal where
 
 import Prelude hiding (lookup)
 
+import Control.Monad.ST (runST)
 import Control.Monad.Zip (mzip)
 import Data.Function (on)
 import Data.IntMap.Strict (IntMap)
 import Data.Kind (Type)
 import Data.List (intercalate, nubBy)
 import Data.Maybe (fromJust)
-import Data.Primitive.Array (Array, indexArray, mapArray')
+import Data.Primitive.Array (Array, freezeArray, indexArray, mapArray', readArray, sizeofArray,
+                             thawArray, writeArray)
 import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
 import Data.Semigroup (Semigroup (..))
 import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
@@ -162,6 +164,28 @@ delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeR
 delete = fromTriples . deleteByFst (typeFp @a) . toTriples
 {-# INLINE delete #-}
 
+{- |
+Update a value at a specific key with the result of the provided function. When
+the key is not a member of the map, the original map is returned.
+
+>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
+>>> lookup @String $ adjust (fmap (++ "ww")) trmap
+Just (Identity "aww")
+-}
+adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
+adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
+    Nothing -> tr
+    Just i  -> tr {trAnys = changeAnyArr i (trAnys tr)}
+  where
+    changeAnyArr :: Int -> Array Any -> Array Any
+    changeAnyArr i trAs = runST $ do
+        let n = sizeofArray trAs
+        mutArr <- thawArray trAs 0 n
+        a <- toAny . fun . fromAny <$> readArray mutArr i
+        writeArray mutArr i a
+        freezeArray mutArr 0 n
+{-# INLINE adjust #-}
+
 {- | Map over the elements of a 'TypeRepMap'.
 
 >>> tm = insert (Identity True) $ one (Identity 'a')