Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse files

Allow 'lookup' to return both values and subtrees depending on the key.

  • Loading branch information...
commit 4a837acac39b5d7976cde8f09c9eec4aa6762938 1 parent 6119a41
@bovinespirit authored
View
55 Data/EnumMapMap/Base.hs
@@ -14,8 +14,8 @@
--
-- Based on Data.IntMap.Base.
--
--- This defines the EnumMapMap (k :& t) v instance, and the Key data types. The
--- terminating key type is K, and the EnumMapMap (K k) v instances are defined
+-- This defines the @'EnumMapMap' (k ':&' t) v@ instance, and the Key data types. The
+-- terminating key type is K, and the @'EnumMapMap' (K k) v@ instances are defined
-- in EnumMapMap.Lazy and EnumMapMap.Strict.
-----------------------------------------------------------------------------
@@ -26,6 +26,7 @@ module Data.EnumMapMap.Base(
-- * Split/Join Keys
IsSplit(..),
Plus,
+ CanSplit(..),
-- * Internal
-- ** IsEMM
EMM(..),
@@ -150,6 +151,38 @@ instance (IsSplit t n, Enum k) => IsSplit (k :& t) (N n) where
type family Plus k1 k2 :: *
type instance Plus (k1 :& t) k2 = k1 :& Plus t k2
+class CanSplit k1 k2 v where
+ type Result k1 k2 v :: *
+ -- | Lookup up the value at a key in the 'EnumMapMap'.
+ --
+ -- > emm = fromList [(3 :& K 1, "a")]
+ -- > lookup (3 :& K 1) emm == Just "a"
+ -- > lookup (2 :& K 1) emm == Nothing
+ --
+ -- If the given key has less dimensions then the 'EnumMapMap' then a submap
+ -- is returned.
+ --
+ -- > emm2 = fromList [(3 :& 2 :& K 1, "a"), (3 :& 2 :& K 4, "a")]
+ -- > lookup (3 :& K 2) emm2 == Just $ fromList [(K 1, "a"), (K 4, "a")]
+ --
+ lookup :: (IsEmm k1, IsEmm k2) =>
+ k1 -> EnumMapMap k2 v -> Maybe (Result k1 k2 v)
+
+instance (Enum k, IsEmm t1, IsEmm t2, CanSplit t1 t2 v) =>
+ CanSplit (k :& t1) (k :& t2) v where
+ type Result (k :& t1) (k :& t2) v = Result t1 t2 v
+ lookup (key' :& nxt) (KCC emm) = key `seq` go emm
+ where
+ go (Bin _ m l r)
+ | zero key m = go l
+ | otherwise = go r
+ go (Tip kx x)
+ = case kx == key of
+ True -> lookup nxt x
+ False -> Nothing
+ go Nil = Nothing
+ key = fromEnum key'
+
class HasSKey k where
type Skey k :: *
-- | Convert a key terminated with 'K' into one terminated with 'S'.
@@ -229,13 +262,6 @@ class (Eq k) => IsEmm k where
--
-- > singleton (5 :& K 3) "a" == fromList [(5 :& K 3, "a")]
singleton :: k -> v -> EnumMapMap k v
- -- | Lookup up the value at a key in the 'EnumMapMap'.
- --
- -- > emm = fromList [(3 :& K 1, "a")]
- -- > lookup (3 :& K 1) emm == Just "a"
- -- > lookup (2 :& K 1) emm == Nothing
- --
- lookup :: k -> EnumMapMap k v -> Maybe v
-- | Insert a new key\/value pair into the 'EnumMapMap'.
insert :: k -> v -> EnumMapMap k v -> EnumMapMap k v
-- | Insert with a combining function.
@@ -372,17 +398,6 @@ instance (Eq k, Enum k, IsEmm t, HasSKey t) => IsEmm (k :& t) where
singleton (key :& nxt) = KCC . Tip (fromEnum key) . singleton nxt
- lookup (key :& nxt) (KCC emm) = go emm
- where
- go (Bin _ m l r)
- | zero (fromEnum key) m = go l
- | otherwise = go r
- go (Tip kx x)
- = case kx == fromEnum key of
- True -> lookup nxt x
- False -> Nothing
- go Nil = Nothing
-
insert (key :& nxt) val (KCC emm)
= KCC $ insertWith_ (insert nxt val) key (singleton nxt val) emm
View
41 Data/EnumMapMap/Lazy.hs
@@ -145,18 +145,6 @@ instance (Enum k, Eq k) => IsEmm (K k) where
singleton !(K key) = KEC . Tip (fromEnum key)
- lookup !(K key') (KEC emm) = go emm
- where
- go (Bin _ m l r)
- | zero key m = go l
- | otherwise = go r
- go (Tip kx x)
- = case kx == key of
- True -> Just x
- False -> Nothing
- go Nil = Nothing
- key = fromEnum key'
-
insert !(K key') val (KEC emm) = KEC $ go emm
where
go t = case t of
@@ -295,3 +283,32 @@ instance IsSplit (k :& t) Z where
type Head (k :& t) Z = K k
type Tail (k :& t) Z = t
splitKey Z (KCC emm) = KEC $ emm
+
+instance (Enum k1, k1 ~ k2) => CanSplit (K k1) (k2 :& t2) v where
+ type Result (K k1) (k2 :& t2) v = EnumMapMap t2 v
+ lookup (K key') (KCC emm) = key `seq` go emm
+ where
+ go (Bin _ m l r)
+ | zero key m = go l
+ | otherwise = go r
+ go (Tip kx x)
+ = case kx == key of
+ True -> Just x
+ False -> Nothing
+ go Nil = Nothing
+ key = fromEnum key'
+
+instance (Enum k) => CanSplit (K k) (K k) v where
+ type Result (K k) (K k) v = v
+ lookup (K key') (KEC emm) = key `seq` go emm
+ where
+ go (Bin _ m l r)
+ | zero key m = go l
+ | otherwise = go r
+ go (Tip kx x)
+ = case kx == key of
+ True -> Just x
+ False -> Nothing
+ go Nil = Nothing
+ key = fromEnum key'
+
View
41 Data/EnumMapMap/Strict.hs
@@ -146,18 +146,6 @@ instance (Enum k, Eq k) => IsEmm (K k) where
singleton !(K key) !val = KEC $ Tip (fromEnum key) val
- lookup !(K key') (KEC emm) = go emm
- where
- go (Bin _ m l r)
- | zero key m = go l
- | otherwise = go r
- go (Tip kx x)
- = case kx == key of
- True -> Just x
- False -> Nothing
- go Nil = Nothing
- key = fromEnum key'
-
insert !(K key') !val (KEC emm) = KEC $ go emm
where
go t = case t of
@@ -296,3 +284,32 @@ instance IsSplit (k :& t) Z where
type Head (k :& t) Z = K k
type Tail (k :& t) Z = t
splitKey Z (KCC emm) = KEC $ emm
+
+instance (Enum k1, k1 ~ k2) => CanSplit (K k1) (k2 :& t2) v where
+ type Result (K k1) (k2 :& t2) v = EnumMapMap t2 v
+ lookup (K key') (KCC emm) = key `seq` go emm
+ where
+ go (Bin _ m l r)
+ | zero key m = go l
+ | otherwise = go r
+ go (Tip kx x)
+ = case kx == key of
+ True -> Just x
+ False -> Nothing
+ go Nil = Nothing
+ key = fromEnum key'
+
+instance (Enum k) => CanSplit (K k) (K k) v where
+ type Result (K k) (K k) v = v
+ lookup (K key') (KEC emm) = key `seq` go emm
+ where
+ go (Bin _ m l r)
+ | zero key m = go l
+ | otherwise = go r
+ go (Tip kx x)
+ = case kx == key of
+ True -> Just x
+ False -> Nothing
+ go Nil = Nothing
+ key = fromEnum key'
+
View
1  Data/EnumMapSet.hs
@@ -31,6 +31,7 @@ module Data.EnumMapSet (
EMS.null,
size,
member,
+ EMS.lookup,
-- * Construction
empty,
singleton,
View
31 Data/EnumMapSet/Base.hs
@@ -1,5 +1,6 @@
-{-# LANGUAGE BangPatterns, CPP, GeneralizedNewtypeDeriving, MagicHash,
- TypeFamilies #-}
+{-# LANGUAGE BangPatterns, CPP, FlexibleContexts, FlexibleInstances,
+ GeneralizedNewtypeDeriving, MagicHash, MultiParamTypeClasses, TypeFamilies,
+ TypeOperators #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-----------------------------------------------------------------------------
@@ -23,6 +24,7 @@ module Data.EnumMapSet.Base (
null,
size,
member,
+ lookup,
-- * Construction
empty,
singleton,
@@ -268,7 +270,6 @@ instance (Enum k, Eq k) => IsEmm (S k) where
insertWith = undefined
insertWithKey = undefined
- lookup = undefined
alter = undefined
foldr = undefined
map = undefined
@@ -301,6 +302,16 @@ size = EMM.size
member ::(IsEmm k) => k -> EnumMapSet k -> Bool
member = EMM.member
+-- | Lookup a subtree in an 'EnumMapSet'.
+--
+-- > ems = fromList [1 :& 2 :& K 3, 1 :& 2 :& K 4]
+-- > lookup (1 :& K 2) ems == fromList [K 3, K 4]
+-- > lookup (1 :& 2 :& K 3) -- ERROR: Use 'member' to check for a key.
+--
+lookup :: (EMM.CanSplit k1 k2 (), IsEmm k1, IsEmm k2) =>
+ k1 -> EnumMapSet k2 -> Maybe (EMM.Result k1 k2 ())
+lookup = EMM.lookup
+
empty :: (IsEmm k) => EnumMapSet k
empty = EMM.empty
@@ -359,6 +370,20 @@ instance EMM.HasSKey (S k) where
toS (S _) = undefined
toK (S _) = undefined
+instance (Enum k1, k1 ~ k2) => EMM.CanSplit (S k1) (k2 :& t2) () where
+ type Result (S k1) (k2 :& t2) () = EnumMapSet t2
+ lookup (S key') (EMM.KCC emm) = key `seq` go emm
+ where
+ go (EMM.Bin _ m l r)
+ | zero key m = go l
+ | otherwise = go r
+ go (EMM.Tip kx x)
+ = case kx == key of
+ True -> Just x
+ False -> Nothing
+ go EMM.Nil = Nothing
+ key = fromEnum key'
+
{---------------------------------------------------------------------
Helper functions
---------------------------------------------------------------------}
View
21 test/UnitEnumMapMap.hs
@@ -96,6 +96,27 @@ main =
it "is the inverse of toList on 2 levels" $
(EMM.fromList $ EMM.toList l2odds) @?= l2odds
+ describe "lookup" $ do
+ let emm3 :: TestEmm3
+ emm3 = EMM.fromList [(ID3 1 :& ID2 2 :& (K $ ID1 3), 4)]
+ key3 = ID3 1 :& ID2 2 :& (K $ ID1 3)
+ describe "looks up a subtree" $ do
+ let emm2 :: EnumMapMap (Int :& K Int) Int
+ emm2 = EMM.fromList [(1 :& K 2, 5)]
+ key1 :: K ID3
+ key1 = K $ ID3 1
+ key2 :: ID3 :& K ID2
+ key2 = ID3 1 :& (K $ ID2 2)
+ it "First level of level 2" $
+ (EMM.lookup (K 1) emm2) @?= (Just $ EMM.fromList [(K 2, 5)])
+ it "1 level of level 3" $
+ (EMM.lookup key1 emm3) @?= (Just $
+ EMM.fromList [(ID2 2 :& (K $ ID1 3), 4)])
+ it "2 levels of level 3" $
+ (EMM.lookup key2 emm3) @?= (Just $ EMM.fromList [(K $ ID1 3, 4)])
+ it "looks up a value" $
+ (EMM.lookup key3 emm3) @?= Just 4
+
describe "insert" $ do
describe "Level 1" $ do
it "creates a value in an empty EMM" $
Please sign in to comment.
Something went wrong with that request. Please try again.