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

import Control.Exception
import Data.Bifunctor (bimap, second)
import Data.Coerce
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)

type KV :: forall kind. (kind -> Type) -> (kind -> Type) -> Type
data KV k v = forall a. KV !(Enum1Info k) !(v a)

-- Invariant: the key-value pairs in a DEnumMap are type-consistent. That is to
-- say: they have the same type-index. Any other type equalities, like between
-- the key argument to 'lookup' and the key-value pairs in the map argument to
-- 'lookup', may /not/ hold, and should be type-checked as much as we're able.
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 v) 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: Wait for DEnumSet.
-- 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, !m') =
        IM.insertLookupWithKey
          (\_ _ (KV inf' v2) -> typeCheck1 k i inf' $ KV inf (f k v (coe1 v2)))
          i (KV inf v) m
     -- Note: type checking unnecessary here, because by the BangPatterns,
     -- evaluating mx evaluates dmap, and the IntMap is strict, so the lambda
     -- will have run and typechecked the old value already.
     -- Second note: the BangPatterns don't do anything operationally because
     -- with the current implementation of IM.insertLookupWithKey, the pair
     -- components are already strict.
  in ((\(KV _ v2) -> coe1 v2) <$> mx, DEnumMap m')

-- * 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, !m') =
        IM.updateLookupWithKey
          (\_ (KV inf v) -> typeCheck1 k i inf $ KV inf <$> f k (coe1 v))
          i m
     -- Note: type checking unnecessary here for the same reason as insertLookupWithKey
  in ((\(KV _ v2) -> coe1 v2) <$> mx, DEnumMap m')

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 ("Data.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, lookupGT, lookupLE, lookupGE
  :: (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 k (DEnumMap m) = let (i, _) = fromEnum1 k in kVToDSum <$> IM.lookupGT i m
lookupLE k (DEnumMap m) = let (i, _) = fromEnum1 k in kVToDSum <$> IM.lookupLE i m
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))

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

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

-- ** Difference

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

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

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)

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)

-- | Because the underlying maps are keyed on integers, it is possible to
-- subtract a map from another even if the key types differ. This function
-- assumes that the @Int@ identifiers of @k1@ and @k2@ are compatible, i.e.
-- that "2" in @k1@ somehow means the same thing as "2" in @k2@.
--
-- Because the key types are different, there is no guarantee whatsoever (even
-- not by 'Enum1' laws) that equal key IDs in @k1@ and @k2@ actually have the
-- same type index (@a@). Hence, the combining function gets key-value pairs
-- with potentially distinct type indices.
differenceWithKey' :: forall k1 k2 v1 v2. (Enum1 k1, Enum1 k2)
                   => (forall a b. k1 a -> v1 a -> k2 b -> v2 b -> Maybe (v1 a))
                   -> DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v1
differenceWithKey' f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.differenceWithKey f' m1 m2)
  where
    f' :: Int -> KV k1 v1 -> KV k2 v2 -> Maybe (KV k1 v1)
    f' i (KV inf1 v1) (KV inf2 v2) = case (toEnum1 i inf1, toEnum1 i inf2) of
      (Some k1, Some k2) -> KV inf1 <$> f k1 (coe1 v1) k2 (coe1 v2)

-- ** Intersection

intersection :: DEnumMap k v1 -> DEnumMap k v2 -> DEnumMap k 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)

-- | Generalises 'intersectionWithKey' in the same way as 'differenceWithKey''
-- generalises 'differenceWithKey'.
intersectionWithKey' :: forall k1 k2 v1 v2 v3. (Enum1 k1, Enum1 k2)
                     => (forall a b. k1 a -> v1 a -> k2 b -> v2 b -> v3 a)
                     -> DEnumMap k1 v1 -> DEnumMap k2 v2 -> DEnumMap k1 v3
intersectionWithKey' f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.intersectionWithKey f' m1 m2)
  where
    f' :: Int -> KV k1 v1 -> KV k2 v2 -> KV k1 v3
    f' i (KV inf1 v1) (KV inf2 v2) = case (toEnum1 i inf1, toEnum1 i inf2) of
      (Some k1, Some k2) -> KV inf1 $ f k1 (coe1 v1) k2 (coe1 v2)

-- ** Disjoint

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

-- ** Compose

compose :: (Enum1 k1, Enum1 k2, TestEquality k2) => DEnumMap k2 v -> DEnumMap k1 k2 -> DEnumMap k1 v
compose m2v (DEnumMap m12) =
  DEnumMap (IM.mapMaybe (\(KV inf1 k2) -> KV inf1 <$> m2v !? k2) m12)

-- ** 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' (coerce g1) (coerce 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)

-- * Traversal
-- ** Map

map :: (Enum1 k, TestEquality k) => (forall a. v1 a -> v2 a) -> DEnumMap k v1 -> DEnumMap k v2
map f = mapWithKey (const f)

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 f, Enum1 k, TestEquality k)
                => (forall a. k a -> v1 a -> f (v2 a)) -> DEnumMap k v1 -> f (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 f, Enum1 k, TestEquality k)
                     => (forall a. k a -> v1 a -> f (Maybe (v2 a))) -> DEnumMap k v1 -> f (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 :: (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

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

foldrWithKey :: Enum1 k => (forall a. k a -> v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldrWithKey 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

foldlWithKey :: Enum1 k => (forall a. acc -> k a -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldlWithKey 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

foldrWithKey' :: Enum1 k => (forall a. k a -> v a -> acc -> acc) -> acc -> DEnumMap k v -> acc
foldrWithKey' 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

foldlWithKey' :: Enum1 k => (forall a. acc -> k a -> v a -> acc) -> acc -> DEnumMap k v -> acc
foldlWithKey' 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: Wait for 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 :: (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: 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)

-- | \(O(\min(n,W)^2)\). Because of the lack of a @takeWhileAntitoneWithValue@
-- operation on 'IntMap', this function performs additional lookups to
-- reconstruct the full keys to pass to the predicate, resulting in a somewhat
-- worse complexity than 'IM.takeWhileAntitone'.
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)

-- | \(O(\min(n,W)^2)\). See 'takeWhileAntitone'.
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)

-- | \(O(\min(n,W)^2)\). See 'takeWhileAntitone'.
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)

mapMaybe :: (Enum1 k, TestEquality k) => (forall a. v1 a -> Maybe (v2 a)) -> DEnumMap k v1 -> DEnumMap k v2
mapMaybe f = mapMaybeWithKey (const f)

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 :: (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 (const f)

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)

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
     -- Note: this coe1 is fine because of the invariant on DEnumMap.
  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

-- TODO: the submap operations can't check any laws because there is no IM.isSubmapOfByKey.
isSubmapOf :: (forall a. Eq (v a)) => DEnumMap k v -> DEnumMap k v -> Bool
isSubmapOf (DEnumMap m1) (DEnumMap m2) = IM.isSubmapOfBy (\(KV _ v1) (KV _ v2) -> v1 == coe1 v2) m1 m2

isSubmapOfBy :: (forall a. v1 a -> v2 a -> Bool) -> DEnumMap k v1 -> DEnumMap k v2 -> Bool
isSubmapOfBy f (DEnumMap m1) (DEnumMap m2) =
  IM.isSubmapOfBy (\(KV _ v1) (KV _ v2) -> f v1 (coe1 v2)) m1 m2

isProperSubmapOf :: (forall a. Eq (v a)) => DEnumMap k v -> DEnumMap k v -> Bool
isProperSubmapOf (DEnumMap m1) (DEnumMap m2) = IM.isProperSubmapOfBy (\(KV _ v1) (KV _ v2) -> v1 == coe1 v2) m1 m2

isProperSubmapOfBy :: (forall a. v1 a -> v2 a -> Bool) -> DEnumMap k v1 -> DEnumMap k v2 -> Bool
isProperSubmapOfBy f (DEnumMap m1) (DEnumMap m2) =
  IM.isProperSubmapOfBy (\(KV _ v1) (KV _ v2) -> f v1 (coe1 v2)) m1 m2

-- * 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 :: Enum1 k => (forall a. v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMin f = updateMinWithKey (const f)

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 :: Enum1 k => (forall a. v a -> Maybe (v a)) -> DEnumMap k v -> DEnumMap k v
updateMax f = updateMaxWithKey (const f)

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