path: root/src/Data/Dependent/EnumMap/Strict/Internal.hs

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Dependent.EnumMap.Strict.Internal where

import Control.Exception
import Data.Bifunctor (bimap, second)
import Data.Dependent.Sum
import qualified Data.Foldable as Foldable
import qualified Data.IntMap.Strict as IM
import Data.Kind (Type)
import Data.Proxy
import Data.Some
import Data.Type.Equality
import Text.Show (showListWith)
import Unsafe.Coerce (unsafeCoerce)

import Prelude hiding (lookup, map)

data KV k v = forall a. KV !(Enum1Info k) !(v a)

newtype DEnumMap k v = DEnumMap (IM.IntMap (KV k v))

instance (Enum1 k, forall a. Show (k a), forall a. Show (v a))
      => Show (DEnumMap (k :: kind -> Type) (v :: kind -> Type)) where
  showsPrec d mp = showParen (d > 10) $
    showString "fromList " . showListWith (\(k :=> v) -> showsPrec 2 k . showString " :=> " . showsPrec 1 v) (toList mp)

-- | This class attempts to generalise 'Enum' to indexed data types: data types
-- with a GADT-like type parameter. Conversion to an 'Int' naturally loses type
-- information, and furthermore it is common to actually need some additional
-- data alongside the 'Int' to be able to reconstruct the original (in
-- 'toEnum1'). This additional data lives in 'Enum1Info'. The laws are:
--
-- [Unique IDs]
--   If @'fst' ('fromEnum1' x) == 'fst' ('fromEnum1' y)@ then @'testEquality' x y == 'Just' 'Refl' && x '==' y@
-- [Persistent IDs]
--   @'fst' ('fromEnum1' ('uncurry' 'toEnum1' ('fromEnum1' x))) == 'fst' ('fromEnum1' x)@
--
-- The "Unique IDs" law states that if the IDs of two values are equal, then
-- the values themselves must have the same type index, and furthermore be
-- equal. If @f@ does not implement 'TestEquality' or 'Eq', the law should
-- morally hold (but most of the API will be unusable).
--
-- The "Persistent IDs" law states that reconstructing a value using 'toEnum1'
-- does not change its ID.
--
-- __Note__: The methods on 'DEnumMap' attempt to check these laws using
-- 'assert' assertions (which are by default __disabled__ when optimisations
-- are on!), but full consistency cannot always be checked; if you break these
-- laws in a sufficiently clever way, the internals of 'DEnumMap' may
-- 'unsafeCoerce' unequal things and engage nasal demons, including crashes and
-- worse.
class Enum1 f where
  type Enum1Info f
  fromEnum1 :: f a -> (Int, Enum1Info f)
  toEnum1 :: Int -> Enum1Info f -> Some f

dSumToKV :: Enum1 k => DSum k v -> (Int, KV k v)
dSumToKV (k :=> v) = let (i, inf) = fromEnum1 k in (i, KV inf v)

-- | Assumes that the input was obtained via 'fromEnum1'.
kVToDSum :: Enum1 k => (Int, KV k v) -> DSum k v
kVToDSum (i, KV inf v) = case toEnum1 i inf of Some k -> k :=> coe1 v

-- * Construction

empty :: DEnumMap k v
empty = DEnumMap IM.empty

singleton :: Enum1 k => k a -> v a -> DEnumMap k v
singleton k v =
  let (i, inf) = fromEnum1 k
  in DEnumMap (IM.singleton i (KV inf v))

-- TODO: probably doesn't make much sense until we have DEnumSet? Or should we ust lists instead?
-- fromSet

-- ** From Unordered Lists

fromList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromList l = DEnumMap (IM.fromList (dSumToKV <$> l))

fromListWith :: (Enum1 k, TestEquality k)
             => (forall a. v a -> v a -> v a)
             -> [DSum k v] -> DEnumMap k v
fromListWith f (l :: [DSum k v]) =
  DEnumMap (IM.fromListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
                typeCheck2 (Proxy @k) i inf1 inf2 $
                  KV inf1 (f v1 (coe1 v2)))
             (dSumToKV <$> l))

fromListWithKey :: (Enum1 k, TestEquality k)
                => (forall a. k a -> v a -> v a -> v a)
                -> [DSum k v] -> DEnumMap k v
fromListWithKey f l =
  DEnumMap (IM.fromListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
                case toEnum1 i inf1 of
                  Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 v2)))
             (dSumToKV <$> l))

-- ** From Ascending Lists

fromAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromAscList l = DEnumMap (IM.fromAscList (dSumToKV <$> l))

fromAscListWith :: (Enum1 k, TestEquality k)
                => (forall a. v a -> v a -> v a)
                -> [DSum k v] -> DEnumMap k v
fromAscListWith f (l :: [DSum k v]) =
  DEnumMap (IM.fromAscListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
               typeCheck2 (Proxy @k) i inf1 inf2 $
                 KV inf1 (f v1 (coe1 v2)))
             (dSumToKV <$> l))

fromAscListWithKey :: (Enum1 k, TestEquality k)
                   => (forall a. k a -> v a -> v a -> v a)
                   -> [DSum k v] -> DEnumMap k v
fromAscListWithKey f l =
  DEnumMap (IM.fromAscListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
               case toEnum1 i inf1 of
                 Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 v2)))
             (dSumToKV <$> l))

fromDistinctAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromDistinctAscList l = DEnumMap (IM.fromDistinctAscList (dSumToKV <$> l))

-- * Insertion

insert :: Enum1 k => k a -> v a -> DEnumMap k v -> DEnumMap k v
insert k v (DEnumMap m) =
  let (i, inf) = fromEnum1 k
  in DEnumMap (IM.insert i (KV inf v) m)

insertWith :: (Enum1 k, TestEquality k)
           => (v a -> v a -> v a)
           -> k a -> v a -> DEnumMap k v -> DEnumMap k v
insertWith = insertWithKey . const

insertWithKey :: (Enum1 k, TestEquality k)
              => (k a -> v a -> v a -> v a)
              -> k a -> v a -> DEnumMap k v -> DEnumMap k v
insertWithKey f k v (DEnumMap m) =
  let (i, inf) = fromEnum1 k
  in DEnumMap (IM.insertWith
                (\_ (KV inf' v2) -> typeCheck1 k i inf' $ KV inf (f k v (coe1 v2)))
                i (KV inf v) m)

insertLookupWithKey :: (Enum1 k, TestEquality k)
                    => (k a -> v a -> v a -> v a)
                    -> k a -> v a -> DEnumMap k v -> (Maybe (v a), DEnumMap k v)
insertLookupWithKey f k v (DEnumMap m) =
  let (i, inf) = fromEnum1 k
      (mx, dmap) =
        IM.insertLookupWithKey
          (\_ _ (KV inf' v2) -> typeCheck1 k i inf' $ KV inf (f k v (coe1 v2)))
          i (KV inf v) m
  in ((\(KV inf' v2) -> typeCheck1 k i inf' $ coe1 v2) <$> mx, DEnumMap dmap)

-- * Deletion\/Update

delete :: Enum1 k => k a -> DEnumMap k v -> DEnumMap k v
delete k (DEnumMap m) = DEnumMap (IM.delete (fst (fromEnum1 k)) m)

adjust :: (Enum1 k, TestEquality k) => (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjust = adjustWithKey . const

adjustWithKey :: (Enum1 k, TestEquality k) => (k a -> v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjustWithKey f k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in DEnumMap (IM.adjust (\(KV inf v) -> typeCheck1 k i inf $ KV inf (f k (coe1 v))) i m)

update :: (Enum1 k, TestEquality k) => (v a -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
update = updateWithKey . const

updateWithKey :: (Enum1 k, TestEquality k) => (k a -> v a -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
updateWithKey f k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in DEnumMap (IM.update (\(KV inf v) -> typeCheck1 k i inf . KV inf <$> f k (coe1 v)) i m)

updateLookupWithKey :: (Enum1 k, TestEquality k) => (k a -> v a -> Maybe (v a)) -> k a -> DEnumMap k v -> (Maybe (v a), DEnumMap k v)
updateLookupWithKey f k (DEnumMap m) =
  let (i, _) = fromEnum1 k
      (mx, dmap) =
        IM.updateLookupWithKey
          (\_ (KV inf v) -> typeCheck1 k i inf . KV inf <$> f k (coe1 v))
          i m
  in ((\(KV inf' v2) -> typeCheck1 k i inf' $ coe1 v2) <$> mx, DEnumMap dmap)

alter :: forall k v a. (Enum1 k, TestEquality k)
      => (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
alter f k (DEnumMap m) = DEnumMap (IM.alter f' i m)
  where
    (i, inf) = fromEnum1 k

    f' :: Maybe (KV k v) -> Maybe (KV k v)
    f' Nothing = KV inf <$> f Nothing
    f' (Just (KV inf' v)) = typeCheck1 k i inf' $ KV inf <$> f (Just (coe1 v))

alterF :: forall k v a f. (Functor f, Enum1 k, TestEquality k)
       => (Maybe (v a) -> f (Maybe (v a))) -> k a -> DEnumMap k v -> f (DEnumMap k v)
alterF f k (DEnumMap m) = DEnumMap <$> IM.alterF f' i m
  where
    (i, inf) = fromEnum1 k

    f' :: Maybe (KV k v) -> f (Maybe (KV k v))
    f' Nothing = fmap (KV inf) <$> f Nothing
    f' (Just (KV inf' v)) = typeCheck1 k i inf' $ fmap (KV inf) <$> f (Just (coe1 v))

-- * Query
-- ** Lookup

lookup :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (v a)
lookup k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in (\(KV inf v) -> typeCheck1 k i inf $ coe1 v) <$> IM.lookup i m

(!?) :: (Enum1 k, TestEquality k) => DEnumMap k v -> k a -> Maybe (v a)
(!?) m k = lookup k m

findWithDefault :: (Enum1 k, TestEquality k) => v a -> k a -> DEnumMap k v -> v a
findWithDefault def k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in case IM.findWithDefault (KV undefined def) i m of
       KV inf' v -> typeCheck1 k i inf' $ coe1 v

find :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> v a
find k = findWithDefault (error ("Dependent.EnumMap.!: key " ++ show (fst $ fromEnum1 k) ++ " is not an element of the map")) k

(!) :: (Enum1 k, TestEquality k) => DEnumMap k v -> k a -> v a
(!) m k = find k m

member :: Enum1 k => k a -> DEnumMap k v -> Bool
member k (DEnumMap m) = IM.member (fst (fromEnum1 k)) m

notMember :: Enum1 k => k a -> DEnumMap k v -> Bool
notMember k m = not $ member k m

lookupLT :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
lookupLT k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in kVToDSum <$> IM.lookupLT i m

lookupGT :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
lookupGT k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in kVToDSum <$> IM.lookupGT i m

lookupLE :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
lookupLE k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in kVToDSum <$> IM.lookupLE i m

lookupGE :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (DSum k v)
lookupGE k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in kVToDSum <$> IM.lookupGE i m

-- ** Size

null :: DEnumMap k v -> Bool
null (DEnumMap m) = IM.null m

size :: DEnumMap k v -> Int
size (DEnumMap m) = IM.size m

-- * Combine

-- ** Union

union :: (Enum1 k, TestEquality k) => DEnumMap k v -> DEnumMap k v -> DEnumMap k v
union = unionWith const  -- if we're type checking, we need unionWith anyway, so might as well just delegate here already

unionWith :: (Enum1 k, TestEquality k)
          => (forall a. v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWith f (DEnumMap m1 :: DEnumMap k v) (DEnumMap m2) = DEnumMap (IM.unionWithKey f' m1 m2)
  where
    f' :: Int -> KV k v -> KV k v -> KV k v
    f' i (KV inf1 v1) (KV inf2 v2) = typeCheck2 (Proxy @k) i inf1 inf2 $ KV inf1 (f v1 (coe1 v2))

unionWithKey :: (Enum1 k, TestEquality k)
             => (forall a. k a -> v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWithKey f (DEnumMap m1 :: DEnumMap k v) (DEnumMap m2) = DEnumMap (IM.unionWithKey f' m1 m2)
  where
    f' :: Int -> KV k v -> KV k v -> KV k v
    f' i (KV inf1 v1) (KV inf2 v2) = case toEnum1 i inf1 of
      Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 v2))
        -- TODO: are the coe1 correct? is the typeCheck1 needed?

unions :: (Foldable f, Enum1 k, TestEquality k) => f (DEnumMap k v) -> DEnumMap k v
unions xs = Foldable.foldl' union empty xs

unionsWith :: (Foldable f, Enum1 k, TestEquality k) => (forall a. v a -> v a -> v a) -> f (DEnumMap k v) -> DEnumMap k v
unionsWith f xs = Foldable.foldl' (unionWith f) empty xs

-- ** Difference

-- TODO: should this be v1, v2 or both v? what about k1 and k2?
difference :: DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1
difference (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.difference m1 m2)

(\\) :: DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1
m1 \\ m2 = difference m1 m2

-- TODO: what about k1 and k2 here?
differenceWith :: forall k v1 v2. (Enum1 k, TestEquality k)
               => (forall a. v1 a -> v2 a -> Maybe (v1 a)) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
differenceWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
  where
    f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v1)
    f' i (KV inf1 v1) (KV inf2 v2) =
      typeCheck2 (Proxy @k) i inf1 inf2 . KV inf1 <$> f (coe1 v1) (coe1 v2)

-- TODO: what about k1 and k2 here?
differenceWithKey :: forall k v1 v2. (Enum1 k, TestEquality k)
                  => (forall a. k a -> v1 a -> v2 a -> Maybe (v1 a)) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v1
differenceWithKey f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
  where
    f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v1)
    f' i (KV inf1 v1) (KV inf2 v2) = case toEnum1 i inf1 of
      Some k1 -> typeCheck1 k1 i inf2 . KV inf1 <$> f k1 (coe1 v1) (coe1 v2)

-- ** Intersection

intersection :: DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1
intersection (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersection m1 m2)

intersectionWith :: forall k v1 v2 v3. (Enum1 k, TestEquality k)
                 => (forall a. v1 a -> v2 a -> v3 a) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3
intersectionWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2)
  where
    f' :: Int -> KV k v1 -> KV k v2 -> KV k v3
    f' i (KV inf1 v1) (KV inf2 v2) =
      typeCheck2 (Proxy @k) i inf1 inf2 $ KV inf1 $ f (coe1 v1) (coe1 v2)

intersectionWithKey :: forall k v1 v2 v3. (Enum1 k, TestEquality k)
                    => (forall a. k a -> v1 a -> v2 a -> v3 a) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3
intersectionWithKey f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2)
  where
    f' :: Int -> KV k v1 -> KV k v2 -> KV k v3
    f' i (KV inf1 v1) (KV inf2 v2) = case toEnum1 i inf1 of
      Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 $ f k1 (coe1 v1) (coe1 v2)

-- ** Disjoint

disjoint :: DEnumMap k v1 -> DEnumMap k v2 -> Bool
disjoint (DEnumMap m1) (DEnumMap m2) = IM.disjoint m1 m2

-- ** Compose

compose :: Enum1 k => DEnumMap k v -> DEnumMap k k -> DEnumMap k v
compose (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.compose m1 (IM.map (\(KV _ v) -> fst $ fromEnum1 v) m2))

-- ** Universal combining function

mergeWithKey :: forall k v1 v2 v3. (Enum1 k, TestEquality k)
             => (forall a. k a -> v1 a -> v2 a -> Maybe (v3 a)) -> (DEnumMap k v1 -> DEnumMap k v3) -> (DEnumMap k v2 -> DEnumMap k v3) -> DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k v3
mergeWithKey f g1 g2 (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.mergeWithKey f' g1' g2' m1 m2)
  where
    f' :: Int -> KV k v1 -> KV k v2 -> Maybe (KV k v3)
    f' i (KV inf1 v1) (KV inf2 v2) = case toEnum1 i inf1 of
      Some k1 -> typeCheck1 k1 i inf2 . KV inf1 <$> f k1 (coe1 v1) (coe1 v2)
    g1' m = let DEnumMap m' = g1 (DEnumMap m) in m'
    g2' m = let DEnumMap m' = g2 (DEnumMap m) in m'

-- * Traversal
-- ** Map

-- TODO: can the kind in the typing be avoided? Or moved to DEnumMap or Enum1 definition or somewhere?
map :: forall c (k :: c -> Type) (v1 :: c -> Type) v2. (Enum1 k, TestEquality k) => (forall a. v1 a -> v2 a) -> DEnumMap k v1 -> DEnumMap k v2
map f = mapWithKey (\_ x -> f x)

mapWithKey :: (Enum1 k, TestEquality k) => (forall a. k a -> v1 a -> v2 a) -> DEnumMap k v1 -> DEnumMap k v2
mapWithKey f (DEnumMap m) =
  DEnumMap (IM.mapWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> KV inf $ f k (coe1 v)) m)

traverseWithKey :: (Applicative t, Enum1 k, TestEquality k)
                => (forall a. k a -> v1 a -> t (v2 a)) -> DEnumMap k v1 -> t (DEnumMap k v2)
traverseWithKey f (DEnumMap m) =
  DEnumMap <$> IM.traverseWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 v)) m

traverseMaybeWithKey :: (Applicative t, Enum1 k, TestEquality k)
                     => (forall a. k a -> v1 a -> t (Maybe (v2 a))) -> DEnumMap k v1 -> t (DEnumMap k v2)
traverseMaybeWithKey f (DEnumMap m) =
  DEnumMap <$> IM.traverseMaybeWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> fmap (KV inf) <$> f k (coe1 v)) m

mapAccum :: forall c (k :: c -> Type) (v1 :: c -> Type) v2 acc. (Enum1 k, TestEquality k) => (forall a. acc -> v1 a -> (acc, v2 a)) -> acc -> DEnumMap k v1 -> (acc, DEnumMap k v2)
mapAccum f = mapAccumWithKey (\x _ y -> f x y)

mapAccumWithKey :: (Enum1 k, TestEquality k) => (forall a. acc -> k a -> v1 a -> (acc, v2 a)) -> acc -> DEnumMap k v1 -> (acc, DEnumMap k v2)
mapAccumWithKey f acc0 (DEnumMap m) =
  second DEnumMap $ IM.mapAccumWithKey (\acc i (KV inf v) -> case toEnum1 i inf of Some k -> second (KV inf) $ f acc k (coe1 v)) acc0 m

mapAccumRWithKey :: (Enum1 k, TestEquality k) => (forall a. acc -> k a -> v1 a -> (acc, v2 a)) -> acc -> DEnumMap k v1 -> (acc, DEnumMap k v2)
mapAccumRWithKey f acc0 (DEnumMap m) =
  second DEnumMap $ IM.mapAccumRWithKey (\acc i (KV inf v) -> case toEnum1 i inf of Some k -> second (KV inf) $ f acc k (coe1 v)) acc0 m

-- TODO: These are hard. Probably we can't avoid using a fold, analogously as in IntMap.
-- mapKeys
-- mapKeysWith
-- mapKeysMonotonic

-- * Folds

-- TODO: did I miss any typeCheck1 or typeCheck2 here or anywhere else? Or are any checks spurious?
foldr :: (forall a. v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldr f acc0 (DEnumMap m) = IM.foldr (\(KV _ v) acc -> f v acc) acc0 m

foldl :: (forall a. acc -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldl f acc0 (DEnumMap m) = IM.foldl (\acc (KV _ v) -> f acc v) acc0 m

foldrWith :: Enum1 k => (forall a. k a -> v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldrWith f acc0 (DEnumMap m) =
  IM.foldrWithKey (\i (KV inf v) acc -> case toEnum1 i inf of Some k -> f k (coe1 v) acc) acc0 m

foldlWith :: Enum1 k => (forall a. acc -> k a -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldlWith f acc0 (DEnumMap m) =
  IM.foldlWithKey (\acc i (KV inf v) -> case toEnum1 i inf of Some k -> f acc k (coe1 v)) acc0 m

foldMapWithKey :: (Monoid m, Enum1 k) => (forall a. k a -> v a -> m) -> DEnumMap k v -> m
foldMapWithKey f (DEnumMap m) =
  IM.foldMapWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> f k (coe1 v)) m

-- ** Strict folds

foldr' :: (forall a. v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldr' f acc0 (DEnumMap m) = IM.foldr' (\(KV _ v) acc -> f v acc) acc0 m

foldl' :: (forall a. acc -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldl' f acc0 (DEnumMap m) = IM.foldl' (\acc (KV _ v) -> f acc v) acc0 m

foldrWith' :: Enum1 k => (forall a. k a -> v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldrWith' f acc0 (DEnumMap m) =
  IM.foldrWithKey' (\i (KV inf v) acc -> case toEnum1 i inf of Some k -> f k (coe1 v) acc) acc0 m

foldlWith' :: Enum1 k => (forall a. acc -> k a -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldlWith' f acc0 (DEnumMap m) =
  IM.foldlWithKey' (\acc i (KV inf v) -> case toEnum1 i inf of Some k -> f acc k (coe1 v)) acc0 m

-- * Conversion

elems :: DEnumMap k v -> [Some v]
elems (DEnumMap m) = (\(KV _ v) -> Some v) <$> IM.elems m

keys :: Enum1 k => DEnumMap k v -> [Some k]
keys (DEnumMap m) = (\(k, KV inf _) -> toEnum1 k inf) <$> IM.assocs m

assocs :: Enum1 k => DEnumMap k v -> [DSum k v]
assocs (DEnumMap m) = kVToDSum <$> IM.assocs m

-- TODO: probably doesn't make much sense until we have DEnumSet?
-- keysSet

-- ** Lists

toList :: Enum1 k => DEnumMap k v -> [DSum k v]
toList = toAscList

-- ** Ordered lists

toAscList :: Enum1 k => DEnumMap k v -> [DSum k v]
toAscList (DEnumMap m) = kVToDSum <$> IM.toAscList m

toDescList :: Enum1 k => DEnumMap k v -> [DSum k v]
toDescList (DEnumMap m) = kVToDSum <$> IM.toDescList m

-- * Filter

-- filter

filter :: (forall a. v a -> Bool) -> DEnumMap k v -> DEnumMap k v
filter f (DEnumMap m) = DEnumMap (IM.filter (\(KV _ v) -> f v) m)

filterWithKey :: Enum1 k => (forall a. k a -> v a -> Bool) -> DEnumMap k v -> DEnumMap k v
filterWithKey f (DEnumMap m) =
  DEnumMap (IM.filterWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> f k (coe1 v)) m)

-- TODO: these use IntSet. Do we use a list instead of wait for DEnumSet?
-- restrictKeys
-- withoutKeys

partition :: (forall a. v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
partition f (DEnumMap m) =
  bimap DEnumMap DEnumMap (IM.partition (\(KV _ v) -> f v) m)

partitionWithKey :: Enum1 k => (forall a. k a -> v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
partitionWithKey f (DEnumMap m) =
  bimap DEnumMap DEnumMap (IM.partitionWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> f k (coe1 v)) m)

-- To make this more efficient, we'd need to define takeWhileAntitoneWithValue
-- for IntMap and use it here.
takeWhileAntitone :: Enum1 k => (forall a. k a -> Bool) -> DEnumMap k v -> DEnumMap k v
takeWhileAntitone f (DEnumMap m) =
  DEnumMap (IM.takeWhileAntitone (\i -> case m IM.! i of KV inf _ -> case toEnum1 i inf of Some k -> f k) m)

dropWhileAntitone :: Enum1 k => (forall a. k a -> Bool) -> DEnumMap k v -> DEnumMap k v
dropWhileAntitone f (DEnumMap m) =
  DEnumMap (IM.dropWhileAntitone (\i -> case m IM.! i of KV inf _ -> case toEnum1 i inf of Some k -> f k) m)

spanAntitone :: Enum1 k => (forall a. k a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
spanAntitone f (DEnumMap m) =
  bimap DEnumMap DEnumMap (IM.spanAntitone (\i -> case m IM.! i of KV inf _ -> case toEnum1 i inf of Some k -> f k) m)

-- TODO: can the kind in the typing be avoided? Or moved to DEnumMap or Enum1 definition or somewhere?
mapMaybe :: forall c (k :: c -> Type) (v1 :: c -> Type) v2. (Enum1 k, TestEquality k) => (forall a. v1 a -> Maybe (v2 a)) -> DEnumMap k v1 -> DEnumMap k v2
mapMaybe f = mapMaybeWithKey (\_ x -> f x)

mapMaybeWithKey :: (Enum1 k, TestEquality k) => (forall a. k a -> v1 a -> Maybe (v2 a)) -> DEnumMap k v1 -> DEnumMap k v2
mapMaybeWithKey f (DEnumMap m) =
  DEnumMap (IM.mapMaybeWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 v)) m)

mapEither :: forall c (k :: c -> Type) (v1 :: c -> Type) v2 v3. (Enum1 k, TestEquality k) => (forall a. v1 a -> Either (v2 a) (v3 a)) -> DEnumMap k v1 -> (DEnumMap k v2, DEnumMap k v3)
mapEither f = mapEitherWithKey (\_ x -> f x)

mapEitherWithKey :: (Enum1 k, TestEquality k) => (forall a. k a -> v1 a -> Either (v2 a) (v3 a)) -> DEnumMap k v1 -> (DEnumMap k v2, DEnumMap k v3)
mapEitherWithKey f (DEnumMap m) =
  bimap DEnumMap DEnumMap (IM.mapEitherWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> bimap (KV inf) (KV inf) $ f k (coe1 v)) m)

split :: Enum1 k => k a -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
split k (DEnumMap m) = bimap DEnumMap DEnumMap (IM.split (fst $ fromEnum1 k) m)

-- TODO: is this coe1 fine or can we readably check that IM doesn't cheat
-- and give us a value with a wrong type?
-- Or Maybe we should return @Some v@ instead of @v a@? But it has to match @k a@, right?
splitLookup :: Enum1 k => k a -> DEnumMap k v -> (DEnumMap k v, Maybe (v a), DEnumMap k v)
splitLookup k (DEnumMap m) =
  let (m1, mkv, m2) = IM.splitLookup (fst $ fromEnum1 k) m
  in (DEnumMap m1, (\(KV _ v) -> coe1 v) <$> mkv, DEnumMap m2)

splitRoot :: DEnumMap k v -> [DEnumMap k v]
splitRoot (DEnumMap m) = DEnumMap <$> IM.splitRoot m

-- * Submap

-- isSubmapOf, isSubmapOfBy
-- isProperSubmapOf, isProperSubmapOfBy

-- * Min\/Max

lookupMin :: Enum1 k => DEnumMap k v -> Maybe (DSum k v)
lookupMin (DEnumMap m) = kVToDSum <$> IM.lookupMin m

lookupMax :: Enum1 k => DEnumMap k v -> Maybe (DSum k v)
lookupMax (DEnumMap m) = kVToDSum <$> IM.lookupMax m

findMin :: Enum1 k => DEnumMap k v -> DSum k v
findMin (DEnumMap m) = kVToDSum $ IM.findMin m

findMax :: Enum1 k => DEnumMap k v -> DSum k v
findMax (DEnumMap m) = kVToDSum $ IM.findMax m

deleteMin :: DEnumMap k v -> DEnumMap k v
deleteMin (DEnumMap m) = DEnumMap $ IM.deleteMin m

deleteMax :: DEnumMap k v -> DEnumMap k v
deleteMax (DEnumMap m) = DEnumMap $ IM.deleteMax m

deleteFindMin :: Enum1 k => DEnumMap k v -> (DSum k v, DEnumMap k v)
deleteFindMin (DEnumMap m) = bimap kVToDSum DEnumMap $ IM.deleteFindMin m

deleteFindMax :: Enum1 k => DEnumMap k v -> (DSum k v, DEnumMap k v)
deleteFindMax (DEnumMap m) = bimap kVToDSum DEnumMap $ IM.deleteFindMax m

updateMin :: forall c (k :: c -> Type) (v :: c -> Type). Enum1 k => (forall a. v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMin f = updateMinWithKey (\_ x -> f x)

updateMinWithKey :: Enum1 k => (forall a. k a -> v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMinWithKey f (DEnumMap m) =
  DEnumMap (IM.updateMinWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 v)) m)

updateMax :: forall c (k :: c -> Type) (v :: c -> Type). Enum1 k => (forall a. v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMax f = updateMaxWithKey (\_ x -> f x)

updateMaxWithKey :: Enum1 k => (forall a. k a -> v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMaxWithKey f (DEnumMap m) =
  DEnumMap (IM.updateMaxWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> KV inf <$> f k (coe1 v)) m)

minView :: DEnumMap k v -> Maybe (v a, DEnumMap k v)
minView (DEnumMap m) = bimap (\(KV _ v) -> coe1 v) DEnumMap <$> IM.minView m

maxView :: DEnumMap k v -> Maybe (v a, DEnumMap k v)
maxView (DEnumMap m) = bimap (\(KV _ v) -> coe1 v) DEnumMap <$> IM.maxView m

minViewWithKey :: Enum1 k => DEnumMap k v -> Maybe (DSum k v, DEnumMap k v)
minViewWithKey (DEnumMap m) = bimap kVToDSum DEnumMap <$> IM.minViewWithKey m

maxViewWithKey :: Enum1 k => DEnumMap k v -> Maybe (DSum k v, DEnumMap k v)
maxViewWithKey (DEnumMap m) = bimap kVToDSum DEnumMap <$> IM.maxViewWithKey m


-- * Helpers

coe1 :: v a -> v b
coe1 = unsafeCoerce

typeCheck1 :: (Enum1 k, TestEquality k)
           => k a -> Int -> Enum1Info k -> r -> r
typeCheck1 k1 i inf2 x =
  assert (case toEnum1 i inf2 of { Some k2 ->
          case testEquality k1 k2 of
            Just Refl -> True
            Nothing -> False })
         x

typeCheck2 :: forall k proxy r. (Enum1 k, TestEquality k)
           => proxy k -> Int -> Enum1Info k -> Enum1Info k -> r -> r
typeCheck2 _ i inf1 inf2 x =
  assert (case toEnum1 @k i inf1 of { Some k1 ->
          case toEnum1 i inf2 of { Some k2 ->
          case testEquality k1 k2 of
            Just Refl -> True
            Nothing -> False }})
         x