content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE ExistentialQuantification #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE KindSignatures #-}
  {-# LANGUAGE BangPatterns #-}
  -- | Bucketing requests by 'DataSource'.
  --
  -- When a request is issued by the client via 'dataFetch', it is placed
  -- in the 'RequestStore'. When we are ready to fetch the current batch
  -- of requests, the 'contents' operation extracts the fetches, bucketed
  -- by 'DataSource'.
  --
  -- This module is provided for access to Haxl internals only; most
  -- users should not need to import it.
  --
  module Haxl.Core.RequestStore
    ( BlockedFetches(..)
    , BlockedFetchInternal(..)
    , RequestStore
    , isEmpty
    , noRequests
    , addRequest
    , contents
    , getSize
    , ReqCountMap(..)
    , emptyReqCounts
    , filterRCMap
    , getMapFromRCMap
    , getSummaryMapFromRCMap
    , addToCountMap
    , subFromCountMap
    ) where
  
  import Haxl.Core.DataSource
  import Haxl.Core.Stats
  import Data.Map (Map)
  import qualified Data.HashMap.Strict as HashMap
  import qualified Data.Map.Strict as Map
  import Data.Proxy
  import Data.Text (Text)
  import Data.Kind (Type)
  import Data.Typeable
  import Unsafe.Coerce
  
  -- | A container for multiple 'BlockedFetch' objects.
  newtype RequestStore u = RequestStore (Map TypeRep (BlockedFetches u))
    -- Since we don't know which data sources we will be using, the store
    -- is dynamically-typed.  It maps the TypeRep of the request to the
    -- 'BlockedFetches' for that 'DataSource'.
  
  newtype BlockedFetchInternal = BlockedFetchInternal CallId
  
  -- | A batch of 'BlockedFetch' objects for a single 'DataSource'
  data BlockedFetches u =
    forall r. (DataSource u r) =>
          BlockedFetches [BlockedFetch r] [BlockedFetchInternal]
  
  isEmpty :: RequestStore u -> Bool
  isEmpty (RequestStore m) = Map.null m
  
  -- | A new empty 'RequestStore'.
  noRequests :: RequestStore u
  noRequests = RequestStore Map.empty
  
  -- | Adds a 'BlockedFetch' to a 'RequestStore'.
  addRequest
    :: forall u r. (DataSource u r)
    => BlockedFetch r -> BlockedFetchInternal -> RequestStore u -> RequestStore u
  addRequest bf bfi (RequestStore m) =
    RequestStore $ Map.insertWith combine ty (BlockedFetches [bf] [bfi]) m
   where
    combine :: BlockedFetches u -> BlockedFetches u -> BlockedFetches u
    combine _ (BlockedFetches bfs bfis)
      | typeOf1 (getR bfs) == ty = BlockedFetches (unsafeCoerce bf:bfs) (bfi:bfis)
      | otherwise                = error "RequestStore.insert"
           -- 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 negligible.
  
    -- a type conversion only, so we can get the type of the reqeusts from
    -- the list of BlockedFetch.
    getR :: [BlockedFetch r1] -> r1 a
    getR _ = undefined
  
    -- The TypeRep of requests for this data source
    ty :: TypeRep
    !ty = typeOf1 (undefined :: r a)
  
  -- | Retrieves the whole contents of the 'RequestStore'.
  contents :: RequestStore u -> [BlockedFetches u]
  contents (RequestStore m) = Map.elems m
  
  getSize :: RequestStore u -> Int
  getSize (RequestStore m) = Map.size m
  
  -- A counter to keep track of outgone requests. Entries are added to this
  -- map as we send requests to datasources, and removed as these fetches
  -- are completed.
  -- This is a 2 level map: the 1st level stores requests for a particular
  -- datasource, the 2nd level stores count of requests per type.
  newtype ReqCountMap = ReqCountMap (Map Text (Map TypeRep Int))
    deriving (Show)
  
  emptyReqCounts :: ReqCountMap
  emptyReqCounts = ReqCountMap Map.empty
  
  addToCountMap
    :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
    => Proxy r
    -> Int -- type and number of requests
    -> ReqCountMap
    -> ReqCountMap
  addToCountMap = updateCountMap (+)
  
  subFromCountMap
    :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
    => Proxy r
    -> Int -- type and number of requests
    -> ReqCountMap
    -> ReqCountMap
  subFromCountMap = updateCountMap (-)
  
  updateCountMap
    :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
    => (Int -> Int -> Int)
    -> Proxy r
    -> Int -- type and number of requests
    -> ReqCountMap
    -> ReqCountMap
  updateCountMap op p n (ReqCountMap m) = ReqCountMap $ Map.insertWith
    (flip (Map.unionWith op)) -- flip is important as "op" is not commutative
    (dataSourceName p) (Map.singleton ty n)
    m
    where
      -- The TypeRep of requests for this data source
      -- The way this is implemented, all elements in the 2nd level map will be
      -- mapped to the same key, as all requests to a datasource have the same
      -- "type". It will be more beneficial to be able to instead map requests
      -- to their names (ie, data constructor) - but there's no cheap way of doing
      -- that.
      ty :: TypeRep
      !ty = typeOf1 (undefined :: r a)
  
  -- Filter all keys with 0 fetches. Since ReqCountMap is a 2-level map, we need
  -- nested filter operations.
  filterRCMap :: ReqCountMap -> ReqCountMap
  filterRCMap (ReqCountMap m) = ReqCountMap $
    Map.filter ((> 0) . Map.size) (Map.filter (> 0) <$> m)
  
  -- Filters the ReqCountMap by default
  getMapFromRCMap :: ReqCountMap -> Map Text (Map TypeRep Int)
  getMapFromRCMap r
    | ReqCountMap m <- filterRCMap r = m
  
  getSummaryMapFromRCMap :: ReqCountMap -> HashMap.HashMap Text Int
  getSummaryMapFromRCMap (ReqCountMap m) = HashMap.fromList
    [ (k, s)
    | (k, v) <- Map.toList m
    , not $ Map.null v
    , let s = sum $ Map.elems v
    , s > 0
    ]
  