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 BangPatterns #-}
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RecordWildCards #-}
  
  -- |
  -- Memoization support. This module is provided for access to Haxl
  -- internals only; most users should import "Haxl.Core" instead.
  --
  module Haxl.Core.Memo
    (
      -- * Basic memoization
      cachedComputation
    , preCacheComputation
  
      -- * High-level memoization
    , memo
    , memoFingerprint
    , MemoFingerprintKey(..)
    , memoize, memoize1, memoize2
    , memoUnique
  
      -- * Local memoization
    , MemoVar
    , newMemo
    , newMemoWith
    , prepareMemo
    , runMemo
    ) where
  
  import Control.Exception as Exception hiding (throw)
  import Data.IORef
  import qualified Data.HashMap.Strict as HashMap
  import Data.Text (Text)
  import Data.Typeable
  import Data.Hashable
  import Data.Int
  import Data.Word
  
  import GHC.Prim (Addr#)
  
  import Haxl.Core.Exception
  import Haxl.Core.DataCache as DataCache
  import Haxl.Core.Flags
  import Haxl.Core.Monad
  import Haxl.Core.Stats
  import Haxl.Core.Profile
  import Haxl.Core.Util (trace_)
  
  -- -----------------------------------------------------------------------------
  -- Memoization
  
  -- | 'cachedComputation' memoizes a Haxl computation.  The key is a
  -- request.
  --
  -- /Note:/ These cached computations will /not/ be included in the output
  -- of 'dumpCacheAsHaskell'.
  --
  cachedComputation
     :: forall req u w a.
        ( Eq (req a)
        , Hashable (req a)
        , Typeable (req a)
        , Monoid w)
     => req a -> GenHaxl u w a -> GenHaxl u w a
  
  cachedComputation req haxl = GenHaxl $ \env@Env{..} -> do
    mbRes <- DataCache.lookup req memoCache
    case mbRes of
      Just (DataCacheItem ivar k) -> do
        ifProfiling flags $ do
          incrementMemoHitCounterFor env k True
        unHaxl (getIVarWithWrites ivar) env
      Nothing -> do
        ivar <- newIVar
        k <- nextCallId env
        -- no need to incremenetMemoHitCounter as execMemo will do it
        DataCache.insertNotShowable req (DataCacheItem ivar k) memoCache
        execMemoNowProfiled env haxl ivar k
  
  
  -- | Like 'cachedComputation', but fails if the cache is already
  -- populated.
  --
  -- Memoization can be (ab)used to "mock" a cached computation, by
  -- pre-populating the cache with an alternative implementation. In
  -- that case we don't want the operation to populate the cache to
  -- silently succeed if the cache is already populated.
  --
  preCacheComputation
    :: forall req u w a.
       ( Eq (req a)
       , Hashable (req a)
       , Typeable (req a)
       , Monoid w)
    => req a -> GenHaxl u w a -> GenHaxl u w a
  preCacheComputation req haxl = GenHaxl $ \env@Env{..} -> do
    mbRes <- DataCache.lookup req memoCache
    case mbRes of
      Just _ -> return $ Throw $ toException $ InvalidParameter
        "preCacheComputation: key is already cached"
      Nothing -> do
        ivar <- newIVar
        k <- nextCallId env
        DataCache.insertNotShowable req (DataCacheItem ivar k) memoCache
        execMemoNowProfiled env haxl ivar k
  
  -- -----------------------------------------------------------------------------
  -- Memoization
  
  newtype MemoVar u w a = MemoVar (IORef (MemoStatus u w a))
  
  data MemoStatus u w a
    = MemoEmpty
    | MemoReady (GenHaxl u w a) CallId
    | MemoRun {-# UNPACK #-} !(IVar u w a) {-# UNPACK #-} !CallId
  
  -- | Create a new @MemoVar@ for storing a memoized computation. The created
  -- @MemoVar@ is initially empty, not tied to any specific computation. Running
  -- this memo (with @runMemo@) without preparing it first (with @prepareMemo@)
  -- will result in an exception.
  newMemo :: GenHaxl u w (MemoVar u w a)
  newMemo = unsafeLiftIO $ MemoVar <$> newIORef MemoEmpty
  
  -- | Store a computation within a supplied @MemoVar@. Any memo stored within the
  -- @MemoVar@ already (regardless of completion) will be discarded, in favor of
  -- the supplied computation. A @MemoVar@ must be prepared before it is run.
  prepareMemo :: MemoVar u w a -> GenHaxl u w a -> GenHaxl u w ()
  prepareMemo (MemoVar memoRef) memoCmp
    = GenHaxl $ \env -> do
        k <- nextCallId env
        writeIORef memoRef (MemoReady memoCmp k)
        return (Done ())
  
  -- | Convenience function, combines @newMemo@ and @prepareMemo@.
  newMemoWith :: GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
  newMemoWith memoCmp = do
    memoVar <- newMemo
    prepareMemo memoVar memoCmp
    return memoVar
  
  -- | Continue the memoized computation within a given @MemoVar@.
  -- Notes:
  --
  --   1. If the memo contains a complete result, return that result.
  --   2. If the memo contains an in-progress computation, continue it as far as
  --      possible for this round.
  --   3. If the memo is empty (it was not prepared), throw an error.
  --
  -- For example, to memoize the computation @one@ given by:
  --
  -- > one :: Haxl Int
  -- > one = return 1
  --
  -- use:
  --
  -- > do
  -- >   oneMemo <- newMemoWith one
  -- >   let memoizedOne = runMemo aMemo one
  -- >   oneResult <- memoizedOne
  --
  -- To memoize mutually dependent computations such as in:
  --
  -- > h :: Haxl Int
  -- > h = do
  -- >   a <- f
  -- >   b <- g
  -- >   return (a + b)
  -- >  where
  -- >   f = return 42
  -- >   g = succ <$> f
  --
  -- without needing to reorder them, use:
  --
  -- > h :: Haxl Int
  -- > h = do
  -- >   fMemoRef <- newMemo
  -- >   gMemoRef <- newMemo
  -- >
  -- >   let f = runMemo fMemoRef
  -- >       g = runMemo gMemoRef
  -- >
  -- >   prepareMemo fMemoRef $ return 42
  -- >   prepareMemo gMemoRef $ succ <$> f
  -- >
  -- >   a <- f
  -- >   b <- g
  -- >   return (a + b)
  --
  runMemo :: Monoid w => MemoVar u w a -> GenHaxl u w a
  runMemo (MemoVar memoRef) = GenHaxl $ \env -> do
    stored <- readIORef memoRef
    case stored of
      -- Memo was not prepared first; throw an exception.
      MemoEmpty -> trace_ "MemoEmpty " $
        raise env $ CriticalError "Attempting to run empty memo."
      -- Memo has been prepared but not run yet
      MemoReady cont k -> trace_ "MemoReady" $ do
        ivar <- newIVar
        writeIORef memoRef (MemoRun ivar k)
        execMemoNowProfiled env cont ivar k
      -- The memo has already been run, get (or wait for) for the result
      MemoRun ivar k -> trace_ "MemoRun" $ do
        ifProfiling (flags env) $ do
          incrementMemoHitCounterFor env k True
        unHaxl (getIVarWithWrites ivar) env
  
  execMemoNowProfiled
    :: Monoid w
    => Env u w
    -> GenHaxl u w a
    -> IVar u w a
    -> CallId
    -> IO (Result u w a)
  execMemoNowProfiled envOuter cont ivar cid =
    if not $ testReportFlag ReportProfiling $ report $ flags envOuter
    then execMemoNow envOuter cont ivar
    else do
      incrementMemoHitCounterFor envOuter cid False
      unHaxl
        (collectMemoData 0 $ GenHaxl $ \e -> execMemoNow e cont ivar)
        envOuter
    where
      addStats :: Env u w -> Int64 -> IO ()
      addStats env acc = modifyIORef' (statsRef env) $ \(Stats s) ->
        Stats (MemoCall cid acc : s)
      collectMemoData :: Int64 -> GenHaxl u w a -> GenHaxl u w a
      collectMemoData acc f = GenHaxl $ \env -> do
        a0 <- getAllocationCounter
        r <- unHaxl f env{memoKey=cid}
        a1 <- getAllocationCounter
        let newTotal = acc + (a0 - a1)
        ret <- case r of
          Done a -> do addStats env newTotal; return (Done a)
          Throw e -> do addStats env newTotal; return (Throw e)
          Blocked ivar k ->
            return (Blocked ivar (Cont (collectMemoData newTotal (toHaxl k))))
        setAllocationCounter a1
        return ret
  
  execMemoNow :: Monoid w => Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
  execMemoNow env cont ivar = do
    wlogs <- newIORef mempty
    let
      !menv = env { writeLogsRef = wlogs }
      -- use an env with empty writes, so we can memoize the extra
      -- writes done as part of 'cont'
    r <- Exception.try $ unHaxl cont menv
  
    case r of
      Left e -> trace_ ("execMemoNow: Left " ++ show e) $ do
        rethrowAsyncExceptions e
        putIVar ivar (ThrowIO e) env
        throwIO e
      Right (Done a) -> trace_ "execMemoNow: Done" $ do
        wt <- readIORef wlogs
        putIVar ivar (Ok a (Just wt)) env
        modifyIORef' (writeLogsRef env) (<> wt)
        return (Done a)
      Right (Throw ex) -> trace_ ("execMemoNow: Throw" ++ show ex) $ do
        wt <- readIORef wlogs
        putIVar ivar (ThrowHaxl ex (Just wt)) env
        modifyIORef' (writeLogsRef env) (<> wt)
        return (Throw ex)
      Right (Blocked ivar' cont) -> trace_ "execMemoNow: Blocked" $ do
        -- We "block" this memoized computation in the new environment 'menv', so
        -- that when it finishes, we can store all the write logs from the env
        -- in the IVar.
        addJob menv (toHaxl cont) ivar ivar'
        -- Now we call @getIVarWithWrites@ to populate the writes in the original
        -- environment 'env'.
        return (Blocked ivar (Cont (getIVarWithWrites ivar)))
  
  -- -----------------------------------------------------------------------------
  -- 1-ary and 2-ary memo functions
  
  newtype MemoVar1 u w a b = MemoVar1 (IORef (MemoStatus1 u w a b))
  newtype MemoVar2 u w a b c = MemoVar2 (IORef (MemoStatus2 u w a b c))
  
  data MemoStatus1 u w a b
    = MemoEmpty1
    | MemoTbl1 (a -> GenHaxl u w b) (HashMap.HashMap a (MemoVar u w b))
  
  data MemoStatus2 u w a b c
    = MemoEmpty2
    | MemoTbl2
        (a -> b -> GenHaxl u w c)
        (HashMap.HashMap a (HashMap.HashMap b (MemoVar u w c)))
  
  newMemo1 :: GenHaxl u w (MemoVar1 u w a b)
  newMemo1 = unsafeLiftIO $ MemoVar1 <$> newIORef MemoEmpty1
  
  newMemoWith1 :: (a -> GenHaxl u w b) -> GenHaxl u w (MemoVar1 u w a b)
  newMemoWith1 f = newMemo1 >>= \r -> prepareMemo1 r f >> return r
  
  prepareMemo1 :: MemoVar1 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
  prepareMemo1 (MemoVar1 r) f
    = unsafeLiftIO $ writeIORef r (MemoTbl1 f HashMap.empty)
  
  runMemo1 :: (Eq a, Hashable a, Monoid w) => MemoVar1 u w a b -> a -> GenHaxl u w b
  runMemo1 (MemoVar1 r) k = unsafeLiftIO (readIORef r) >>= \case
    MemoEmpty1 -> throw $ CriticalError "Attempting to run empty memo."
    MemoTbl1 f h -> case HashMap.lookup k h of
      Nothing -> do
        x <- newMemoWith (f k)
        unsafeLiftIO $ writeIORef r (MemoTbl1 f (HashMap.insert k x h))
        runMemo x
      Just v -> runMemo v
  
  newMemo2 :: GenHaxl u w (MemoVar2 u w a b c)
  newMemo2 = unsafeLiftIO $ MemoVar2 <$> newIORef MemoEmpty2
  
  newMemoWith2 :: (a -> b -> GenHaxl u w c) -> GenHaxl u w (MemoVar2 u w a b c)
  newMemoWith2 f = newMemo2 >>= \r -> prepareMemo2 r f >> return r
  
  prepareMemo2 :: MemoVar2 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w ()
  prepareMemo2 (MemoVar2 r) f
    = unsafeLiftIO $ writeIORef r (MemoTbl2 f HashMap.empty)
  
  runMemo2 :: (Eq a, Hashable a, Eq b, Hashable b, Monoid w)
           => MemoVar2 u w a b c
           -> a -> b -> GenHaxl u w c
  runMemo2 (MemoVar2 r) k1 k2 = unsafeLiftIO (readIORef r) >>= \case
    MemoEmpty2 -> throw $ CriticalError "Attempting to run empty memo."
    MemoTbl2 f h1 -> case HashMap.lookup k1 h1 of
      Nothing -> do
        v <- newMemoWith (f k1 k2)
        unsafeLiftIO $ writeIORef r
          (MemoTbl2 f (HashMap.insert k1 (HashMap.singleton k2 v) h1))
        runMemo v
      Just h2 -> case HashMap.lookup k2 h2 of
        Nothing -> do
          v <- newMemoWith (f k1 k2)
          unsafeLiftIO $ writeIORef r
            (MemoTbl2 f (HashMap.insert k1 (HashMap.insert k2 v h2) h1))
          runMemo v
        Just v -> runMemo v
  
  -- -----------------------------------------------------------------------------
  -- A key type that can be used for memoizing computations by a Text key
  
  -- | Memoize a computation using an arbitrary key.  The result will be
  -- calculated once; the second and subsequent time it will be returned
  -- immediately.  It is the caller's responsibility to ensure that for
  -- every two calls @memo key haxl@, if they have the same @key@ then
  -- they compute the same result.
  memo
    :: (Typeable a, Typeable k, Hashable k, Eq k, Monoid w)
    => k -> GenHaxl u w a -> GenHaxl u w a
  memo key = cachedComputation (MemoKey key)
  
  {-# RULES
  "memo/Text" memo = memoText :: (Typeable a, Monoid w) =>
              Text -> GenHaxl u w a -> GenHaxl u w a
   #-}
  
  {-# NOINLINE memo #-}
  
  -- | Memoize a computation using its location and a Fingerprint. This ensures
  -- uniqueness across computations.
  memoUnique
    :: (Typeable a, Typeable k, Hashable k, Eq k, Monoid w)
    => MemoFingerprintKey a -> Text -> k -> GenHaxl u w a -> GenHaxl u w a
  memoUnique fp label key = withLabel label . memo (fp, key)
  
  {-# NOINLINE memoUnique #-}
  
  data MemoKey k a where
    MemoKey :: (Typeable k, Hashable k, Eq k) => k -> MemoKey k a
    deriving Typeable
  
  deriving instance Eq (MemoKey k a)
  
  instance Hashable (MemoKey k a) where
    hashWithSalt s (MemoKey t) = hashWithSalt s t
  
  -- An optimised memo key for Text keys.  This is used automatically
  -- when the key is Text, due to the RULES pragma above.
  
  data MemoTextKey a where
    MemoText :: Text -> MemoTextKey a
    deriving Typeable
  
  deriving instance Eq (MemoTextKey a)
  
  instance Hashable (MemoTextKey a) where
    hashWithSalt s (MemoText t) = hashWithSalt s t
  
  memoText :: (Typeable a, Monoid w) => Text -> GenHaxl u w a -> GenHaxl u w a
  memoText key = withLabel key . cachedComputation (MemoText key)
  
  -- | A memo key derived from a 128-bit MD5 hash.  Do not use this directly,
  -- it is for use by automatically-generated memoization.
  data MemoFingerprintKey a where
    MemoFingerprintKey
      :: {-# UNPACK #-} !Word64
      -> {-# UNPACK #-} !Word64
      -> Addr# -> Addr#
      -> MemoFingerprintKey a
    deriving Typeable
  
  deriving instance Eq (MemoFingerprintKey a)
  
  instance Hashable (MemoFingerprintKey a) where
    hashWithSalt s (MemoFingerprintKey x _ _ _) =
      hashWithSalt s (fromIntegral x :: Int)
  
  -- This is optimised for cheap call sites: when we have a call
  --
  --   memoFingerprint (MemoFingerprintKey 1234 5678 "module"# "name"#) e
  --
  -- then the MemoFingerprintKey constructor will be statically
  -- allocated (with two 64-bit fields and pointers to cstrings for the names),
  -- and shared by all calls to memo. So the memo call will not allocate,
  -- unlike memoText.
  --
  {-# NOINLINE memoFingerprint #-}
  memoFingerprint
    :: (Typeable a, Monoid w) => MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
  memoFingerprint key@(MemoFingerprintKey _ _ mnPtr nPtr) =
    withFingerprintLabel mnPtr nPtr . cachedComputation key
  
  -- * Generic memoization machinery.
  
  -- | Transform a Haxl computation into a memoized version of itself.
  --
  -- Given a Haxl computation, @memoize@ creates a version which stores its result
  -- in a @MemoVar@ (which @memoize@ creates), and returns the stored result on
  -- subsequent invocations. This permits the creation of local memos, whose
  -- lifetimes are scoped to the current function, rather than the entire request.
  memoize :: Monoid w => GenHaxl u w a -> GenHaxl u w (GenHaxl u w a)
  memoize a = runMemo <$> newMemoWith a
  
  -- | Transform a 1-argument function returning a Haxl computation into a
  -- memoized version of itself.
  --
  -- Given a function @f@ of type @a -> GenHaxl u w b@, @memoize1@ creates a version
  -- which memoizes the results of @f@ in a table keyed by its argument, and
  -- returns stored results on subsequent invocations with the same argument.
  --
  -- e.g.:
  --
  -- > allFriends :: [Int] -> GenHaxl u w [Int]
  -- > allFriends ids = do
  -- >   memoizedFriendsOf <- memoize1 friendsOf
  -- >   concat <$> mapM memoizeFriendsOf ids
  --
  -- The above implementation will not invoke the underlying @friendsOf@
  -- repeatedly for duplicate values in @ids@.
  memoize1 :: (Eq a, Hashable a, Monoid w)
           => (a -> GenHaxl u w b)
           -> GenHaxl u w (a -> GenHaxl u w b)
  memoize1 f = runMemo1 <$> newMemoWith1 f
  
  -- | Transform a 2-argument function returning a Haxl computation, into a
  -- memoized version of itself.
  --
  -- The 2-ary version of @memoize1@, see its documentation for details.
  memoize2 :: (Eq a, Hashable a, Eq b, Hashable b, Monoid w)
           => (a -> b -> GenHaxl u w c)
           -> GenHaxl u w (a -> b -> GenHaxl u w c)
  memoize2 f = runMemo2 <$> newMemoWith2 f
  