{-# LANGUAGE DeriveFunctor #-}

module Subscription
(
  SubscriptionTree (..),
  broadcast,
  broadcast',
  empty,
  subscribe,
  unsubscribe,
  showTree,
)
where

import Control.Monad (void)
import Control.Monad.Writer (Writer, tell, execWriter)
import Data.Aeson (Value)
import Data.Foldable (for_, traverse_)
import Data.HashMap.Strict (HashMap)
import Data.Hashable (Hashable)
import Data.Maybe (fromMaybe)
import Data.Monoid ((<>))
import Data.Text (Text)

import qualified Control.Concurrent.Async as Async
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Text as Text

import qualified Store

-- Keeps subscriptions in a tree data structure, so we can efficiently determine
-- which clients need to be notified for a given update.
data SubscriptionTree id conn =
  SubscriptionTree (HashMap id conn) (HashMap Text (SubscriptionTree id conn))
  deriving (SubscriptionTree id conn -> SubscriptionTree id conn -> Bool
(SubscriptionTree id conn -> SubscriptionTree id conn -> Bool)
-> (SubscriptionTree id conn -> SubscriptionTree id conn -> Bool)
-> Eq (SubscriptionTree id conn)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall id conn.
(Eq id, Eq conn) =>
SubscriptionTree id conn -> SubscriptionTree id conn -> Bool
/= :: SubscriptionTree id conn -> SubscriptionTree id conn -> Bool
$c/= :: forall id conn.
(Eq id, Eq conn) =>
SubscriptionTree id conn -> SubscriptionTree id conn -> Bool
== :: SubscriptionTree id conn -> SubscriptionTree id conn -> Bool
$c== :: forall id conn.
(Eq id, Eq conn) =>
SubscriptionTree id conn -> SubscriptionTree id conn -> Bool
Eq, a -> SubscriptionTree id b -> SubscriptionTree id a
(a -> b) -> SubscriptionTree id a -> SubscriptionTree id b
(forall a b.
 (a -> b) -> SubscriptionTree id a -> SubscriptionTree id b)
-> (forall a b.
    a -> SubscriptionTree id b -> SubscriptionTree id a)
-> Functor (SubscriptionTree id)
forall a b. a -> SubscriptionTree id b -> SubscriptionTree id a
forall a b.
(a -> b) -> SubscriptionTree id a -> SubscriptionTree id b
forall id a b. a -> SubscriptionTree id b -> SubscriptionTree id a
forall id a b.
(a -> b) -> SubscriptionTree id a -> SubscriptionTree id b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> SubscriptionTree id b -> SubscriptionTree id a
$c<$ :: forall id a b. a -> SubscriptionTree id b -> SubscriptionTree id a
fmap :: (a -> b) -> SubscriptionTree id a -> SubscriptionTree id b
$cfmap :: forall id a b.
(a -> b) -> SubscriptionTree id a -> SubscriptionTree id b
Functor, Int -> SubscriptionTree id conn -> ShowS
[SubscriptionTree id conn] -> ShowS
SubscriptionTree id conn -> String
(Int -> SubscriptionTree id conn -> ShowS)
-> (SubscriptionTree id conn -> String)
-> ([SubscriptionTree id conn] -> ShowS)
-> Show (SubscriptionTree id conn)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall id conn.
(Show id, Show conn) =>
Int -> SubscriptionTree id conn -> ShowS
forall id conn.
(Show id, Show conn) =>
[SubscriptionTree id conn] -> ShowS
forall id conn.
(Show id, Show conn) =>
SubscriptionTree id conn -> String
showList :: [SubscriptionTree id conn] -> ShowS
$cshowList :: forall id conn.
(Show id, Show conn) =>
[SubscriptionTree id conn] -> ShowS
show :: SubscriptionTree id conn -> String
$cshow :: forall id conn.
(Show id, Show conn) =>
SubscriptionTree id conn -> String
showsPrec :: Int -> SubscriptionTree id conn -> ShowS
$cshowsPrec :: forall id conn.
(Show id, Show conn) =>
Int -> SubscriptionTree id conn -> ShowS
Show)

empty :: SubscriptionTree id conn
empty :: SubscriptionTree id conn
empty = HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
forall id conn.
HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
SubscriptionTree HashMap id conn
forall k v. HashMap k v
HashMap.empty HashMap Text (SubscriptionTree id conn)
forall k v. HashMap k v
HashMap.empty

isEmpty :: SubscriptionTree id conn -> Bool
isEmpty :: SubscriptionTree id conn -> Bool
isEmpty (SubscriptionTree HashMap id conn
here HashMap Text (SubscriptionTree id conn)
inner) = HashMap id conn -> Bool
forall k v. HashMap k v -> Bool
HashMap.null HashMap id conn
here Bool -> Bool -> Bool
&& HashMap Text (SubscriptionTree id conn) -> Bool
forall k v. HashMap k v -> Bool
HashMap.null HashMap Text (SubscriptionTree id conn)
inner

subscribe
  :: (Eq id, Hashable id)
  => [Text]
  -> id
  -> conn
  -> SubscriptionTree id conn
  -> SubscriptionTree id conn
subscribe :: [Text]
-> id
-> conn
-> SubscriptionTree id conn
-> SubscriptionTree id conn
subscribe [Text]
path id
subid conn
subval (SubscriptionTree HashMap id conn
here HashMap Text (SubscriptionTree id conn)
inner) =
  case [Text]
path of
    [] -> HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
forall id conn.
HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
SubscriptionTree (id -> conn -> HashMap id conn -> HashMap id conn
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert id
subid conn
subval HashMap id conn
here) HashMap Text (SubscriptionTree id conn)
inner
    Text
key : [Text]
pathTail ->
      let
        subscribeInner :: SubscriptionTree id conn -> SubscriptionTree id conn
subscribeInner = [Text]
-> id
-> conn
-> SubscriptionTree id conn
-> SubscriptionTree id conn
forall id conn.
(Eq id, Hashable id) =>
[Text]
-> id
-> conn
-> SubscriptionTree id conn
-> SubscriptionTree id conn
subscribe [Text]
pathTail id
subid conn
subval
        newInner :: HashMap Text (SubscriptionTree id conn)
newInner = (Maybe (SubscriptionTree id conn)
 -> Maybe (SubscriptionTree id conn))
-> Text
-> HashMap Text (SubscriptionTree id conn)
-> HashMap Text (SubscriptionTree id conn)
forall k v.
(Eq k, Hashable k) =>
(Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
HashMap.alter (SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
forall a. a -> Maybe a
Just (SubscriptionTree id conn -> Maybe (SubscriptionTree id conn))
-> (Maybe (SubscriptionTree id conn) -> SubscriptionTree id conn)
-> Maybe (SubscriptionTree id conn)
-> Maybe (SubscriptionTree id conn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubscriptionTree id conn -> SubscriptionTree id conn
subscribeInner (SubscriptionTree id conn -> SubscriptionTree id conn)
-> (Maybe (SubscriptionTree id conn) -> SubscriptionTree id conn)
-> Maybe (SubscriptionTree id conn)
-> SubscriptionTree id conn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubscriptionTree id conn
-> Maybe (SubscriptionTree id conn) -> SubscriptionTree id conn
forall a. a -> Maybe a -> a
fromMaybe SubscriptionTree id conn
forall id conn. SubscriptionTree id conn
empty) Text
key HashMap Text (SubscriptionTree id conn)
inner
      in
        HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
forall id conn.
HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
SubscriptionTree HashMap id conn
here HashMap Text (SubscriptionTree id conn)
newInner

unsubscribe
  :: (Eq id, Hashable id)
  => [Text]
  -> id
  -> SubscriptionTree id conn
  -> SubscriptionTree id conn
unsubscribe :: [Text]
-> id -> SubscriptionTree id conn -> SubscriptionTree id conn
unsubscribe [Text]
path id
subid (SubscriptionTree HashMap id conn
here HashMap Text (SubscriptionTree id conn)
inner) =
  case [Text]
path of
    [] -> HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
forall id conn.
HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
SubscriptionTree (id -> HashMap id conn -> HashMap id conn
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HashMap.delete id
subid HashMap id conn
here) HashMap Text (SubscriptionTree id conn)
inner
    Text
key : [Text]
pathTail ->
      let
        -- Remove the tail from the inner tree (if it exists). If that left the
        -- inner tree empty, remove the key altogether to keep the tree clean.
        justNotEmpty :: SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
justNotEmpty SubscriptionTree id conn
tree = if SubscriptionTree id conn -> Bool
forall id conn. SubscriptionTree id conn -> Bool
isEmpty SubscriptionTree id conn
tree then Maybe (SubscriptionTree id conn)
forall a. Maybe a
Nothing else SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
forall a. a -> Maybe a
Just SubscriptionTree id conn
tree
        unsubscribeInner :: SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
unsubscribeInner = SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
forall id conn.
SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
justNotEmpty (SubscriptionTree id conn -> Maybe (SubscriptionTree id conn))
-> (SubscriptionTree id conn -> SubscriptionTree id conn)
-> SubscriptionTree id conn
-> Maybe (SubscriptionTree id conn)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text]
-> id -> SubscriptionTree id conn -> SubscriptionTree id conn
forall id conn.
(Eq id, Hashable id) =>
[Text]
-> id -> SubscriptionTree id conn -> SubscriptionTree id conn
unsubscribe [Text]
pathTail id
subid
        newInner :: HashMap Text (SubscriptionTree id conn)
newInner = (SubscriptionTree id conn -> Maybe (SubscriptionTree id conn))
-> Text
-> HashMap Text (SubscriptionTree id conn)
-> HashMap Text (SubscriptionTree id conn)
forall k a.
(Eq k, Hashable k) =>
(a -> Maybe a) -> k -> HashMap k a -> HashMap k a
HashMap.update SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
forall conn.
SubscriptionTree id conn -> Maybe (SubscriptionTree id conn)
unsubscribeInner Text
key HashMap Text (SubscriptionTree id conn)
inner
      in
        HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
forall id conn.
HashMap id conn
-> HashMap Text (SubscriptionTree id conn)
-> SubscriptionTree id conn
SubscriptionTree HashMap id conn
here HashMap Text (SubscriptionTree id conn)
newInner

-- Invoke f for all subscribers to the path. The subscribers get passed the
-- subvalue at the path that they are subscribed to.
broadcast :: (conn -> Value -> IO ()) -> [Text] -> Value -> SubscriptionTree id conn -> IO ()
broadcast :: (conn -> Value -> IO ())
-> [Text] -> Value -> SubscriptionTree id conn -> IO ()
broadcast conn -> Value -> IO ()
f [Text]
path Value
value SubscriptionTree id conn
tree =
  -- We broadcast concurrently since all updates are independent of each other
  ((conn, Value) -> IO ()) -> [(conn, Value)] -> IO ()
forall (f :: * -> *) a b. Foldable f => (a -> IO b) -> f a -> IO ()
Async.mapConcurrently_ ((conn -> Value -> IO ()) -> (conn, Value) -> IO ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry conn -> Value -> IO ()
f) [(conn, Value)]
notifications
  where notifications :: [(conn, Value)]
notifications = [Text] -> Value -> SubscriptionTree id conn -> [(conn, Value)]
forall id conn.
[Text] -> Value -> SubscriptionTree id conn -> [(conn, Value)]
broadcast' [Text]
path Value
value SubscriptionTree id conn
tree

-- Like broadcast, but return a list of notifications rather than invoking an
-- effect on each of them.
broadcast' :: [Text] -> Value -> SubscriptionTree id conn -> [(conn, Value)]
broadcast' :: [Text] -> Value -> SubscriptionTree id conn -> [(conn, Value)]
broadcast' = \[Text]
path Value
value SubscriptionTree id conn
tree -> Writer [(conn, Value)] () -> [(conn, Value)]
forall w a. Writer w a -> w
execWriter (Writer [(conn, Value)] () -> [(conn, Value)])
-> Writer [(conn, Value)] () -> [(conn, Value)]
forall a b. (a -> b) -> a -> b
$ [Text]
-> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
forall id conn.
[Text]
-> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
loop [Text]
path Value
value SubscriptionTree id conn
tree
  where
  loop :: [Text] -> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
  loop :: [Text]
-> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
loop [Text]
path Value
value (SubscriptionTree HashMap id conn
here HashMap Text (SubscriptionTree id conn)
inner) = do
    case [Text]
path of
      [] -> do
        -- When the path is empty, all subscribers that are "here" or at a deeper
        -- level should receive a notification.
        (conn -> Writer [(conn, Value)] ())
-> HashMap id conn -> Writer [(conn, Value)] ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (\conn
v -> [(conn, Value)] -> Writer [(conn, Value)] ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [(conn
v, Value
value)]) HashMap id conn
here
        let broadcastInner :: Text -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
broadcastInner Text
key = [Text]
-> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
forall id conn.
[Text]
-> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
loop [] ([Text] -> Value -> Value
Store.lookupOrNull [Text
key] Value
value)
        WriterT [(conn, Value)] Identity (HashMap Text ())
-> Writer [(conn, Value)] ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (WriterT [(conn, Value)] Identity (HashMap Text ())
 -> Writer [(conn, Value)] ())
-> WriterT [(conn, Value)] Identity (HashMap Text ())
-> Writer [(conn, Value)] ()
forall a b. (a -> b) -> a -> b
$ (Text -> SubscriptionTree id conn -> Writer [(conn, Value)] ())
-> HashMap Text (SubscriptionTree id conn)
-> WriterT [(conn, Value)] Identity (HashMap Text ())
forall (f :: * -> *) k v1 v2.
Applicative f =>
(k -> v1 -> f v2) -> HashMap k v1 -> f (HashMap k v2)
HashMap.traverseWithKey Text -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
forall id conn.
Text -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
broadcastInner HashMap Text (SubscriptionTree id conn)
inner

      Text
key : [Text]
pathTail -> do
        (conn -> Writer [(conn, Value)] ())
-> HashMap id conn -> Writer [(conn, Value)] ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (\conn
v -> [(conn, Value)] -> Writer [(conn, Value)] ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [(conn
v, Value
value)]) HashMap id conn
here
        Maybe (SubscriptionTree id conn)
-> (SubscriptionTree id conn -> Writer [(conn, Value)] ())
-> Writer [(conn, Value)] ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Text
-> HashMap Text (SubscriptionTree id conn)
-> Maybe (SubscriptionTree id conn)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
key HashMap Text (SubscriptionTree id conn)
inner) ((SubscriptionTree id conn -> Writer [(conn, Value)] ())
 -> Writer [(conn, Value)] ())
-> (SubscriptionTree id conn -> Writer [(conn, Value)] ())
-> Writer [(conn, Value)] ()
forall a b. (a -> b) -> a -> b
$ \SubscriptionTree id conn
subs ->
          [Text]
-> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
forall id conn.
[Text]
-> Value -> SubscriptionTree id conn -> Writer [(conn, Value)] ()
loop [Text]
pathTail ([Text] -> Value -> Value
Store.lookupOrNull [Text
key] Value
value) SubscriptionTree id conn
subs

-- Show subscriptions, for debugging purposes.
showTree :: Show id => SubscriptionTree id conn -> String
showTree :: SubscriptionTree id conn -> String
showTree SubscriptionTree id conn
tree =
  let
    withPrefix :: String -> SubscriptionTree a v -> String
withPrefix String
prefix (SubscriptionTree HashMap a v
here HashMap Text (SubscriptionTree a v)
inner) =
      let
        strHere :: String
        strHere :: String
strHere = (a -> String) -> [a] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\a
cid -> String
" * " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> (a -> String
forall a. Show a => a -> String
show a
cid) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"\n") (HashMap a v -> [a]
forall k v. HashMap k v -> [k]
HashMap.keys HashMap a v
here)
        showInner :: String -> SubscriptionTree a v -> String
showInner String
iPrefix SubscriptionTree a v
t = String
iPrefix String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"\n" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String -> SubscriptionTree a v -> String
withPrefix String
iPrefix SubscriptionTree a v
t
        strInner :: String
        strInner :: String
strInner = HashMap Text String -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (HashMap Text String -> String) -> HashMap Text String -> String
forall a b. (a -> b) -> a -> b
$ (Text -> SubscriptionTree a v -> String)
-> HashMap Text (SubscriptionTree a v) -> HashMap Text String
forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapWithKey (\Text
key -> String -> SubscriptionTree a v -> String
showInner (String
prefix String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"/" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
Text.unpack Text
key)) HashMap Text (SubscriptionTree a v)
inner
      in
        String
strHere String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
strInner
  in
    String
"/\n" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> (String -> SubscriptionTree id conn -> String
forall a v. Show a => String -> SubscriptionTree a v -> String
withPrefix String
"" SubscriptionTree id conn
tree)