Haxl.hs

{-# LANGUAGE Trustworthy #-}

-- |
-- Module      : Effectful.Haxl
-- Copyright   : (c) 2026 Institute for Digital Autonomy
-- License     : EUPL-1.2
-- Maintainer  : IDA
--
-- <https://hackage.haskell.org/package/haxl Haxl> is a library for
-- concurrent, batched, and cached data fetching.
--
-- This library provides the 'Haxl' effect can initialise an 'Env', as well
-- as a lifted 'GenHaxl' monad that runs computation with a single @Env@.
--
-- Haxl can parallelise independent requests using the 'Applicative' instance,
-- but not the 'Monad' instance.
-- To benefit from parallelisation in @do@ blocks, enable the @ApplicativeDo@
-- language extension.
--
-- = Example usage
--
-- == Defining a data source
--
-- A data source describes one kind of request and how to fetch it. At
-- minimum this means a request type indexed by its result, and instances of
-- 'DataSourceName' and 'DataSource':
--
-- > data ExampleReq a where
-- >   CountAardvarks :: String -> ExampleReq Int
-- >
-- > deriving stock instance Eq (ExampleReq a)
-- > deriving stock instance Show (ExampleReq a)
-- > instance ShowP ExampleReq where showp = show
-- > instance Hashable (ExampleReq a) where
-- >   hashWithSalt s (CountAardvarks a) = hashWithSalt s a
-- >
-- > instance StateKey ExampleReq where
-- >   data State ExampleReq = ExampleState
-- >
-- > instance DataSourceName ExampleReq where
-- >   dataSourceName _ = "ExampleDataSource"
-- >
-- > instance DataSource u ExampleReq where
-- >   fetch _state _flags _user = SyncFetch $ mapM_ \(BlockedFetch (CountAardvarks str) result) ->
-- >     putSuccess result (length (filter (== 'a') str))
-- >
-- > countAardvarks :: String -> GenHaxl u w es Int
-- > countAardvarks = dataFetch . CountAardvarks
--
-- See the <https://hackage.haskell.org/package/haxl Haxl> package documentation
-- for the full 'DataSource' contract, including asynchronous fetching and batching
-- multiple requests per round.
--
-- == Running a computation
--
-- Build the 'StateStore' from every data source's 'State', then use
-- 'runHaxlWith' to discharge the 'Haxl' effect and 'haxl' to run a 'GenHaxl'
-- computation inside it:
--
-- > import Effectful
-- > import Effectful.Haxl
-- >
-- > main :: IO ()
-- > main = runEff . runHaxlWith initialState () defaultFlags $ do
-- >   n <- inject . haxl @() @() $ countAardvarks "abcabc"
-- >   liftIO $ print n
-- >   where
-- >     initialState = pure $ stateSet ExampleState stateEmpty
--
-- If you only have one data source and no user environment or writes, 'runHaxl'
-- is a shorter alternative to 'runHaxlWith'.
--
-- == Batching independent fetches
--
-- Requests combined applicatively (with @('<*>')@, @('+')@ via the 'Num'
-- instance, 'traverse', etc.) run in parallel in a single round trip:
--
-- > -- One round: both fetches are independent.
-- > total :: GenHaxl u w es Int
-- > total = countAardvarks "abc" + countAardvarks "def"
--
-- Requests sequenced with @('>>=')@ each wait for the previous result and so
-- take a separate round:
--
-- > -- Two rounds: the second fetch depends on the first result.
-- > chained :: GenHaxl u w es Int
-- > chained = countAardvarks "abc" >>= \n -> countAardvarks (replicate n 'a')
--
-- These functions can also be expressed with @ApplicativeDo@:
--
-- > -- One round with ApplicativeDo, two rounds otherwise
-- > total :: GenHaxl u w es Int
-- > total = do
-- >     a <- countAardvarks "abc"
-- >     b <- countAardvarks "def"
-- >     pure (a + b)
-- >
-- > -- Two rounds even with ApplicativeDo
-- > chained :: GenHaxl u w es Int
-- > chained = do
-- >     n <- countAardvarks "abc"
-- >     countAardvarks (replicate n 'a')
--
-- Use 'haxlWithStats' to inspect 'Stats' and confirm how many rounds and
-- fetches a computation actually used:
--
-- > (_, totalStats) <- inject $ haxlWithStats @() @() total
-- > numFetches totalStats `shouldBe` 1 -- one round: 1 fetch batch of size 2
-- >
-- > (_, chainedStats) <- inject $ haxlWithStats @() @() chained
-- > numFetches chainedStats `shouldBe` 2 -- two rounds: 1 fetch each
--
-- == Caching and memoization
--
-- Haxl automatically caches each unique request within a run, so fetching
-- the same request twice only issues one fetch. 'cacheRequest' seeds the
-- cache with a known result to skip a fetch entirely, and 'cachedComputation'
-- \/ 'memo' deduplicate an entire sub-computation (not just a single request)
-- across multiple uses:
--
-- > deduped :: GenHaxl u w es Int
-- > deduped = memo (CountAardvarks "ababa") $ do
-- >   -- runs only once even if `deduped` itself is used more than once
-- >   countAardvarks "ababa"
--
-- == Handling exceptions
--
-- 'GenHaxl' has its own 'catch' and 'try', distinct from the ones in
-- @Effectful.Exception@ or @Control.Exception@.
-- An exception thrown inside a 'GenHaxl' computation will not stop execution
-- at the point of the call, but only when its round of concurrent, batched
-- fetches is dispatched and completed.
-- Demanding the corresponding result will re-raise the exception.
--
-- > safeCount :: GenHaxl u w es (Either SomeException Int)
-- > safeCount = try (countAardvarks "BANG")
module Effectful.Haxl
    ( -- * Effect
      Haxl
    , runHaxlWith
    , runHaxl
    , haxl
    , haxlWithStats
    , haxlWithWrites

      -- * Env
    , initEnv
    , env
    , withEnv
    , GenHaxl (..)
    , liftGenHaxl
    , dataFetch
    , uncachedRequest
    , cacheRequest
    , dupableCacheRequest
    , cacheResult
    , cacheResultWithShow
    , cachedComputation
    , preCacheComputation
    , memoize
    , memoize1
    , memoize2
    , memo
    , memoFingerprint

      -- * Exceptions
    , throw
    , catch
    , catchIf
    , try
    , rethrowAsyncExceptions
    , tryWithRethrow

      -- * Dumping the cache
    , dumpCacheAsHaskell
    , dumpCacheAsHaskellFn

      -- * Re-exports
    , module Haxl.Core
    , module Haxl.Core.DataSource
    , module Haxl.Core.Exception
    , module Haxl.Core.Flags
    , module Haxl.Core.Monad
    , module Haxl.Core.ShowP
    , module Haxl.Core.StateStore
    , module Haxl.Core.Stats
    )
where

import Data.Functor ((<&>))
import Data.Hashable (Hashable)
import Data.IORef (readIORef)
import Data.Kind (Type)
import Data.Typeable (Typeable)
import Effectful
import Effectful.Dispatch.Static
import Effectful.Exception (Exception, SomeException)
import Effectful.Exception qualified as Effectful
import Haxl.Core (Env (..), Request)
import Haxl.Core qualified as Haxl
import Haxl.Core.DataSource
import Haxl.Core.Exception hiding (MonadFail (..), rethrowAsyncExceptions, tryWithRethrow)
import Haxl.Core.Fetch (ShowReq)
import Haxl.Core.Flags
import Haxl.Core.Memo (MemoFingerprintKey)
import Haxl.Core.Monad (Result (..))
import Haxl.Core.Monad qualified as Env
import Haxl.Core.Monad qualified as Haxl
import Haxl.Core.ShowP
import Haxl.Core.StateStore
import Haxl.Core.Stats
import Haxl.Prelude (IfThenElse (..))
import Prelude

data Haxl (u :: Type) (w :: Type) :: Effect

type instance DispatchOf (Haxl _ _) = 'Static 'WithSideEffects

data instance StaticRep (Haxl u _) = Haxl
    { stateStore :: !StateStore
    , userEnv :: !u
    , flags :: !Flags
    }

-- | Discharge the 'Haxl' effect for the duration of an action.
--
-- Use this when you need a user environment (@u@), collected writes (@w@),
-- non-default 'Flags', or non-trivial state store initialisation.
-- If none of that applies, 'runHaxl' is a shorter alternative.
runHaxlWith
    :: forall u w es a
     . (IOE :> es)
    => Eff es StateStore
    -> u
    -> Flags
    -> Eff (Haxl u w ': es) a
    -> Eff es a
runHaxlWith initialiseStore userEnv flags eff = do
    stateStore <- initialiseStore
    evalStaticRep Haxl{..} eff

-- | Discharge the 'Haxl' effect for a single data source, with no user
-- environment, no writes, and default 'Flags'.
--
-- This covers the common case of one data source and no need to tune
-- fetch reporting or batching behaviour.
runHaxl
    :: (StateKey k, IOE :> es)
    => State k
    -> Eff (Haxl () () ': es) a
    -> Eff es a
runHaxl s = runHaxlWith (pure $ stateSet s stateEmpty) () defaultFlags

-- | Run a 'GenHaxl' computation to completion inside the 'Haxl' effect with a fresh 'Env'.
haxl :: forall u w es a. (Monoid w) => GenHaxl u w es a -> Eff (Haxl u w ': es) a
haxl g = do
    env <- initEnv
    inject $ unsafeConcUnliftIO Persistent Unlimited \unlift ->
        Haxl.runHaxl env $ Haxl.GenHaxl (unlift . unHaxl g)

-- | Like 'haxl', but also return everything written via the 'w' monoid during the run.
haxlWithWrites :: forall u w es a. (Monoid w) => GenHaxl u w es a -> Eff (Haxl u w ': es) (a, w)
haxlWithWrites g = do
    env <- initEnv
    inject $ unsafeConcUnliftIO Persistent Unlimited \unlift ->
        Haxl.runHaxlWithWrites env $ Haxl.GenHaxl (unlift . unHaxl g)

-- | Like 'haxl', but also return the 'Stats' collected during the run.
--
-- Use this when you want to inspect batching behaviour,
-- e.g. asserting that a computation fetches in one round rather than many.
haxlWithStats :: forall u w es a. (Monoid w) => GenHaxl u w es a -> Eff (Haxl u w ': es) (a, Stats)
haxlWithStats g = do
    env <- initEnv
    a <- inject $ unsafeConcUnliftIO Persistent Unlimited \unlift ->
        Haxl.runHaxl env $ Haxl.GenHaxl (unlift . unHaxl g)
    stats <- unsafeEff_ . readIORef $ statsRef env
    pure (a, stats)

-- | Build a fresh 'Env' from the current 'Haxl' effect's state store, user
-- environment, and 'Flags'.
--
-- Each 'Env' has its own request cache, so an 'Env' obtained here starts
-- with an empty cache even if other 'Env's have already been used within
-- the same 'Haxl' effect scope.
-- You normally don't need this directly, as 'haxl' and @haxlWith*@ call it for you.
initEnv :: forall u w es. (Monoid w, Haxl u w :> es) => Eff es (Env u w)
initEnv = do
    Haxl{..} <- getStaticRep @(Haxl u w)
    unsafeEff_ $ Haxl.initEnv stateStore userEnv <&> \env -> env{Env.flags}

-- | The monad in which data-fetching computations are written.
--
-- Note that 'GenHaxl' is not expressed in terms of 'Eff'.
-- Its 'Applicative' instance batches independent fetches combined with @('<*>')@ into a single round trip.
newtype GenHaxl u w es a = GenHaxl {unHaxl :: Env u w -> Eff es (Haxl.Result u w a)}

instance Functor (GenHaxl u w es) where
    fmap f (GenHaxl m) = GenHaxl $ (fmap . fmap . fmap) f m

instance Applicative (GenHaxl u w es) where
    pure = GenHaxl . const . pure . Done
    f <*> a = GenHaxl $ \env -> do
        unsafeConcUnliftIO Ephemeral Unlimited $ \unlift ->
            Haxl.unHaxl
                ( Haxl.GenHaxl (unlift . unHaxl f)
                    <*> Haxl.GenHaxl (unlift . unHaxl a)
                )
                env

instance Monad (GenHaxl u w es) where
    return = pure
    m >>= k = GenHaxl $ \env -> do
        unsafeConcUnliftIO Ephemeral Unlimited $ \unlift ->
            Haxl.unHaxl
                ( Haxl.GenHaxl (unlift . unHaxl m)
                    >>= \a -> Haxl.GenHaxl (unlift . unHaxl (k a))
                )
                env

instance MonadFail (GenHaxl u w es) where
    fail = liftGenHaxl . fail

instance (Semigroup a) => Semigroup (GenHaxl u w es a) where
    (<>) = liftA2 (<>)

instance (Monoid a) => Monoid (GenHaxl u w es a) where
    mempty = pure mempty

instance IfThenElse (GenHaxl u w es Bool) (GenHaxl u w es a) where
    ifThenElse fb t e = do
        b <- fb
        if b then t else e

instance (Num a) => Num (GenHaxl u w es a) where
    (+) = liftA2 (+)
    (-) = liftA2 (-)
    (*) = liftA2 (*)
    fromInteger = pure . fromInteger
    abs = fmap abs
    signum = fmap signum
    negate = fmap negate

instance (Fractional a) => Fractional (GenHaxl u w es a) where
    (/) = liftA2 (/)
    recip = fmap recip
    fromRational = return . fromRational

-- | Lift a plain 'Haxl.GenHaxl' computation into 'GenHaxl'.
--
-- This is an escape hatch for the rare cases where you need to work directly
-- with 'Haxl.GenHaxl'; for example when you need precise control over
-- Haxl's fetch scheduling that the 'Eff' wrapper does not preserve, or when
-- integrating with third-party code that returns 'Haxl.GenHaxl' values
-- directly.
--
-- Warning: 'Haxl.GenHaxl' exposes functions such as 'Haxl.unsafeLiftIO'
-- that allow arbitrary 'IO' to be injected without it being reflected in @es@.
liftGenHaxl :: Haxl.GenHaxl u w a -> GenHaxl u w es a
liftGenHaxl = GenHaxl . (unsafeEff_ .) . Haxl.unHaxl

unliftGenHaxl :: GenHaxl u w es a -> Eff es (Haxl.GenHaxl u w a)
unliftGenHaxl g =
    unsafeConcUnliftIO Ephemeral Unlimited $ pure . Haxl.GenHaxl . (. unHaxl g)

unliftGenHaxl1 :: (a -> GenHaxl u w es b) -> Eff es (a -> Haxl.GenHaxl u w b)
unliftGenHaxl1 f =
    unsafeConcUnliftIO Ephemeral Unlimited \unlift -> pure \a ->
        Haxl.GenHaxl $ unlift . unHaxl (f a)

unliftGenHaxl2 :: (a -> b -> GenHaxl u w es c) -> Eff es (a -> b -> Haxl.GenHaxl u w c)
unliftGenHaxl2 f =
    unsafeConcUnliftIO Ephemeral Unlimited \unlift -> pure \a b ->
        Haxl.GenHaxl $ unlift . unHaxl (f a b)

mapGenHaxl :: (Haxl.GenHaxl u w a -> Haxl.GenHaxl u w b) -> GenHaxl u w es a -> GenHaxl u w es b
mapGenHaxl f g = GenHaxl \env -> do
    g <- unliftGenHaxl g
    unsafeEff_ $ Haxl.unHaxl (f g) env

mapGenHaxl1
    :: ((a -> Haxl.GenHaxl u w b) -> Haxl.GenHaxl u w c)
    -> (a -> GenHaxl u w es b)
    -> GenHaxl u w es c
mapGenHaxl1 f g = GenHaxl \env -> do
    g <- unliftGenHaxl1 g
    unsafeEff_ $ Haxl.unHaxl (f g) env

mapGenHaxl2
    :: ((a -> b -> Haxl.GenHaxl u w c) -> Haxl.GenHaxl u w d)
    -> (a -> b -> GenHaxl u w es c)
    -> GenHaxl u w es d
mapGenHaxl2 f g = GenHaxl \env -> do
    g <- unliftGenHaxl2 g
    unsafeEff_ $ Haxl.unHaxl (f g) env

mapGenHaxlCatch
    :: (Haxl.GenHaxl u w a -> (e -> Haxl.GenHaxl u w a) -> Haxl.GenHaxl u w a)
    -> GenHaxl u w es a
    -> (e -> GenHaxl u w es a)
    -> GenHaxl u w es a
mapGenHaxlCatch f g h = GenHaxl \env -> do
    g <- unliftGenHaxl g
    h <- unliftGenHaxl1 h
    unsafeEff_ $ Haxl.unHaxl (f g h) env

-- | Lifted 'Haxl.env'
env :: (Env u w -> a) -> GenHaxl u w es a
env = liftGenHaxl . Haxl.env

-- | Lifted 'Haxl.withEnv'
withEnv :: Env u w -> GenHaxl u w es a -> GenHaxl u w es a
withEnv = mapGenHaxl . Haxl.withEnv

-- | Lifted 'Haxl.dataFetch'
dataFetch :: (DataSource u r, Request r a) => r a -> GenHaxl u w es a
dataFetch = liftGenHaxl . Haxl.dataFetch

-- | Lifted 'Haxl.uncachedRequest'
uncachedRequest :: (DataSource u r, Request r a) => r a -> GenHaxl u w es a
uncachedRequest = liftGenHaxl . Haxl.uncachedRequest

-- | Lifted 'Haxl.cacheRequest'
cacheRequest :: (Request req a) => req a -> Either SomeException a -> GenHaxl u w es ()
cacheRequest = (liftGenHaxl .) . Haxl.cacheRequest

-- | Lifted 'Haxl.dupableCacheRequest'
dupableCacheRequest :: (Request req a) => req a -> Either SomeException a -> GenHaxl u w es ()
dupableCacheRequest = (liftGenHaxl .) . Haxl.dupableCacheRequest

-- | Lifted 'Haxl.cacheResult'
cacheResult :: (Request r a) => r a -> Eff es a -> GenHaxl u w es a
cacheResult r eff =
    GenHaxl \env ->
        unsafeConcUnliftIO Ephemeral Unlimited \unlift ->
            Haxl.unHaxl (Haxl.cacheResult r (unlift eff)) env

-- | Lifted 'Haxl.cacheResultWithShow'
cacheResultWithShow
    :: (Hashable (r a), Typeable (r a))
    => ShowReq r a
    -> r a
    -> Eff es a
    -> GenHaxl u w es a
cacheResultWithShow showReq req eff =
    GenHaxl \env ->
        unsafeConcUnliftIO Ephemeral Unlimited \unlift ->
            Haxl.unHaxl (Haxl.cacheResultWithShow showReq req (unlift eff)) env

-- | Lifted 'Haxl.cachedComputation'
cachedComputation
    :: (Hashable (req a), Typeable (req a), Monoid w)
    => req a
    -> GenHaxl u w es a
    -> GenHaxl u w es a
cachedComputation = mapGenHaxl . Haxl.cachedComputation

-- | Lifted 'Haxl.preCacheComputation'
preCacheComputation
    :: (Hashable (req a), Typeable (req a), Monoid w)
    => req a
    -> GenHaxl u w es a
    -> GenHaxl u w es a
preCacheComputation = mapGenHaxl . Haxl.preCacheComputation

-- | Lifted 'Haxl.memoize'
memoize :: (Monoid w) => GenHaxl u w es a -> GenHaxl u w es (GenHaxl u w es a)
memoize = fmap liftGenHaxl . mapGenHaxl Haxl.memoize

-- | Lifted 'Haxl.memoize1'
memoize1
    :: (Hashable a, Monoid w)
    => (a -> GenHaxl u w es b)
    -> GenHaxl u w es (a -> GenHaxl u w es b)
memoize1 = fmap (liftGenHaxl .) . mapGenHaxl1 Haxl.memoize1

-- | Lifted 'Haxl.memoize2'
memoize2
    :: (Hashable a, Hashable b, Monoid w)
    => (a -> b -> GenHaxl u w es c)
    -> GenHaxl u w es (a -> b -> GenHaxl u w es c)
memoize2 = fmap ((liftGenHaxl .) .) . mapGenHaxl2 Haxl.memoize2

-- | Lifted 'Haxl.memo'
memo
    :: ( Typeable a
       , Typeable k
       , Hashable k
       , Monoid w
       )
    => k
    -> GenHaxl u w es a
    -> GenHaxl u w es a
memo = mapGenHaxl . Haxl.memo

-- | Lifted 'Haxl.memoFingerprint'
memoFingerprint
    :: (Typeable a, Monoid w)
    => MemoFingerprintKey a
    -> GenHaxl u w es a
    -> GenHaxl u w es a
memoFingerprint = mapGenHaxl . Haxl.memoFingerprint

-- | Lifted 'Haxl.rethrowAsyncExceptions'
rethrowAsyncExceptions :: SomeException -> Eff es ()
rethrowAsyncExceptions = unsafeEff_ . Haxl.rethrowAsyncExceptions

tryWithRethrow :: Eff es a -> Eff es (Either SomeException a)
tryWithRethrow eff = (Right <$> eff) `Effectful.catch` \e -> rethrowAsyncExceptions e >> pure (Left e)

-- | Lifted 'Haxl.throw'
throw :: (Exception e) => e -> GenHaxl u w es a
throw = liftGenHaxl . Haxl.throw

-- | Lifted 'Haxl.catch'
catch :: (Exception e) => GenHaxl u w es a -> (e -> GenHaxl u w es a) -> GenHaxl u w es a
catch = mapGenHaxlCatch Haxl.catch

-- | Lifted 'Haxl.catchIf'
catchIf
    :: (Exception e)
    => (e -> Bool)
    -> GenHaxl u w es a
    -> (e -> GenHaxl u w es a)
    -> GenHaxl u w es a
catchIf = mapGenHaxlCatch . Haxl.catchIf

-- | Lifted 'Haxl.try'
try :: (Exception e) => GenHaxl u w es a -> GenHaxl u w es (Either e a)
try = mapGenHaxl Haxl.try

-- | Lifted 'Haxl.dumpCacheAsHaskell'
dumpCacheAsHaskell :: GenHaxl u w es String
dumpCacheAsHaskell = dumpCacheAsHaskellFn "loadCache" "GenHaxl u w es ()" "cacheRequest"

-- | Lifted 'Haxl.dumpCacheAsHaskellFn'
dumpCacheAsHaskellFn :: String -> String -> String -> GenHaxl u w es String
dumpCacheAsHaskellFn = ((liftGenHaxl .) .) . Haxl.dumpCacheAsHaskellFn