{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}

module Language.LSP.Protocol.Utils.SMethodMap (
  SMethodMap,
  singleton,
  insert,
  delete,
  member,
  lookup,
  map,
) where

import Data.IntMap (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Kind (Type)
import Data.Map (Map)
import Data.Map.Strict qualified as Map
import GHC.Exts (Any, Int (..), dataToTag#)
import Unsafe.Coerce (unsafeCoerce)
import Prelude hiding (lookup, map)

import GHC.TypeLits (symbolVal)
import Language.LSP.Protocol.Message (Method (..), SMethod (..))

-- This type exists to avoid a dependency on 'dependent-map'. It is less
-- safe (since we use 'unsafeCoerce') but much simpler and hence easier to include.

-- | A specialized alternative to a full dependent map for use with 'SMethod'.
data SMethodMap (v :: Method f t -> Type)
  = -- This works by using an 'IntMap' indexed by constructor tag for the majority
    -- of 'SMethod's, which have no parameters, and hence can only appear once as keys
    -- in the map. We do not attempt to be truly dependent here, and instead exploit
    -- 'usafeCoerce' to go to and from 'v Any'.
    -- The sole exception is 'SCustomMethod', for which we keep a separate map from
    -- its 'Text' parameter
    SMethodMap !(IntMap (v Any)) !(Map String (v Any))

toIx :: SMethod a -> Int
toIx :: forall {f :: MessageDirection} {t :: MessageKind}
       (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k = Int# -> Int
I# (forall a. a -> Int#
dataToTag# SMethod a
k)

singleton :: SMethod a -> v a -> SMethodMap v
singleton :: forall {f :: MessageDirection} {t :: MessageKind} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> v a -> SMethodMap v
singleton (SMethod_CustomMethod Proxy s
t) v a
v = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap forall a. Monoid a => a
mempty (forall k a. k -> a -> Map k a
Map.singleton (forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal Proxy s
t) (forall a b. a -> b
unsafeCoerce v a
v))
singleton SMethod a
k v a
v = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap (forall a. Int -> a -> IntMap a
IntMap.singleton (forall {f :: MessageDirection} {t :: MessageKind}
       (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k) (forall a b. a -> b
unsafeCoerce v a
v)) forall a. Monoid a => a
mempty

insert :: SMethod a -> v a -> SMethodMap v -> SMethodMap v
insert :: forall {f :: MessageDirection} {t :: MessageKind} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> v a -> SMethodMap v -> SMethodMap v
insert (SMethod_CustomMethod Proxy s
t) v a
v (SMethodMap IntMap (v Any)
xs Map String (v Any)
ys) = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap IntMap (v Any)
xs (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal Proxy s
t) (forall a b. a -> b
unsafeCoerce v a
v) Map String (v Any)
ys)
insert SMethod a
k v a
v (SMethodMap IntMap (v Any)
xs Map String (v Any)
ys) = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap (forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert (forall {f :: MessageDirection} {t :: MessageKind}
       (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k) (forall a b. a -> b
unsafeCoerce v a
v) IntMap (v Any)
xs) Map String (v Any)
ys

delete :: SMethod a -> SMethodMap v -> SMethodMap v
delete :: forall {f :: MessageDirection} {t :: MessageKind}
       {f :: MessageDirection} {t :: MessageKind} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> SMethodMap v -> SMethodMap v
delete (SMethod_CustomMethod Proxy s
t) (SMethodMap IntMap (v Any)
xs Map String (v Any)
ys) = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap IntMap (v Any)
xs (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal Proxy s
t) Map String (v Any)
ys)
delete SMethod a
k (SMethodMap IntMap (v Any)
xs Map String (v Any)
ys) = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap (forall a. Int -> IntMap a -> IntMap a
IntMap.delete (forall {f :: MessageDirection} {t :: MessageKind}
       (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k) IntMap (v Any)
xs) Map String (v Any)
ys

member :: SMethod a -> SMethodMap v -> Bool
member :: forall {f :: MessageDirection} {t :: MessageKind}
       {f :: MessageDirection} {t :: MessageKind} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> SMethodMap v -> Bool
member (SMethod_CustomMethod Proxy s
t) (SMethodMap IntMap (v Any)
_ Map String (v Any)
ys) = forall k a. Ord k => k -> Map k a -> Bool
Map.member (forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal Proxy s
t) Map String (v Any)
ys
member SMethod a
k (SMethodMap IntMap (v Any)
xs Map String (v Any)
_) = forall a. Int -> IntMap a -> Bool
IntMap.member (forall {f :: MessageDirection} {t :: MessageKind}
       (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k) IntMap (v Any)
xs

lookup :: SMethod a -> SMethodMap v -> Maybe (v a)
lookup :: forall {f :: MessageDirection} {t :: MessageKind} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> SMethodMap v -> Maybe (v a)
lookup (SMethod_CustomMethod Proxy s
t) (SMethodMap IntMap (v Any)
_ Map String (v Any)
ys) = forall a b. a -> b
unsafeCoerce (forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal Proxy s
t) Map String (v Any)
ys)
lookup SMethod a
k (SMethodMap IntMap (v Any)
xs Map String (v Any)
_) = forall a b. a -> b
unsafeCoerce (forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (forall {f :: MessageDirection} {t :: MessageKind}
       (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k) IntMap (v Any)
xs)

map :: (forall a. u a -> v a) -> SMethodMap u -> SMethodMap v
map :: forall {f :: MessageDirection} {t :: MessageKind}
       (u :: Method f t -> *) (v :: Method f t -> *).
(forall (a :: Method f t). u a -> v a)
-> SMethodMap u -> SMethodMap v
map forall (a :: Method f t). u a -> v a
f (SMethodMap IntMap (u Any)
xs Map String (u Any)
ys) = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap (forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map forall (a :: Method f t). u a -> v a
f IntMap (u Any)
xs) (forall a b k. (a -> b) -> Map k a -> Map k b
Map.map forall (a :: Method f t). u a -> v a
f Map String (u Any)
ys)

instance Semigroup (SMethodMap v) where
  SMethodMap IntMap (v Any)
xs Map String (v Any)
ys <> :: SMethodMap v -> SMethodMap v -> SMethodMap v
<> SMethodMap IntMap (v Any)
xs' Map String (v Any)
ys' = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap (IntMap (v Any)
xs forall a. Semigroup a => a -> a -> a
<> IntMap (v Any)
xs') (Map String (v Any)
ys forall a. Semigroup a => a -> a -> a
<> Map String (v Any)
ys')

instance Monoid (SMethodMap v) where
  mempty :: SMethodMap v
mempty = forall (f :: MessageDirection) (t :: MessageKind)
       (v :: Method f t -> *).
IntMap (v Any) -> Map String (v Any) -> SMethodMap v
SMethodMap forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty