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

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

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

import Language.LSP.Types.Method (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 (and where we can get the type indices right).
  SMethodMap !(IntMap (v Any)) !(Map Text (v 'CustomMethod))

toIx :: SMethod a -> Int
toIx :: forall {f :: From} {t :: MethodType} (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 :: From} {t :: MethodType} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> v a -> SMethodMap v
singleton (SCustomMethod Text
t) v a
v = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap forall a. Monoid a => a
mempty (forall k a. k -> a -> Map k a
Map.singleton Text
t v a
v)
singleton SMethod a
k v a
v = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap (forall a. Int -> a -> IntMap a
IntMap.singleton (forall {f :: From} {t :: MethodType} (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 :: From} {t :: MethodType} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> v a -> SMethodMap v -> SMethodMap v
insert (SCustomMethod Text
t) v a
v (SMethodMap IntMap (v Any)
xs Map Text (v 'CustomMethod)
ys) = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap IntMap (v Any)
xs (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
t v a
v Map Text (v 'CustomMethod)
ys)
insert SMethod a
k v a
v (SMethodMap IntMap (v Any)
xs Map Text (v 'CustomMethod)
ys) = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap (forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert (forall {f :: From} {t :: MethodType} (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 Text (v 'CustomMethod)
ys

delete :: SMethod a -> SMethodMap v -> SMethodMap v
delete :: forall {f :: From} {t :: MethodType} {f :: From} {t :: MethodType}
       (a :: Method f t) (v :: Method f t -> *).
SMethod a -> SMethodMap v -> SMethodMap v
delete (SCustomMethod Text
t) (SMethodMap IntMap (v Any)
xs Map Text (v 'CustomMethod)
ys) = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap IntMap (v Any)
xs (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Text
t Map Text (v 'CustomMethod)
ys)
delete SMethod a
k (SMethodMap IntMap (v Any)
xs Map Text (v 'CustomMethod)
ys) = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap (forall a. Int -> IntMap a -> IntMap a
IntMap.delete (forall {f :: From} {t :: MethodType} (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k) IntMap (v Any)
xs) Map Text (v 'CustomMethod)
ys

member :: SMethod a -> SMethodMap v -> Bool
member :: forall {f :: From} {t :: MethodType} {f :: From} {t :: MethodType}
       (a :: Method f t) (v :: Method f t -> *).
SMethod a -> SMethodMap v -> Bool
member (SCustomMethod Text
t) (SMethodMap IntMap (v Any)
_ Map Text (v 'CustomMethod)
ys) = forall k a. Ord k => k -> Map k a -> Bool
Map.member Text
t Map Text (v 'CustomMethod)
ys
member SMethod a
k (SMethodMap IntMap (v Any)
xs Map Text (v 'CustomMethod)
_) = forall a. Int -> IntMap a -> Bool
IntMap.member (forall {f :: From} {t :: MethodType} (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 :: From} {t :: MethodType} (a :: Method f t)
       (v :: Method f t -> *).
SMethod a -> SMethodMap v -> Maybe (v a)
lookup (SCustomMethod Text
t) (SMethodMap IntMap (v Any)
_ Map Text (v 'CustomMethod)
ys) = forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
t Map Text (v 'CustomMethod)
ys
lookup SMethod a
k (SMethodMap IntMap (v Any)
xs Map Text (v 'CustomMethod)
_) = forall a b. a -> b
unsafeCoerce (forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (forall {f :: From} {t :: MethodType} (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 :: From} {t :: MethodType} (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 Text (u 'CustomMethod)
ys) = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> 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 Text (u 'CustomMethod)
ys)

instance Semigroup (SMethodMap v) where
  SMethodMap IntMap (v Any)
xs Map Text (v 'CustomMethod)
ys <> :: SMethodMap v -> SMethodMap v -> SMethodMap v
<> SMethodMap IntMap (v Any)
xs' Map Text (v 'CustomMethod)
ys' = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap (IntMap (v Any)
xs forall a. Semigroup a => a -> a -> a
<> IntMap (v Any)
xs') (Map Text (v 'CustomMethod)
ys forall a. Semigroup a => a -> a -> a
<> Map Text (v 'CustomMethod)
ys')

instance Monoid (SMethodMap v) where
  mempty :: SMethodMap v
mempty = forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty