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

module Data.IxMap where

import qualified Data.Map.Strict as M
import Data.Some
import Data.Kind
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 = Map (Base k) (Some f) -> IxMap k f
forall a (k :: a -> *) (f :: a -> *).
Map (Base k) (Some f) -> IxMap k f
IxMap Map (Base k) (Some f)
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 (k m -> Base k
forall (a :: a). k a -> Base k
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)
  | Base k -> Map (Base k) (Some f) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.notMember Base k
i Map (Base k) (Some f)
m = IxMap k f -> Maybe (IxMap k f)
forall a. a -> Maybe a
Just (IxMap k f -> Maybe (IxMap k f)) -> IxMap k f -> Maybe (IxMap k f)
forall a b. (a -> b) -> a -> b
$ Map (Base k) (Some f) -> IxMap k f
forall a (k :: a -> *) (f :: a -> *).
Map (Base k) (Some f) -> IxMap k f
IxMap (Map (Base k) (Some f) -> IxMap k f)
-> Map (Base k) (Some f) -> IxMap k f
forall a b. (a -> b) -> a -> b
$ Base k -> Some f -> Map (Base k) (Some f) -> Map (Base k) (Some f)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Base k
i (f m -> Some f
forall {k} (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome f m
x) Map (Base k) (Some f)
m
  | Bool
otherwise = Maybe (IxMap k f)
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 Base k -> Map (Base k) (Some f) -> Maybe (Some f)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (k m -> Base k
forall (a :: a). k a -> Base k
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) -> f m -> Maybe (f m)
forall a. a -> Maybe a
Just (f m -> Maybe (f m)) -> f m -> Maybe (f m)
forall a b. (a -> b) -> a -> b
$ f a -> f m
forall a b. a -> b
unsafeCoerce f a
v
    Maybe (Some f)
Nothing -> Maybe (f m)
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 (Base k -> Some f -> Maybe (Some f))
-> Base k
-> Map (Base k) (Some f)
-> (Maybe (Some f), Map (Base k) (Some f))
forall k a.
Ord k =>
(k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
M.updateLookupWithKey (\Base k
_ Some f
_ -> Maybe (Some f)
forall a. Maybe a
Nothing) (k m -> Base k
forall (a :: a). k a -> Base k
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') -> (Maybe (f m)
forall a. Maybe a
Nothing,Map (Base k) (Some f) -> IxMap k f
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') -> (f m -> Maybe (f m)
forall a. a -> Maybe a
Just (f a -> f m
forall a b. a -> b
unsafeCoerce f a
k),Map (Base k) (Some f) -> IxMap k f
forall a (k :: a -> *) (f :: a -> *).
Map (Base k) (Some f) -> IxMap k f
IxMap Map (Base k) (Some f)
m')