on main help sign in
[#24] Add union and unionWith (#29)
979991bDmitrii Kovanikov
internal/Data/TypeRep/CacheMap.hs+22
@@ -21,6 +21,8 @@ module Data.TypeRep.CacheMap
21 , insert
21 , insert
22 , delete
22 , delete
23 , hoist
23 , hoist
24+ , unionWith
25+ , union
24 , lookup
26 , lookup
25 , member
27 , member
26 , size
28 , size
@@ -51,6 +53,7 @@ import GHC.Word (Word64 (..))
51import Unsafe.Coerce (unsafeCoerce)
53import Unsafe.Coerce (unsafeCoerce)
52
54
53import qualified Data.IntMap.Strict as IM
55import qualified Data.IntMap.Strict as IM
56+import qualified Data.Map.Strict as Map
54import qualified GHC.Exts as GHC (fromList, toList)
57import qualified GHC.Exts as GHC (fromList, toList)
55
58
56-- | Map-like data structure that keeps types as keys.
59-- | Map-like data structure that keeps types as keys.
@@ -123,6 +126,25 @@ Just "a"
123-}
126-}
124hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
127hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
125hoist f (TypeRepMap as bs ans) = TypeRepMap as bs $ mapArray' (toAny . f . fromAny) ans
128hoist f (TypeRepMap as bs ans) = TypeRepMap as bs $ mapArray' (toAny . f . fromAny) ans
129+{-# INLINE hoist #-}
130+
131+-- | The union with a combining function.
132+unionWith :: (forall x. f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
133+unionWith f m1 m2 = fromListPairs
134+ $ Map.toList
135+ $ Map.unionWith combine
136+ (Map.fromList $ toPairList m1)
137+ (Map.fromList $ toPairList m2)
138+ where
139+ combine :: Any -> Any -> Any
140+ combine a b = toAny $ f (fromAny a) (fromAny b)
141+{-# INLINE unionWith #-}
142+
143+-- | The (left-biased) union of two maps. It prefers the first map when
144+-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
145+union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
146+union = unionWith const
147+{-# INLINE union #-}
126
148
127{- | Returns 'True' if there exist value of given type.
149{- | Returns 'True' if there exist value of given type.
128
150
src/Data/TypeRep/Map.hs+2
@@ -12,6 +12,8 @@ module Data.TypeRep.Map
12 , insert
12 , insert
13 , delete
13 , delete
14 , hoist
14 , hoist
15+ , unionWith
16+ , union
15
17
16 -- * Query
18 -- * Query
17 , lookup
19 , lookup
test/Test/TypeRep/CacheMap.hs+8−1
@@ -5,7 +5,7 @@ import Prelude hiding (lookup)
5import Data.Functor.Identity (Identity (..))
5import Data.Functor.Identity (Identity (..))
6import Test.Tasty.Hspec (Spec, describe, it, shouldBe)
6import Test.Tasty.Hspec (Spec, describe, it, shouldBe)
7
7
8−import Data.TypeRep.CacheMap (TF (..), TypeRepMap, empty, fromList, insert, lookup, size)
8+import Data.TypeRep.CacheMap (TF (..), TypeRepMap, empty, fromList, insert, lookup, size, union)
9
9
10-- Simple test for 'lookup', 'insert' and 'size' functions.
10-- Simple test for 'lookup', 'insert' and 'size' functions.
11spec_insertLookup :: Spec
11spec_insertLookup :: Spec
@@ -26,6 +26,13 @@ spec_insertLookup = do
26 it "returns 10 when 10 different types are inserted" $
26 it "returns 10 when 10 different types are inserted" $
27 size mapOf10 `shouldBe` 10
27 size mapOf10 `shouldBe` 10
28
28
29+ describe "Union test" $ do
30+ let m = fromList [TF $ Identity 'a', TF $ Identity True] `union` fromList [TF $ Identity 'b']
31+ it "lookup works on union as expected" $ do
32+ lookup m `shouldBe` Just (Identity 'a')
33+ lookup m `shouldBe` Just (Identity True)
34+ lookup @Int m `shouldBe` Nothing
35+
29
36
30mapOf10 :: TypeRepMap Identity
37mapOf10 :: TypeRepMap Identity
31mapOf10 = insert (Identity True)
38mapOf10 = insert (Identity True)