{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE ViewPatterns #-}

module Data.IxMap where

import Data.Kind
import Data.Map.Strict qualified as M
import Data.Some
import Unsafe.Coerce

-- a `compare` b <=> toBase a `compare` toBase b
-- toBase (i :: f a) == toBase (j :: f b) <=> a ~ b
class Ord (Base f) => IxOrd f where
  type Base f
  toBase :: forall a. f a -> Base f

newtype IxMap (k :: a -> Type) (f :: a -> Type) = IxMap {forall a (k :: a -> *) (f :: a -> *).
IxMap k f -> Map (Base k) (Some f)
getMap :: M.Map (Base k) (Some f)}

emptyIxMap :: IxMap k f
emptyIxMap :: forall {a} (k :: a -> *) (f :: a -> *). IxMap k f
emptyIxMap = forall a (k :: a -> *) (f :: a -> *).
Map (Base k) (Some f) -> IxMap k f
IxMap forall k a. Map k a
M.empty

insertIxMap :: IxOrd k => k m -> f m -> IxMap k f -> Maybe (IxMap k f)
insertIxMap :: forall {a} (k :: a -> *) (m :: a) (f :: a -> *).
IxOrd k =>
k m -> f m -> IxMap k f -> Maybe (IxMap k f)
insertIxMap (forall {k} (f :: k -> *) (a :: k). IxOrd f => f a -> Base f
toBase -> Base k
i) f m
x (IxMap Map (Base k) (Some f)
m)
  | forall k a. Ord k => k -> Map k a -> Bool
M.notMember Base k
i Map (Base k) (Some f)
m = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a (k :: a -> *) (f :: a -> *).
Map (Base k) (Some f) -> IxMap k f
IxMap forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Base k
i (forall {k} (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome f m
x) Map (Base k) (Some f)
m
  | Bool
otherwise = forall a. Maybe a
Nothing

lookupIxMap :: IxOrd k => k m -> IxMap k f -> Maybe (f m)
lookupIxMap :: forall {a} (k :: a -> *) (m :: a) (f :: a -> *).
IxOrd k =>
k m -> IxMap k f -> Maybe (f m)
lookupIxMap k m
i (IxMap Map (Base k) (Some f)
m) =
  case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall {k} (f :: k -> *) (a :: k). IxOrd f => f a -> Base f
toBase k m
i) Map (Base k) (Some f)
m of
    Just (Some f a
v) -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. a -> b
unsafeCoerce f a
v
    Maybe (Some f)
Nothing -> forall a. Maybe a
Nothing

pickFromIxMap :: IxOrd k => k m -> IxMap k f -> (Maybe (f m), IxMap k f)
pickFromIxMap :: forall {a} (k :: a -> *) (m :: a) (f :: a -> *).
IxOrd k =>
k m -> IxMap k f -> (Maybe (f m), IxMap k f)
pickFromIxMap k m
i (IxMap Map (Base k) (Some f)
m) =
  case forall k a.
Ord k =>
(k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
M.updateLookupWithKey (\Base k
_ Some f
_ -> forall a. Maybe a
Nothing) (forall {k} (f :: k -> *) (a :: k). IxOrd f => f a -> Base f
toBase k m
i) Map (Base k) (Some f)
m of
    (Maybe (Some f)
Nothing, !Map (Base k) (Some f)
m') -> (forall a. Maybe a
Nothing, forall a (k :: a -> *) (f :: a -> *).
Map (Base k) (Some f) -> IxMap k f
IxMap Map (Base k) (Some f)
m')
    (Just (Some f a
k), !Map (Base k) (Some f)
m') -> (forall a. a -> Maybe a
Just (forall a b. a -> b
unsafeCoerce f a
k), forall a (k :: a -> *) (f :: a -> *).
Map (Base k) (Some f) -> IxMap k f
IxMap Map (Base k) (Some f)
m')