-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}

-- |
-- Most users should import "Haxl.Core" instead of importing this
-- module directly.
--
module Haxl.Core.StateStore
  ( StateKey(..)
  , StateStore
  , stateGet
  , stateSet
  , stateEmpty
  ) where

import Data.Map (Map)
import Data.Kind (Type)
import qualified Data.Map.Strict as Map
#if __GLASGOW_HASKELL__ < 804
import Data.Monoid
#endif
import Data.Typeable
import Unsafe.Coerce

-- | 'StateKey' maps one type to another type. A type that is an
-- instance of 'StateKey' can store and retrieve information from a
-- 'StateStore'.
--
class Typeable f => StateKey (f :: Type -> Type) where
  data State f

  -- | We default this to typeOf1, but if f is itself a complex type that is
  -- already applied to some paramaters, we want to be able to use the same
  -- state by using typeOf2, etc
  getStateType :: Proxy f -> TypeRep
  getStateType = Proxy f -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep

-- | The 'StateStore' maps a 'StateKey' to the 'State' for that type.
newtype StateStore = StateStore (Map TypeRep StateStoreData)

#if __GLASGOW_HASKELL__ >= 804
instance Semigroup StateStore where
  <> :: StateStore -> StateStore -> StateStore
(<>) = StateStore -> StateStore -> StateStore
forall a. Monoid a => a -> a -> a
mappend
#endif

instance Monoid StateStore where
  mempty :: StateStore
mempty = StateStore
stateEmpty
  -- Left-biased union
  mappend :: StateStore -> StateStore -> StateStore
mappend (StateStore Map TypeRep StateStoreData
m1) (StateStore Map TypeRep StateStoreData
m2) = Map TypeRep StateStoreData -> StateStore
StateStore (Map TypeRep StateStoreData -> StateStore)
-> Map TypeRep StateStoreData -> StateStore
forall a b. (a -> b) -> a -> b
$ Map TypeRep StateStoreData
m1 Map TypeRep StateStoreData
-> Map TypeRep StateStoreData -> Map TypeRep StateStoreData
forall a. Semigroup a => a -> a -> a
<> Map TypeRep StateStoreData
m2

-- | Encapsulates the type of 'StateStore' data so we can have a
-- heterogeneous collection.
data StateStoreData = forall f. StateKey f => StateStoreData (State f)

-- | A `StateStore` with no entries.
stateEmpty :: StateStore
stateEmpty :: StateStore
stateEmpty = Map TypeRep StateStoreData -> StateStore
StateStore Map TypeRep StateStoreData
forall k a. Map k a
Map.empty

-- | Inserts a `State` in the `StateStore` container.
stateSet :: forall f . StateKey f => State f -> StateStore -> StateStore
stateSet :: State f -> StateStore -> StateStore
stateSet State f
st (StateStore Map TypeRep StateStoreData
m) =
  Map TypeRep StateStoreData -> StateStore
StateStore (TypeRep
-> StateStoreData
-> Map TypeRep StateStoreData
-> Map TypeRep StateStoreData
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Proxy f -> TypeRep
forall (f :: * -> *). StateKey f => Proxy f -> TypeRep
getStateType (Proxy f
forall k (t :: k). Proxy t
Proxy :: Proxy f)) (State f -> StateStoreData
forall (f :: * -> *). StateKey f => State f -> StateStoreData
StateStoreData State f
st) Map TypeRep StateStoreData
m)

-- | Retrieves a `State` from the `StateStore` container.
stateGet :: forall r . StateKey r => StateStore -> Maybe (State r)
stateGet :: StateStore -> Maybe (State r)
stateGet (StateStore Map TypeRep StateStoreData
m) =
  case TypeRep -> Map TypeRep StateStoreData -> Maybe StateStoreData
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TypeRep
ty Map TypeRep StateStoreData
m of
     Maybe StateStoreData
Nothing -> Maybe (State r)
forall a. Maybe a
Nothing
     Just (StateStoreData (State f
st :: State f))
       | Proxy f -> TypeRep
forall (f :: * -> *). StateKey f => Proxy f -> TypeRep
getStateType (Proxy f
forall k (t :: k). Proxy t
Proxy :: Proxy f) TypeRep -> TypeRep -> Bool
forall a. Eq a => a -> a -> Bool
== TypeRep
ty -> State r -> Maybe (State r)
forall a. a -> Maybe a
Just (State f -> State r
forall a b. a -> b
unsafeCoerce State f
st)
       | Bool
otherwise             -> Maybe (State r)
forall a. Maybe a
Nothing
          -- the dynamic type check here should be unnecessary, but if
          -- there are bugs in `Typeable` or `Map` then we'll get an
          -- error instead of a crash.  The overhead is a few percent.
 where
  ty :: TypeRep
ty = Proxy r -> TypeRep
forall (f :: * -> *). StateKey f => Proxy f -> TypeRep
getStateType (Proxy r
forall k (t :: k). Proxy t
Proxy :: Proxy r)