diff --git a/CHANGELOG.md b/CHANGELOG.md index 9dc4f80..fd3b4a3 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/Data/TMap.hs b/src/Data/TMap.hs index 6ae2b9d..b650ac6 100644 --- a/src/Data/TMap.hs +++ b/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 diff --git a/src/Data/TypeRepMap.hs b/src/Data/TypeRepMap.hs index 1234932..030d3a8 100644 --- a/src/Data/TypeRepMap.hs +++ b/src/Data/TypeRepMap.hs @@ -51,6 +51,7 @@ module Data.TypeRepMap -- * Modification , insert , delete + , adjust , hoist , hoistA , hoistWithKey diff --git a/src/Data/TypeRepMap/Internal.hs b/src/Data/TypeRepMap/Internal.hs index 4dd82ef..3a2716d 100644 --- a/src/Data/TypeRepMap/Internal.hs +++ b/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')