{-# 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 :: SMethod a -> Int
toIx SMethod a
k = Int# -> Int
I# (SMethod a -> Int#
forall a. a -> Int#
dataToTag# SMethod a
k)

singleton :: SMethod a -> v a -> SMethodMap v
singleton :: SMethod a -> v a -> SMethodMap v
singleton (SCustomMethod Text
t) v a
v = IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap IntMap (v Any)
forall a. Monoid a => a
mempty (Text -> v a -> Map Text (v a)
forall k a. k -> a -> Map k a
Map.singleton Text
t v a
v)
singleton SMethod a
k v a
v = IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
forall (f :: From) (t :: MethodType) (v :: Method f t -> *).
IntMap (v Any) -> Map Text (v 'CustomMethod) -> SMethodMap v
SMethodMap (Int -> v Any -> IntMap (v Any)
forall a. Int -> a -> IntMap a
IntMap.singleton (SMethod a -> Int
forall (f :: From) (t :: MethodType) (a :: Method f t).
SMethod a -> Int
toIx SMethod a
k) (v a -> v Any
forall a b. a -> b
unsafeCoerce v a
v)) Map Text (v 'CustomMethod)
forall a. Monoid a => a
mempty

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

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