1 files changed, 324 insertions, 0 deletions

diff --git a/lib/IntMapClass.hs b/lib/IntMapClass.hs
new file mode 100644
index 0000000..c5e6a88
--- /dev/null
+++ b/lib/IntMapClass.hs
@@ -0,0 +1,324 @@
+{-# LANGUAGE CPP,
+             FlexibleContexts,
+             MultiParamTypeClasses,
+             GeneralizedNewtypeDeriving,
+             DeriveTraversable,
+             DeriveDataTypeable #-}
+module IntMapClass where
+
+import qualified Data.IntMap.Strict as IntMap
+import Data.IntMap.Strict ( IntMap )
+import Data.Typeable ( Typeable )
+import Data.Data     ( Data )
+import Data.Foldable ( Foldable )
+import Data.Traversable ( Traversable )
+import Data.Monoid   ( Monoid )
+import Control.DeepSeq ( NFData )
+import Control.Applicative ( Applicative )
+import Data.Coerce
+
+newtype IMap k a = IMap { intmap :: IntMap a }
+ deriving
+    ( Functor
+    , Typeable
+    , Foldable
+    , Traversable
+    , Eq
+    , Data
+    , Ord
+    , Read
+    , Show
+    , Semigroup
+    , Monoid
+    , NFData
+    )
+
+adaptm_k_a_m f k a m = IMap $ adapt_k_a_m f k a m
+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
+adapt_m_k f (IMap m) k = f m (coerce k)
+
+adapt_k_m :: Coercible k Int => (Int -> IntMap a -> x) -> k -> IMap k a -> x
+adapt_k_m f k (IMap m) = f (coerce k) m
+-- adapt_k_m2 :: Coercible k Int => (Int -> IntMap a -> x) -> k -> IMap k a -> x
+-- adapt_k_m2 f k m = (adapt_k f) k (intmap m)
+
+adaptm_k_m
+  :: Coercible k Int =>
+     (Int -> IntMap a -> IntMap a) -> k -> IMap k a -> IMap k a
+adaptm_k_m f k m = IMap $ adapt_k_m f k m
+
+adapt_k :: Coercible k Int => (Int -> x) -> k -> x
+adapt_k f k = f (coerce k)
+
+adapt_m_m :: (IntMap a -> IntMap a -> x) -> IMap k a -> IMap k a -> x
+adapt_m_m f m = adapt_m (adapt_m f m)
+
+adaptm_m_m :: (IntMap a -> IntMap a -> IntMap a) -> IMap k a -> IMap k a -> IMap k a
+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 = adapt_m (IntMap.null)
+size = adapt_m (IntMap.size)
+
+member :: Coercible k Int => k -> IMap k a -> Bool
+member = adapt_k_m (IntMap.member)
+
+notMember :: Coercible k Int => k -> IMap k a -> Bool
+notMember = adapt_k_m (IntMap.notMember)
+
+lookup :: Coercible k Int => k -> IMap k a -> Maybe a
+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
+
+empty = IMap IntMap.empty
+
+singleton :: Coercible k Int => k -> a -> IMap k a
+singleton = (IMap .) . adapt_k IntMap.singleton
+
+insert :: Coercible k Int => k -> a -> IMap k a -> IMap k a
+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 f = adaptm_k_a_m (IntMap.insertWithKey $ f . coerce)
+
+insertLookupWithKey :: (Coercible Int k, Coercible k Int) =>
+    (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
+
+delete :: Coercible k Int => k -> IMap k a -> IMap k a
+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 ) =>
+     (k -> a -> a) -> k -> IMap k a -> IMap k a
+adjustWithKey f = adaptm_k_m (IntMap.adjustWithKey $ f . coerce)
+
+update
+  :: Coercible k Int => (a -> Maybe a) -> k -> IMap k a -> IMap k a
+update f = adaptm_k_m (IntMap.update f)
+
+
+updateWithKey
+  :: (Coercible k Int, 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) =>
+  (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
+
+alter :: Coercible k Int => (Maybe a -> Maybe a) -> k -> IMap k a -> IMap k a
+alter f = adaptm_k_m (IntMap.alter f)
+
+union :: IMap k a -> IMap k a -> IMap k a
+union = adaptm_m_m IntMap.union
+
+unionWith :: (a -> a -> a) -> IMap k a -> IMap k a -> IMap k a
+unionWith f = adaptm_m_m (IntMap.unionWith f)
+
+
+unionWithKey :: Coercible Int k => (k -> a -> a -> a) -> IMap k a -> IMap k a -> IMap k a
+unionWithKey f = adaptm_m_m (IntMap.unionWithKey $ f . coerce)
+
+-- unions :: Coercible [IMap k a] [IntMap a] => [IMap k a] -> IMap k a
+-- unions ms = IMap $ IntMap.unions (coerce ms)
+
+-- 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 = adaptm_m_m IntMap.difference
+
+differenceWith f = adaptm_m_m (IntMap.differenceWith f)
+
+differenceWithKey ::
+  Coercible Int k =>
+  (k -> a -> a -> Maybe a) -> IMap k a -> IMap k a -> IMap k a
+differenceWithKey f = adaptm_m_m (IntMap.differenceWithKey $ f . coerce)
+
+intersection = adaptm_m_m IntMap.intersection
+intersectionWith f = adaptm_m_m (IntMap.intersectionWith f)
+
+mergeWithKey ::
+  Coercible Int k =>
+  (k -> a -> b -> Maybe c)
+  -> (IMap k a -> IMap k c)
+  -> (IMap k b -> IMap k c)
+  -> IMap k a
+  -> IMap k b
+  -> IMap k c
+mergeWithKey f g1 g2 = adaptm_m_m (IntMap.mergeWithKey f' g1' g2')
+ where f'  = f . coerce
+       g1' = intmap . g1 . IMap
+       g2' = intmap . g2 . IMap
+       adapt_m_m f m = adapt_m (adapt_m f m)
+       adaptm_m_m f a b = IMap $ adapt_m_m f a b
+
+map :: (a -> b) -> IMap k a -> IMap k b
+map f = IMap . adapt_m (IntMap.map f)
+
+mapWithKey :: Coercible Int k => (k -> a -> b) -> IMap k a -> IMap k b
+mapWithKey f = IMap . adapt_m (IntMap.mapWithKey $ f . coerce)
+
+-- FIXME: fmap IMap ?
+traverseWithKey ::
+  (Applicative f, Coercible Int k) =>
+  (k -> a -> f b) -> IMap k a -> f (IMap k b)
+traverseWithKey f = fmap IMap . adapt_m (IntMap.traverseWithKey $ f . coerce)
+
+mapAccum :: (t -> b -> (t, a)) -> t -> IMap k b -> (t, IMap k a)
+mapAccum f a m = second IMap $ IntMap.mapAccum f a (intmap m)
+
+mapAccumWithKey :: Coercible Int k =>
+    (t -> k -> b -> (t, a)) -> t -> IMap k b -> (t, IMap k a)
+mapAccumWithKey f a m = second IMap $ IntMap.mapAccumWithKey f' a (intmap m)
+ where f' a k b = f a (coerce k) b
+
+
+mapAccumRWithKey :: Coercible Int k =>
+    (t -> k -> b -> (t, a)) -> t -> IMap k b -> (t, IMap k a)
+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) =>
+    (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) =>
+    (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) =>
+    (k1 -> k2) -> IMap k1 a -> IMap k2 a
+mapKeysMonotonic f = IMap . adapt_m (IntMap.mapKeysMonotonic (coerce . f . coerce))
+
+foldr :: (a -> x -> x) -> x -> IMap k a -> x
+foldr f b = adapt_m (IntMap.foldr f b)
+
+foldl :: (x -> a -> x) -> x -> IMap k a -> x
+foldl f a = adapt_m (IntMap.foldl f a)
+
+foldrWithKey :: Coercible Int b => (b -> a -> x -> x) -> x -> IMap k a -> x
+foldrWithKey f b = adapt_m (IntMap.foldrWithKey (f . coerce) b)
+
+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 f = adapt_m (IntMap.foldMapWithKey $ f . coerce)
+
+foldr' :: (a -> x -> x) -> x -> IMap k a -> x
+foldr' f b = adapt_m (IntMap.foldr' f b)
+
+foldl' :: (a -> x -> a) -> a -> IMap k x -> a
+foldl' f b = adapt_m (IntMap.foldl' f b)
+
+foldrWithKey' :: Coercible Int b => (b -> a -> x -> x) -> x -> IMap k a -> x
+foldrWithKey' f b = adapt_m (IntMap.foldrWithKey' (f . coerce) b)
+
+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
+
+elems :: IMap k a -> [a]
+elems = IntMap.elems . intmap
+
+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 = 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 = coerce . IntMap.toList . intmap
+
+fromList :: Coercible [(k,a)] [(Int,a)] => [(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 f = IMap . IntMap.fromListWith f . coerce
+
+fromListWithKey :: ( Coercible Int k
+                   , Coercible [(k,a)] [(Int,a)] ) =>
+    (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 (IMap m) = coerce $ IntMap.toAscList m
+
+toDescList :: Coercible [(Int,a)] [(k,a)] => IMap k a -> [(k,a)]
+toDescList (IMap m) = coerce $ IntMap.toDescList m
+
+fromAscList :: Coercible [(k,a)] [(Int,a)] => [(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 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 f = IMap . IntMap.fromAscListWithKey (f . coerce) . coerce
+
+fromDistinctAscList :: Coercible [(k,a)] [(Int,a)] => [(k, a)] -> IMap k a
+fromDistinctAscList = IMap . IntMap.fromDistinctAscList . coerce
+
+filter :: (a -> Bool) -> IMap k a -> IMap k a
+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 f m = coerce $ IntMap.partition f (intmap m)
+
+
+partitionWithKey :: ( Coercible Int k
+                    , Coercible (IntMap a,IntMap a) (IMap k a,IMap k a) )
+    => (k -> a -> Bool) -> IMap k a -> (IMap k a, IMap k a)
+partitionWithKey f m = coerce $ IntMap.partitionWithKey (f . coerce) (intmap m)
+
+mapMaybe :: (a -> Maybe b) -> IMap k a -> IMap k b
+mapMaybe f = IMap . IntMap.mapMaybe f . intmap
+
+