Updates to IntMapClass.

1 files changed, 53 insertions, 45 deletions

diff --git a/lib/IntMapClass.hs b/lib/IntMapClass.hs
index 3d08e46..b44e05c 100644
--- a/lib/IntMapClass.hs
+++ b/lib/IntMapClass.hs
@@ -5,6 +5,7 @@
              DeriveDataTypeable #-}
 module IntMapClass where
 
+import Control.Arrow (second)
 import qualified Data.IntMap.Strict as IntMap
 import Data.IntMap.Strict ( IntMap )
 import Data.Typeable ( Typeable )
@@ -32,7 +33,13 @@ newtype IMap k a = IMap { intmap :: IntMap a }
     , NFData
     )
 
+adapt_k_a_m :: Coercible k1 Int =>
+                     (Int -> t -> IntMap a -> x) -> k1 -> t -> IMap k2 a -> x
 adaptm_k_a_m f k a m = IMap $ adapt_k_a_m f k a m
+
+adaptm_k_a_m :: Coercible k1 Int =>
+                      (Int -> t -> IntMap a1 -> IntMap a2)
+                      -> k1 -> t -> IMap k2 a1 -> IMap k3 a2
 adapt_k_a_m f k a m = adapt_m (adapt_k f k a) m
 
 adapt_m_k :: Coercible k Int => (IntMap a -> Int -> x) -> IMap k a -> k -> x
@@ -60,17 +67,16 @@ adaptm_m_m f a b = IMap $ adapt_m_m f a b
 adapt_m :: (IntMap a -> x) -> IMap k a -> x
 adapt_m f (IMap m) = f m
 
-first f (x,y) = (f x,y)
-second f (x,y) = (x,f y)
-
-
 (!) :: Coercible k Int => IMap k a -> k -> a
 (!) = adapt_m_k (IntMap.!)
 
 (\\) :: IMap k a -> IMap k a -> IMap k a
 (\\) a b = IMap $ adapt_m_m (IntMap.\\) a b
 
+null :: IMap k a -> Bool
 null = adapt_m (IntMap.null)
+
+size :: IMap k a -> Int
 size = adapt_m (IntMap.size)
 
 member :: Coercible k Int => k -> IMap k a -> Bool
@@ -85,16 +91,19 @@ lookup = adapt_k_m (IntMap.lookup)
 findWithDefault :: Coercible k Int => x -> k -> IMap k x -> x
 findWithDefault a = adapt_k_m (IntMap.findWithDefault a)
 
--- FIXME: fmap (first coerce) probably incurs cost
-lookupLT :: ( Coercible Int k, Coercible k Int ) => k -> IMap k a -> Maybe (k, a)
-lookupLT k m = fmap (first coerce) $ adapt_k_m (IntMap.lookupLT) k m
-lookupGT :: ( Coercible Int k, Coercible k Int ) => k -> IMap k a -> Maybe (k, a)
-lookupGT k m = fmap (first coerce) $ adapt_k_m (IntMap.lookupGT) k m
-lookupLE :: ( Coercible Int k, Coercible k Int ) => k -> IMap k a -> Maybe (k, a)
-lookupLE k m = fmap (first coerce) $ adapt_k_m (IntMap.lookupLE) k m
-lookupGE :: ( Coercible Int k, Coercible k Int ) => k -> IMap k a -> Maybe (k, a)
-lookupGE k m = fmap (first coerce) $ adapt_k_m (IntMap.lookupGE) k m
+lookupLT :: Coercible Int k => k -> IMap k a -> Maybe (k, a)
+lookupLT k (IMap m) = coerce $ IntMap.lookupLT (coerce k) m
 
+lookupGT :: Coercible Int k => k -> IMap k a -> Maybe (k, a)
+lookupGT k (IMap m) = coerce $ IntMap.lookupGT (coerce k) m
+
+lookupLE :: Coercible Int k => k -> IMap k a -> Maybe (k, a)
+lookupLE k (IMap m) = coerce $ IntMap.lookupLE (coerce k) m
+
+lookupGE :: Coercible Int k => k -> IMap k a -> Maybe (k, a)
+lookupGE k (IMap m) = coerce $ IntMap.lookupGE (coerce k) m
+
+empty :: IMap k a
 empty = IMap IntMap.empty
 
 singleton :: Coercible k Int => k -> a -> IMap k a
@@ -106,10 +115,10 @@ insert = adaptm_k_a_m IntMap.insert
 insertWith :: Coercible k Int => (a -> a -> a) -> k -> a -> IMap k a -> IMap k a
 insertWith f = adaptm_k_a_m (IntMap.insertWith f)
 
-insertWithKey :: (Coercible Int k, Coercible k Int) => (k -> a -> a -> a) -> k -> a -> IMap k a -> IMap k a
+insertWithKey :: Coercible Int k => (k -> a -> a -> a) -> k -> a -> IMap k a -> IMap k a
 insertWithKey f = adaptm_k_a_m (IntMap.insertWithKey $ f . coerce)
 
-insertLookupWithKey :: (Coercible Int k, Coercible k Int) =>
+insertLookupWithKey :: Coercible Int k =>
     (k -> a -> a -> a) -> k -> a -> IMap k a -> (Maybe a, IMap k a)
 insertLookupWithKey f k a m = second IMap $ adapt_k_a_m (IntMap.insertLookupWithKey $ f . coerce) k a m
 
@@ -119,7 +128,7 @@ delete = adaptm_k_m IntMap.delete
 adjust :: Coercible k Int => (a -> a) -> k -> IMap k a -> IMap k a
 adjust f = adaptm_k_m (IntMap.adjust f)
 
-adjustWithKey :: ( Coercible Int k, Coercible k Int ) =>
+adjustWithKey :: Coercible Int k =>
      (k -> a -> a) -> k -> IMap k a -> IMap k a
 adjustWithKey f = adaptm_k_m (IntMap.adjustWithKey $ f . coerce)
 
@@ -128,13 +137,11 @@ update
 update f = adaptm_k_m (IntMap.update f)
 
 
-updateWithKey
-  :: (Coercible k Int, Coercible Int k) =>
+updateWithKey :: Coercible Int k =>
      (k -> a -> Maybe a) -> k -> IMap k a -> IMap k a
 updateWithKey f = adaptm_k_m (IntMap.updateWithKey $ f . coerce)
 
-updateLookupWithKey ::
-  (Coercible k Int, Coercible Int k) =>
+updateLookupWithKey :: Coercible k Int =>
   (k -> a -> Maybe a) -> k -> IMap k a -> (Maybe a, IMap k a)
 updateLookupWithKey f k m =
     second IMap $ adapt_k_m (IntMap.updateLookupWithKey $ f . coerce) k m
@@ -158,8 +165,11 @@ unionWithKey f = adaptm_m_m (IntMap.unionWithKey $ f . coerce)
 -- unionsWith :: Coercible [IMap k a] [IntMap a] => (a->a->a) -> [IMap k a] -> IMap k a
 -- unionsWith f ms = IMap $ IntMap.unionsWith f (coerce ms)
 
+difference :: IMap k b -> IMap k b -> IMap k b
 difference = adaptm_m_m IntMap.difference
 
+differenceWith :: (b -> b -> Maybe b)
+                        -> IMap k b -> IMap k b -> IMap k b
 differenceWith f = adaptm_m_m (IntMap.differenceWith f)
 
 differenceWithKey ::
@@ -167,7 +177,10 @@ differenceWithKey ::
   (k -> a -> a -> Maybe a) -> IMap k a -> IMap k a -> IMap k a
 differenceWithKey f = adaptm_m_m (IntMap.differenceWithKey $ f . coerce)
 
+intersection :: IMap k b -> IMap k b -> IMap k b
 intersection = adaptm_m_m IntMap.intersection
+
+intersectionWith :: (a -> a -> a) -> IMap k a -> IMap k a -> IMap k a
 intersectionWith f = adaptm_m_m (IntMap.intersectionWith f)
 
 mergeWithKey ::
@@ -211,15 +224,15 @@ mapAccumRWithKey :: Coercible Int k =>
 mapAccumRWithKey f a m = second IMap $ IntMap.mapAccumRWithKey f' a (intmap m)
  where f' a k b = f a (coerce k) b
 
-mapKeys :: (Coercible Int k1, Coercible k2 Int) =>
+mapKeys :: (Coercible Int k1, Coercible Int k2) =>
     (k1 -> k2) -> IMap k1 a -> IMap k2 a
 mapKeys f = IMap . adapt_m (IntMap.mapKeys (coerce . f . coerce))
 
-mapKeysWith :: (Coercible Int k1, Coercible k2 Int) =>
+mapKeysWith :: (Coercible Int k1, Coercible Int k2) =>
     (a->a->a) -> (k1 -> k2) -> IMap k1 a -> IMap k2 a
 mapKeysWith c f = IMap . adapt_m (IntMap.mapKeysWith c (coerce . f . coerce))
 
-mapKeysMonotonic :: (Coercible Int k1, Coercible k2 Int) =>
+mapKeysMonotonic :: (Coercible Int k1, Coercible Int k2) =>
     (k1 -> k2) -> IMap k1 a -> IMap k2 a
 mapKeysMonotonic f = IMap . adapt_m (IntMap.mapKeysMonotonic (coerce . f . coerce))
 
@@ -236,7 +249,7 @@ foldlWithKey ::
   Coercible Int k => (x -> k -> a -> x) -> x -> IMap k a -> x
 foldlWithKey f a = adapt_m (IntMap.foldlWithKey f' a) where f' a = f a . coerce
 
-foldMapWithKey :: (Monoid m, Coercible k Int) => (k -> a -> m) -> IMap k a -> m
+foldMapWithKey :: (Monoid m, Coercible Int k) => (k -> a -> m) -> IMap k a -> m
 foldMapWithKey f = adapt_m (IntMap.foldMapWithKey $ f . coerce)
 
 foldr' :: (a -> x -> x) -> x -> IMap k a -> x
@@ -255,50 +268,47 @@ foldlWithKey' f a = adapt_m (IntMap.foldlWithKey' f' a) where f' a = f a . coerc
 elems :: IMap k a -> [a]
 elems = IntMap.elems . intmap
 
-keys :: Coercible [Int] [k] => IMap k a -> [k]
+keys :: Coercible Int k => IMap k a -> [k]
 keys = coerce . IntMap.keys . intmap
 
-assocs :: Coercible [(Int,a)] [(k,a)] => IMap k a -> [(k, a)]
+assocs :: Coercible Int k => IMap k a -> [(k, a)]
 assocs = coerce . IntMap.assocs . intmap
 
 -- Not implementing... (doing it right requires wrapping IntSet)
 -- keysSet :: IntMap a -> IntSet
 -- fromSet :: (Key -> a) -> IntSet -> IntMap a
 
-toList :: Coercible [(Int,a)] [(k,a)] => IMap k a -> [(k, a)]
+toList :: Coercible Int k => IMap k a -> [(k, a)]
 toList = coerce . IntMap.toList . intmap
 
-fromList :: Coercible [(k,a)] [(Int,a)] => [(k, a)] -> IMap k a
+fromList :: Coercible Int k => [(k, a)] -> IMap k a
 fromList = IMap . IntMap.fromList . coerce
 
-fromListWith :: Coercible [(k,a)] [(Int,a)] => (a -> a -> a) -> [(k, a)] -> IMap k a
+fromListWith :: Coercible Int k => (a -> a -> a) -> [(k, a)] -> IMap k a
 fromListWith f = IMap . IntMap.fromListWith f . coerce
 
-fromListWithKey :: ( Coercible Int k
-                   , Coercible [(k,a)] [(Int,a)] ) =>
+fromListWithKey :: Coercible Int k =>
     (k -> a -> a -> a) -> [(k, a)] -> IMap k a
 fromListWithKey f = IMap . IntMap.fromListWithKey (f . coerce) . coerce
 
-toAscList :: Coercible [(Int,a)] [(k,a)] => IMap k a -> [(k,a)]
+toAscList :: Coercible Int k => IMap k a -> [(k,a)]
 toAscList (IMap m) = coerce $ IntMap.toAscList m
 
-toDescList :: Coercible [(Int,a)] [(k,a)] => IMap k a -> [(k,a)]
+toDescList :: Coercible Int k => IMap k a -> [(k,a)]
 toDescList (IMap m) = coerce $ IntMap.toDescList m
 
-fromAscList :: Coercible [(k,a)] [(Int,a)] => [(k, a)] -> IMap k a
+fromAscList :: Coercible Int k => [(k, a)] -> IMap k a
 fromAscList = IMap . IntMap.fromAscList . coerce
 
-fromAscListWith ::
-  Coercible [(k,a)] [(Int, a)] =>
-  (a -> a -> a) -> [(k,a)] -> IMap k a
+fromAscListWith :: Coercible Int k
+  => (a -> a -> a) -> [(k,a)] -> IMap k a
 fromAscListWith f = IMap . IntMap.fromAscListWith f . coerce
 
-fromAscListWithKey ::
-  (Coercible Int k, Coercible [(k,a)] [(Int, a)]) =>
-  (k -> a -> a -> a) -> [(k,a)] -> IMap k a
+fromAscListWithKey :: Coercible Int k
+  => (k -> a -> a -> a) -> [(k,a)] -> IMap k a
 fromAscListWithKey f = IMap . IntMap.fromAscListWithKey (f . coerce) . coerce
 
-fromDistinctAscList :: Coercible [(k,a)] [(Int,a)] => [(k, a)] -> IMap k a
+fromDistinctAscList :: Coercible Int k => [(k, a)] -> IMap k a
 fromDistinctAscList = IMap . IntMap.fromDistinctAscList . coerce
 
 filter :: (a -> Bool) -> IMap k a -> IMap k a
@@ -307,13 +317,11 @@ filter f = IMap . adapt_m (IntMap.filter f)
 filterWithKey :: Coercible Int k => (k -> a -> Bool) -> IMap k a -> IMap k a
 filterWithKey f = IMap . adapt_m (IntMap.filterWithKey $ f . coerce)
 
-partition :: Coercible (IntMap a,IntMap a) (IMap k a,IMap k a)
-    => (a -> Bool) -> IMap k a -> (IMap k a, IMap k a)
+partition :: (a -> Bool) -> IMap k a -> (IMap k a, IMap k a)
 partition f m = coerce $ IntMap.partition f (intmap m)
 
 
-partitionWithKey :: ( Coercible Int k
-                    , Coercible (IntMap a,IntMap a) (IMap k a,IMap k a) )
+partitionWithKey :: Coercible Int k
     => (k -> a -> Bool) -> IMap k a -> (IMap k a, IMap k a)
 partitionWithKey f m = coerce $ IntMap.partitionWithKey (f . coerce) (intmap m)