packages feed

fractal-layer-0.1.0.0: src/Fractal/Layer.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- |
-- Module      : Fractal.Layer
-- Description : Composable, type-safe resource management and dependency injection
-- Stability   : experimental
--
-- = What is a Layer?
--
-- A @'Layer' m deps env@ describes one layer of an application: it
-- requires some services as input (@deps@) and produces some services
-- as output (@env@).  You can think of it as an effectful function
-- with managed resource lifetimes:
--
-- @
-- type Layer m deps env  ≈  deps -> m env   (+ automatic cleanup)
-- @
--
-- For example, a @Layer IO (Socket, Config) Database@ says: "give me a
-- @Socket@ and a @Config@, and I will build you a @Database@, cleaning
-- it up when you're done."
--
-- A layer with no dependencies uses @()@ for @deps@:
--
-- @
-- configLayer :: Layer IO () Config
-- @
--
-- This means @configLayer@ is a self-contained recipe for producing a
-- @Config@ value. It doesn't need anything from the outside world
-- other than @IO@.
--
-- = Why Layers?
--
-- == The Problem: God Records
--
-- Most Haskell applications centralise their dependencies in a single
-- record:
--
-- @
-- data App = App
--   { appConfig  :: Config
--   , appDB      :: Database
--   , appCache   :: Cache
--   , appWeb     :: WebServer
--   , appMetrics :: MetricsCollector
--   }
--
-- initApp :: IO App
-- initApp = do
--   cfg     <- loadConfig
--   db      <- connectDB cfg        -- what if this throws?
--   cache   <- connectCache cfg     -- do we close db?
--   web     <- startServer cfg db   -- what about cache?
--   metrics <- startMetrics cfg     -- cleanup order?
--   pure App {..}
-- @
--
-- This works for small programs. As the application grows, three
-- problems compound:
--
-- 1. __No local reasoning.__  Every function that accepts @App@ can
--    touch every field. Changing the cache implementation forces you
--    to audit every consumer of @App@, even those that never use the
--    cache, because the type doesn't tell you what is actually needed.
--
-- 2. __Fragile resource safety.__  If @connectCache@ throws, is @db@
--    closed?  If @startServer@ throws, is @cache@ released?  The
--    compiler gives you no help. Maintaining the cleanup order by hand,
--    and keeping it correct as the application evolves, is tedious and
--    error-prone. Mistakes show up as leaked connections and file
--    handles in production.
--
-- 3. __Hidden sharing.__  Two subsystems that both need a connection
--    pool either get two independent pools (wasteful, and potentially
--    incorrect if they need to share a transaction context) or share
--    one via convention. Nothing in the type system enforces or even
--    documents the sharing.
--
-- == The Solution
--
-- With layers, you break the god record into independent pieces, each
-- declaring exactly what it needs:
--
-- @
-- configLayer :: Layer IO () Config
-- dbLayer     :: Layer IO Config Database
-- cacheLayer  :: Layer IO Config Cache
-- webLayer    :: Layer IO (Config, Database) WebServer
-- @
--
-- Each layer is self-contained. You can read, test, and refactor
-- @dbLayer@ without looking at @cacheLayer@ or @webLayer@. The type
-- checker enforces that dependencies are satisfied at every
-- composition site. If you change the signature of @dbLayer@, every
-- composition that uses it produces a type error until it is updated.
--
-- = Properties of Layers
--
-- == Recipes for Creating Services
--
-- A layer describes /how/ to build a service from its dependencies.
-- @Layer IO Config Database@ is a recipe that says: "given a @Config@,
-- I know how to create a @Database@." The recipe includes both the
-- initialisation logic and any associated cleanup.
--
-- == An Alternative to Manual Initialisation
--
-- A layer plays the same role as a hand-written @initFoo :: Config -> IO Foo@
-- function, but with important extras: it tracks resource lifetimes, composes
-- with other layers via standard typeclasses, and participates in parallel
-- startup and fallback. You get the same "build a value from its inputs"
-- workflow, without any of the manual plumbing.
--
-- == Effectful and Resourceful
--
-- Layer construction is fully effectful. Layers can acquire resources
-- (open connections, start threads) and register finalisers that are
-- __guaranteed__ to run when the layer's scope ends, even under
-- exceptions or cancellation. You never write a manual cleanup chain.
--
-- For example, a @Database@ layer can open a connection pool during
-- construction and drain it on shutdown:
--
-- @
-- dbLayer :: Layer IO Config Database
-- dbLayer = do
--   cfg <- ask
--   resource (createPool cfg) drainPool
-- @
--
-- == Asynchronous
--
-- Layer construction is fully asynchronous and non-blocking. This
-- matters when initialisation involves network calls, waiting for
-- readiness checks, or other long-running work. You can construct
-- an entire application's dependency graph without blocking the main
-- thread on any single subsystem.
--
-- == Parallel
--
-- Independent layers are acquired in parallel automatically. When you
-- compose two layers with @&&&@, both sides are started concurrently
-- via @async@. For applications with many independent subsystems, this
-- can significantly reduce startup time compared to sequential
-- initialisation.
--
-- @
-- -- db and cache start concurrently, both using the same config
-- backend :: Layer IO Config (Database, Cache)
-- backend = dbLayer &&& cacheLayer
-- @
--
-- == Composable
--
-- Layers compose through standard Haskell typeclasses. You don't need
-- to learn a new composition API; if you know 'Category', 'Arrow',
-- and 'Alternative', you already know how to wire layers together:
--
-- [@>>>@]  __Sequential__: pipe the output of one layer as the input
--    to the next.
--
--    @configLayer >>> dbLayer  ::  Layer IO () Database@
--
-- [@&&&@]  __Parallel__: run two layers concurrently on the same
--    input, returning both results as a tuple.
--
--    @dbLayer &&& cacheLayer  ::  Layer IO Config (Database, Cache)@
--
-- [@***@]  __Split__: run two layers concurrently on different inputs.
--
-- [@\<|\>@]  __Fallback__: try the first layer; if it fails, clean up
--    its resources and try the second.
--
--    @remoteDB \<|\> localDB  ::  Layer IO Config Database@
--
-- [@+++@, @|||@]  __Choice__: route @Either@ inputs through different
--    layers.
--
-- A complete example wiring up an application:
--
-- @
-- appLayer :: Layer IO () (Config, (Database, WebServer))
-- appLayer = configLayer >>> (dbLayer &&& webLayer)
-- @
--
-- This reads as: "build a @Config@ from nothing, then use it to build
-- a @Database@ and a @WebServer@ in parallel." The compiler verifies
-- that the output type of @configLayer@ (@Config@) matches the input
-- type expected by @dbLayer@ and @webLayer@.
--
-- == Resilient via Fallback
--
-- Layer construction can be made resilient using 'Alternative'. If
-- the primary layer fails, @\<|\>@ cleans up any resources the failed
-- attempt acquired and tries an alternative:
--
-- @
-- databaseLayer :: Layer IO Config Database
-- databaseLayer = remotePostgres \<|\> localSQLite
-- @
--
-- The failed branch's resources (partially-opened connections, temp
-- files) are released before the fallback begins. This is safe by
-- construction; you do not need to write any cleanup code for the
-- failure path.
--
-- = Running Layers
--
-- Once you have composed your layer graph, you run it with 'withLayer':
--
-- @
-- main :: IO ()
-- main = withLayer () appLayer $ \\(cfg, (db, ws)) -> do
--   putStrLn "All services ready."
--   serveRequests ws
--   -- when this callback returns (or throws), all resources
--   -- are released in reverse acquisition order
-- @
--
-- 'withLayer' acquires every resource in the composition, passes the
-- resulting environment to your callback, and releases everything when
-- the callback finishes. This guarantee holds even if the callback
-- throws an exception.
--
-- For tests and REPL experiments where you don't need the resources to
-- outlive construction, 'runLayer' is a simpler alternative:
--
-- @
-- cfg <- runLayer () configLayer
-- @
--
-- = Services: Shared Singletons
--
-- Real applications have diamond dependencies. Both @webLayer@ and
-- @apiLayer@ need a connection pool, but you want exactly one pool,
-- not two.
--
-- 'Service' solves this. Wrap any layer with 'mkService' to mark it
-- as shared, and consume it with 'service':
--
-- @
-- poolService :: Service IO Config ConnectionPool
-- poolService = mkService $ do
--   cfg <- ask
--   resource (createPool cfg) drainPool
-- @
--
-- Now every consumer gets the same instance:
--
-- @
-- webLayer :: Layer IO Config WebServer
-- webLayer = do
--   pool <- service poolService   -- first call creates the pool
--   resource (startWeb pool) stopWeb
--
-- apiLayer :: Layer IO Config ApiServer
-- apiLayer = do
--   pool <- service poolService   -- second call reuses the same pool
--   resource (startApi pool) stopApi
--
-- -- one pool, two consumers, automatic cleanup
-- app :: Layer IO Config (WebServer, ApiServer)
-- app = webLayer &&& apiLayer
-- @
--
-- Key properties of services:
--
-- * __At-most-once initialisation.__ The first call to 'service'
--   within a composition creates the instance. Every subsequent call
--   returns the same value.
--
-- * __Thread-safe.__ If @webLayer@ and @apiLayer@ race to create the
--   pool (as they will under @&&&@), only one thread performs the
--   initialisation and the other blocks until it completes.
--
-- * __Failure caching.__ If initialisation fails, the failure is
--   cached. Later calls to 'service' will re-throw the same exception
--   rather than retrying.
--
-- * __Automatic cleanup.__ The service's finaliser runs when the
--   enclosing 'withLayer' scope ends, just like any other resource.
--
-- = Observability
--
-- Every resource acquisition, effect, service lookup, and composition
-- step fires callbacks on a 'LayerInterceptor'. You can plug in
-- logging, metrics, or distributed tracing without modifying any layer
-- code. The companion module "Fractal.Layer.Diagnostics" turns these
-- callbacks into a tree you can render to the terminal or export as
-- JSON.
--
-- @
-- main :: IO ()
-- main = withLayerAndInterceptor loggingInterceptor () appLayer $ \\env ->
--   runApp env
-- @
--
-- = Creating Layers
--
-- There are three ways to create a layer, depending on the lifecycle
-- of the value it produces:
--
-- ['effect']  For values that don't need cleanup: configuration,
--   computed values, or any action whose result can be garbage
--   collected normally.
--
-- ['resource']  For values with an acquire\/release lifecycle:
--   database connections, file handles, thread pools. The release
--   action is guaranteed to run.
--
-- ['bracketed']  For adapting existing @with@-style functions from
--   third-party libraries, e.g. @withSocketsDo@, @withConnection@.
--
-- In all three cases, dependencies are accessed via 'ask' from
-- 'MonadReader', not through function arguments:
--
-- @
-- configLayer :: Layer IO () Config
-- configLayer = effect loadConfig
--
-- dbLayer :: Layer IO Config Database
-- dbLayer = do
--   cfg <- ask
--   resource (connectDB cfg) closeDB
--
-- poolLayer :: Layer IO () ConnectionPool
-- poolLayer = bracketed withConnectionPool
-- @
--
-- = Design Influences
--
-- The layer model is inspired by ZIO's @ZLayer@
-- (<https://zio.dev/reference/contextual/zlayer>), adapted to
-- idiomatic Haskell using 'Arrow', 'Category', and 'Alternative'.
module Fractal.Layer
  ( -- * Core type
    Layer,

    -- * Building layers
    -- $building
    effect,
    resource,
    bracketed,

    -- * Running layers
    -- $running
    withLayer,
    runLayer,

    -- * Services (shared singletons)
    -- $services
    Service,
    mkService,
    service,
    uncached,

    -- * Composition helpers
    -- $composition
    composeLayer,
    zipLayer,
    mapLayer,

    -- * Interceptor-aware runners
    runLayerWithInterceptor,
    withLayerAndInterceptor,

    -- * Interceptors (re-exported from "Fractal.Layer.Interceptor")
    LayerInterceptor (..),
    nullInterceptor,
    combineInterceptors,

    -- * Exceptions
    EmptyLayer (..),

    -- * Re-exports
    MonadResource,
  )
where

import Control.Applicative
import Control.Arrow
import Control.Category
import Control.Monad
import Control.Monad.Reader.Class
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Resource (createInternalState, runInternalState, withInternalState)
import Control.Monad.Trans.Resource.Internal (ReleaseMap (..), stateCleanupChecked)
import Control.Selective
import Data.List.NonEmpty (NonEmpty(..))
import Data.Semigroup (Semigroup(..))
import Data.IntMap.Strict (mapKeysMonotonic)
import Data.Profunctor
import Data.Profunctor.Traversing
import qualified Data.Text as T
import Data.Time.Clock (getCurrentTime, diffUTCTime)
import Data.Typeable
import Fractal.Layer.Interceptor
import Fractal.Layer.Internal
import UnliftIO
import UnliftIO.Resource
import Prelude hiding ((.))

-- $building
--
-- Three primitives cover the full spectrum of resource lifecycles:
--
-- * 'effect': for values that don't need cleanup (config, computed
--   values). Analogous to @ZLayer.fromZIO@ or @ZLayer.succeed@ in ZIO.
-- * 'resource': for values with an explicit acquire\/release pair.
--   The finaliser is guaranteed to run even under exceptions.
--   Analogous to @ZLayer.scoped@ with @acquireRelease@.
-- * 'bracketed': for adapting existing @withFoo@-style bracket
--   functions from third-party libraries.
--
-- In all three cases, dependencies come from 'ask'
-- ('Control.Monad.Reader.Class.MonadReader'), not through the
-- primitive's function arguments. This keeps the creation API uniform
-- and lets you compose layers via 'Monad' @do@-notation.

-- $running
--
-- 'withLayer' is the primary entry point for production code. It
-- acquires every resource in the layer graph, passes the resulting
-- environment to your callback, and releases everything when the
-- callback returns or throws. Resources are released in reverse
-- acquisition order.
--
-- 'runLayer' is a simpler alternative for tests and REPL sessions. It
-- builds the layer, returns the output, and immediately releases all
-- resources. Any handles or connections in the output will be invalid
-- after 'runLayer' returns.

-- $services
--
-- Services solve the diamond-dependency problem: when two independent
-- layers both need the same expensive resource (a connection pool, a
-- thread manager), you want exactly one instance, not two.
--
-- Wrap any layer with 'mkService' to declare it as shared, then
-- consume it with 'service'. The first call initialises the service
-- and caches the result. Every subsequent call returns the cached
-- instance. Initialisation is thread-safe and failure is cached.
--
-- See the module introduction for a complete worked example.

-- $composition
--
-- Most of the time you compose layers via their typeclass instances
-- ('Category' for @>>>@, 'Arrow' for @&&&@/@***@, 'ArrowChoice' for
-- @+++@/@|||@, 'Alternative' for @\<|\>@). The functions below are the
-- named equivalents that implement those instances. They are exported
-- for cases where you need to pass them as first-class values or where
-- the operator syntax is ambiguous.


getOrCreateCachedService :: forall m deps env. (MonadUnliftIO m, Typeable env) => Service m deps env -> Layer m deps env
getOrCreateCachedService (Service m) = Layer $ \lenv deps -> do
  let rep = typeRep (Proxy @env)
      serviceName = T.pack (show rep)

  -- Fast path: check without write lock
  snapshot <- liftIO $ readMVar lenv.serviceStates
  case lookupTypeMap @env snapshot of
    Just (Initialized x) -> do
      lift $ onServiceReuse (interceptor lenv) serviceName rep
      pure x
    Just (Failed e) -> throwIO e
    _ -> join $ modifyMVar lenv.serviceStates $ \currentStates -> do
      case lookupTypeMap @env currentStates of
        Just (Initialized x) -> do
          lift $ onServiceReuse (interceptor lenv) serviceName rep
          pure (currentStates, pure x)
        Just (Failed e) -> pure (currentStates, throwIO e)
        Nothing -> do
          let ctx = OperationContext
                { operationName = serviceName
                , operationType = Just rep
                , operationMetadata = []
                }
          lift $ onServiceCreate (interceptor lenv) ctx

          eRes <- try $ build m lenv deps
          case eRes of
            Left e -> do
              let states' = insertTypeMap @env (Failed e) currentStates
              pure (states', throwIO e)
            Right x -> do
              let states' = insertTypeMap @env (Initialized x) currentStates
              pure (states', pure x)

-------------------------------------------------------------------------------
-- Construction helpers
-------------------------------------------------------------------------------

-- | Acquire a resource and register its finaliser.
--
-- Use 'resource' when you have a matching acquire\/release pair.  The
-- release action runs when the enclosing 'withLayer' scope ends,
-- whether normally or via exception, so you never need to write
-- manual cleanup code.
--
-- __Database connection:__
--
-- @
-- dbLayer :: Layer IO Config Database
-- dbLayer = do
--   cfg <- ask
--   resource (connectDB cfg) closeDB
-- @
--
-- __Thread pool:__
--
-- @
-- workerPool :: Layer IO PoolConfig ThreadPool
-- workerPool = do
--   cfg <- ask
--   resource (newThreadPool cfg.numWorkers) shutdownPool
-- @
--
-- __File handle (release guaranteed even on exception):__
--
-- @
-- logFile :: Layer IO FilePath Handle
-- logFile = do
--   path <- ask
--   resource (openFile path AppendMode) hClose
-- @
--
-- The acquire action runs inside 'ResourceT', so any exception during
-- acquisition prevents the finaliser from being registered (there is
-- nothing to clean up). Once acquisition succeeds, the finaliser is
-- registered and will run no matter what happens later.
resource ::
  forall m deps env. (MonadUnliftIO m, Typeable env) =>
  -- | Acquire resource
  m env ->
  -- | Release resource
  (env -> m ()) ->
  Layer m deps env
resource acq rel = Layer $ \lenv _deps -> do
  let rep = typeRep (Proxy @env)
      ctx = OperationContext
        { operationName = T.pack (show rep)
        , operationType = Just rep
        , operationMetadata = []
        }
  lift $ onResourceAcquire (interceptor lenv) ctx
  startTime <- liftIO getCurrentTime

  (_, env) <- allocateU
    (lift acq)
    (\x -> do
      lift $ onResourceRelease (interceptor lenv) (operationName ctx)
      lift $ rel x)

  endTime <- liftIO getCurrentTime
  let duration = diffUTCTime endTime startTime
  lift $ onResourceAcquireComplete (interceptor lenv) (operationName ctx) duration

  pure env
{-# INLINABLE resource #-}

-- | Lift a monadic action into a layer.  No finaliser is registered,
-- so this is the right choice for values that can be garbage collected
-- normally: configuration, computed values, or one-shot effects whose
-- result doesn't hold on to external resources.
--
-- __Loading configuration:__
--
-- @
-- configLayer :: Layer IO () Config
-- configLayer = effect $
--   Config \<$\> readEnvVar \"PORT\"
--          \<*\> readEnvVar \"HOST\"
-- @
--
-- __Pure computation (no IO needed):__
--
-- @
-- seedLayer :: Layer IO () Seed
-- seedLayer = effect (pure (Seed 42))
-- @
--
-- __Reading a file at startup:__
--
-- @
-- templatesLayer :: Layer IO Config Templates
-- templatesLayer = do
--   cfg <- ask
--   effect $ loadTemplates cfg.templateDir
-- @
--
-- If the value you are producing holds on to an external resource
-- (a file handle, a socket, a connection), use 'resource' instead
-- so the resource is cleaned up automatically.
effect ::
  forall m deps env. (MonadUnliftIO m, Typeable env) =>
  -- | Run effect to produce environment
  m env ->
  Layer m deps env
effect f = Layer $ \lenv _deps -> do
  let rep = typeRep (Proxy @env)
      ctx = OperationContext
        { operationName = T.pack (show rep)
        , operationType = Just rep
        , operationMetadata = []
        }
  lift $ onEffectRun (interceptor lenv) ctx
  startTime <- liftIO getCurrentTime

  result <- lift f

  endTime <- liftIO getCurrentTime
  let duration = diffUTCTime endTime startTime
  lift $ onEffectComplete (interceptor lenv) (operationName ctx) duration

  pure result
{-# INLINABLE effect #-}

-- | Adapt a @with@-style bracket function into a layer.
--
-- Many Haskell libraries expose resource management in continuation-
-- passing style: @withFoo :: (Foo -> m a) -> m a@. The library
-- handles acquisition and cleanup internally, and you use the
-- resource within the continuation.
--
-- 'bracketed' captures the resource from this pattern and extends its
-- lifetime to match the enclosing 'withLayer' scope, so it
-- participates in the same cleanup guarantees as 'resource'.
--
-- __Connection pool from a library:__
--
-- @
-- poolLayer :: Layer IO () ConnectionPool
-- poolLayer = bracketed withConnectionPool
-- @
--
-- __TLS manager:__
--
-- @
-- tlsLayer :: Layer IO TlsConfig TlsManager
-- tlsLayer = do
--   cfg <- ask
--   bracketed (withTlsManager cfg)
-- @
--
-- __Network socket:__
--
-- @
-- socketLayer :: Layer IO SocketConfig Socket
-- socketLayer = do
--   cfg <- ask
--   bracketed (withBoundSocket cfg.host cfg.port)
-- @
--
-- Prefer 'resource' when you have separate acquire and release
-- functions; 'bracketed' is specifically for libraries that only
-- expose the @with@ pattern.
bracketed ::
  forall m deps env. MonadUnliftIO m =>
  (forall a. (env -> m a) -> m a) ->
  Layer m deps env
bracketed f = Layer $ \_ _deps -> do
  (resourceVar :: MVar (Either SomeException env)) <- newEmptyMVar
  cleanupVar <- newEmptyMVar
  void $ allocateU
    ( async do
      lift $ f (\env -> putMVar resourceVar (Right env) >> takeMVar cleanupVar)
        `catch` (\(e :: SomeException) -> putMVar resourceVar (Left e) >> throwIO e)
    )
    (\h -> do
      putMVar cleanupVar ()
      wait h
    )
  takeMVar resourceVar >>= either throwIO pure

-- | Wrap a layer so it can be shared via 'service'.
--
-- 'mkService' does not change the layer's behaviour. It simply tags
-- the layer so that 'service' knows to cache its result. Define your
-- service once at the top level, and reference it from every consumer:
--
-- @
-- poolService :: Service IO Config ConnectionPool
-- poolService = mkService $ do
--   cfg <- ask
--   resource (createPool cfg) drainPool
-- @
mkService :: Layer m deps env -> Service m deps env
mkService = Service

-- | Obtain a shared instance of a 'Service'.
--
-- The first call within a composition initialises the service and
-- caches the result, keyed by the output 'TypeRep'. Subsequent calls
-- return the cached instance. Initialisation is thread-safe:
-- concurrent requests block until the first completes.
--
-- __Two consumers sharing one pool:__
--
-- @
-- webLayer :: Layer IO Config WebServer
-- webLayer = do
--   pool <- service poolService  -- first call creates the pool
--   resource (startWeb pool) stopWeb
--
-- apiLayer :: Layer IO Config ApiServer
-- apiLayer = do
--   pool <- service poolService  -- reuses the same pool
--   resource (startApi pool) stopApi
--
-- app :: Layer IO Config (WebServer, ApiServer)
-- app = webLayer &&& apiLayer   -- one pool, two consumers
-- @
--
-- __Using a service in a sequential pipeline:__
--
-- @
-- appLayer :: Layer IO () App
-- appLayer = configLayer >>> do
--   pool <- service poolService
--   db   <- service dbService
--   effect $ pure (App pool db)
-- @
--
-- __Caveat__: caching is keyed by output type. Two @Service@ values
-- that produce the same type will share a cache entry. Use a newtype
-- wrapper if you need distinct instances of the same underlying type.
service :: forall m deps env. (MonadUnliftIO m, Typeable env) => Service m deps env -> Layer m deps env
service = getOrCreateCachedService

-------------------------------------------------------------------------------
-- Typeclass instances --------------------------------------------------------
-------------------------------------------------------------------------------

instance Functor m => Functor (Layer m deps) where
  fmap f (Layer l) = Layer $ \lenv deps -> fmap f (l lenv deps)

instance MonadUnliftIO m => Applicative (Layer m deps) where
  pure = pureLayer
  lf <*> la = dimap (\d -> (d, d)) (uncurry ($)) (zipLayer lf la)

instance MonadUnliftIO m => Selective (Layer m deps) where
  select = selectM

instance MonadUnliftIO m => Monad (Layer m deps) where
  return = pure
  (Layer l) >>= f = Layer $ \lenv deps -> do
    a <- l lenv deps
    let Layer l' = f a
    l' lenv deps

instance (MonadUnliftIO m, Semigroup env) => Semigroup (Layer m deps env) where
  lf <> la = liftA2 (<>) lf la
  sconcat (h :| t) = foldl (<>) h t
  stimes n x
    | n <= 0 = error "stimes: positive multiplier expected"
    | n == 1 = x
    | otherwise = x <> stimes (n - 1) x

instance (MonadUnliftIO m, Monoid env) => Monoid (Layer m deps env) where
  mempty = pure mempty
  mappend = (<>)
  mconcat = foldl (<>) mempty

-- | Contravariantly map over the input. Named form of 'lmap'.
--
-- Useful when a layer expects a different input type than you have,
-- and you can project from one to the other:
--
-- @
-- -- dbLayer needs Config, but we have AppEnv
-- dbFromAppEnv :: Layer IO AppEnv Database
-- dbFromAppEnv = mapLayer (.config) dbLayer
-- @
mapLayer :: (b -> a) -> Layer m a c -> Layer m b c
mapLayer f (Layer l) = Layer $ \lenv env -> l lenv (f env)

instance (MonadUnliftIO m) => MonadReader deps (Layer m deps) where
  ask = Layer $ \_ deps -> pure deps
  local f (Layer l) = Layer $ \lenv deps -> l lenv (f deps)

-- | Sequential (vertical) composition. Named form of @>>>@.
--
-- Feeds the output of the upstream layer as dependencies to the
-- downstream layer. This is how you express "build A, then use A to
-- build B":
--
-- @
-- -- using the operator:
-- appLayer = configLayer >>> dbLayer
--
-- -- equivalent named version:
-- appLayer = composeLayer configLayer dbLayer
-- @
composeLayer ::
  MonadUnliftIO m =>
  -- | "Upstream" layer
  Layer m a b ->
  -- | "Downstream" layer
  Layer m b c ->
  Layer m a c
composeLayer (Layer upper) (Layer lower) = Layer $ \lenv deps -> do
  lift $ onCompositionStart (interceptor lenv) Sequential
  startTime <- liftIO getCurrentTime
  b <- upper lenv deps
  c <- lower lenv b
  endTime <- liftIO getCurrentTime
  lift $ onCompositionEnd (interceptor lenv) Sequential (diffUTCTime endTime startTime)
  pure c

instance MonadUnliftIO m => Category (Layer m) where
  id = Layer $ \_ deps -> pure deps
  (.) = flip composeLayer

instance MonadUnliftIO m => Profunctor (Layer m) where
  dimap f g (Layer l) = Layer $ \lenv a -> fmap g (l lenv (f a))
  lmap  = mapLayer
  rmap = fmap

instance MonadUnliftIO m => Strong (Layer m) where
  first' :: Layer m a b -> Layer m (a, c) (b, c)
  first' (Layer l) = Layer $ \lenv (a, c) -> do
    b <- l lenv a
    pure (b, c)

  second' :: Layer m a b -> Layer m (c, a) (c, b)
  second' (Layer l) = Layer $ \lenv (c, a) -> do
    b <- l lenv a
    pure (c, b)

instance MonadUnliftIO m => Choice (Layer m) where
  left' :: Layer m a b -> Layer m (Either a c) (Either b c)
  left' (Layer l) = Layer $ \lenv -> \case
    Left a -> Left <$> l lenv a
    Right c -> pure (Right c)

  right' :: Layer m a b -> Layer m (Either c a) (Either c b)
  right' (Layer l) = Layer $ \lenv -> \case
    Right a -> Right <$> l lenv a
    Left c -> pure (Left c)

instance MonadUnliftIO m => Traversing (Layer m) where
  traverse' :: Traversable f => Layer m a b -> Layer m (f a) (f b)
  traverse' (Layer l) = Layer $ \lenv fa -> traverse (l lenv) fa

instance MonadUnliftIO m => Arrow (Layer m) where
  arr f = Layer $ \_ a -> pure (f a)
  first = first'
  second = second'

instance MonadUnliftIO m => ArrowChoice (Layer m) where
  left = left'
  right = right'

-------------------------------------------------------------------------------

-- | Thrown by 'empty' \/ 'zeroArrow'.  Caught internally by @\<|\>@ and
-- @\<+\>@ to implement fallback behavior. You should not need to catch
-- this yourself; it is an implementation detail of the 'Alternative'
-- instance.
data EmptyLayer = EmptyLayer
  deriving stock (Show)

instance Exception EmptyLayer

instance MonadUnliftIO m => Alternative (Layer m deps) where
  empty = Layer $ \_ _ -> throwIO EmptyLayer
  Layer l1 <|> Layer l2 = Layer $ \lenv deps ->
    tryWithCleanup (l1 lenv deps) (l2 lenv deps)

instance MonadUnliftIO m => MonadPlus (Layer m deps) where
  mzero = empty
  mplus = (<|>)

instance MonadUnliftIO m => ArrowZero (Layer m) where
  zeroArrow = Layer $ \_ _ -> throwIO EmptyLayer

instance MonadUnliftIO m => ArrowPlus (Layer m) where
  (Layer l1) <+> (Layer l2) = Layer $ \lenv deps ->
    tryWithCleanup (l1 lenv deps) (l2 lenv deps)

instance MonadUnliftIO m => ArrowApply (Layer m) where
  app = Layer $ \lenv (Layer l, b) -> l lenv b

tryWithCleanup :: MonadUnliftIO m => ResourceT m a -> ResourceT m a -> ResourceT m a
tryWithCleanup l r = mask $ \restore -> do
  tmpState <- liftIO createInternalState
  try (restore $ runInternalState (lift l) tmpState) >>= \case
    Left e -> do
      liftIO $ stateCleanupChecked (Just (e :: SomeException)) tmpState
      restore r
    Right result -> withInternalState $ \stateMain -> do
      tmpState' <- liftIO $ readIORef tmpState
      atomicModifyIORef' stateMain $ \state -> (unsafelyMergeReleaseMap state tmpState', ())
      pure result

-------------------------------------------------------------------------------
-- Composition
-------------------------------------------------------------------------------

-- | Parallel (horizontal) composition. Powers @&&&@ and @\<*\>@.
--
-- Both layers run concurrently via @async@ in isolated resource
-- scopes. On success, the scopes merge into the parent. On failure,
-- the surviving branch is cleaned up before the exception propagates.
--
-- This means @&&&@ is not just a convenience operator; it actually
-- reduces startup time by running independent layers in parallel.
--
-- @
-- -- named form:
-- fullSystem :: Layer IO Config (Database, WebServer)
-- fullSystem = zipLayer dbLayer webLayer
--
-- -- operator form (more common):
-- fullSystem :: Layer IO Config (Database, WebServer)
-- fullSystem = dbLayer &&& webLayer
-- @
--
-- If one branch throws, the other is cancelled and its resources are
-- released before the exception propagates to the caller.
zipLayer ::
  (MonadUnliftIO m) =>
  Layer m d1 o1 ->
  Layer m d2 o2 ->
  Layer m (d1, d2) (o1, o2)
zipLayer (Layer l1) (Layer l2) = Layer $ \lenv (d1, d2) -> do
  -- For concurrent execution, we need to:
  -- 1. Run each layer in isolation to prevent cross-contamination if one fails
  -- 2. Capture their resources and transfer to parent context if successful

  -- Create isolated states for each branch
  stateA <- liftIO createInternalState
  stateB <- liftIO createInternalState

  mask $ \restore -> do
    asyncA <- async $ restore $ runInternalState (lift $ l1 lenv d1) stateA
    asyncB <- async $ restore $ runInternalState (lift $ l2 lenv d2) stateB
    eRes <- try (restore (atomically $ (,) <$> waitCatchSTM asyncA <*> waitCatchSTM asyncB))
    stateA' <- liftIO $ readIORef stateA
    stateB' <- liftIO $ readIORef stateB
    case eRes of
      Left e -> do
        uninterruptibleCancel asyncA
        uninterruptibleCancel asyncB
        liftIO (stateCleanupChecked (Just e) stateA `finally` stateCleanupChecked (Just e) stateB)
        throwIO e
      Right res -> case res of
        (Left eA, Left eB) -> do
          liftIO (stateCleanupChecked (Just eA) stateA `finally` stateCleanupChecked (Just eB) stateB)
          throwIO eA
        (Left eA, Right _) -> do
          liftIO (stateCleanupChecked (Just eA) stateA `finally` stateCleanupChecked Nothing stateB)
          throwIO eA
        (Right _envA, Left eB) -> do
          liftIO (stateCleanupChecked (Just eB) stateB `finally` stateCleanupChecked Nothing stateA)
          throwIO eB
        (Right envA, Right envB) -> withInternalState $ \stateMain -> do
          let mergedStates = unsafelyMergeReleaseMap stateA' stateB'
          atomicModifyIORef' stateMain $ \state -> (unsafelyMergeReleaseMap state mergedStates, ())
          pure (envA, envB)

-- The first map is the one that is being merged into. Its refcount is not affected.
unsafelyMergeReleaseMap :: ReleaseMap -> ReleaseMap -> ReleaseMap
unsafelyMergeReleaseMap (ReleaseMap nextKey refCount stateA) (ReleaseMap nextKey' _ stateB) = ReleaseMap (nextKey + nextKey') refCount (stateA <> mapKeysMonotonic (+ nextKey) stateB)
unsafelyMergeReleaseMap l ReleaseMapClosed = l
unsafelyMergeReleaseMap ReleaseMapClosed r = r

-------------------------------------------------------------------------------
-- Running layers
-------------------------------------------------------------------------------

-- | Build a layer and return its output, releasing resources immediately.
--
-- This is useful in tests and REPL sessions where you want a quick
-- result without managing a callback scope:
--
-- @
-- -- in a test:
-- cfg <- runLayer () configLayer
-- cfg.port \`shouldBe\` 8080
--
-- -- in GHCi:
-- >>> runLayer () configLayer
-- Config {port = 8080, host = "localhost"}
-- @
--
-- __Warning__: any handles or connections in the output will be
-- invalid after this returns, because resources are released as
-- soon as the layer finishes building.  Use 'withLayer' for real
-- applications.
runLayer :: MonadUnliftIO m => deps -> Layer m deps env -> m env
runLayer deps (Layer l) = do
  lenv <- LayerEnv <$> newMVar emptyTypeMap <*> pure nullInterceptor
  runResourceT (l lenv deps)

-- | Build a layer, use its environment, and guarantee cleanup.
--
-- This is the primary way to run a layer in production code.
-- Resources are released in reverse acquisition order when the
-- callback returns or throws.
--
-- __Typical main:__
--
-- @
-- main :: IO ()
-- main = withLayer () appLayer $ \\(cfg, (db, ws)) ->
--   serveRequests ws
-- @
--
-- __Nested scopes:__
--
-- @
-- main :: IO ()
-- main = withLayer () configLayer $ \\cfg ->
--   withLayer cfg dbLayer $ \\db -> do
--     rows <- query db "SELECT 1"
--     print rows
-- @
--
-- Prefer the composed form (@configLayer >>> dbLayer@) over nesting
-- when possible, since it is both more concise and enables parallel
-- composition.
withLayer ::
  MonadUnliftIO m =>
  -- | Dependencies
  deps ->
  -- | Layer to run
  Layer m deps env ->
  -- | Action that needs the env
  (env -> m r) ->
  m r
withLayer deps (Layer l) useEnv = runResourceT $ do
  lenv <- LayerEnv <$> newMVar emptyTypeMap <*> pure nullInterceptor
  env <- l lenv deps
  lift $ useEnv env

-- | Like 'runLayer', but with a custom 'LayerInterceptor' for
-- observing construction events.
--
-- @
-- diags <- runLayerWithInterceptor loggingInterceptor () appLayer
-- @
runLayerWithInterceptor ::
  MonadUnliftIO m =>
  -- | Custom interceptor
  LayerInterceptor m ->
  -- | Dependencies
  deps ->
  -- | Layer to run
  Layer m deps env ->
  m env
runLayerWithInterceptor i deps (Layer l) = do
  lenv <- LayerEnv <$> newMVar emptyTypeMap <*> pure i
  runResourceT (l lenv deps)

-- | Like 'withLayer', but with a custom 'LayerInterceptor'.
--
-- @
-- main :: IO ()
-- main = withLayerAndInterceptor loggingInterceptor () appLayer $ \\env ->
--   runApp env
-- @
withLayerAndInterceptor ::
  MonadUnliftIO m =>
  -- | Custom interceptor
  LayerInterceptor m ->
  -- | Dependencies
  deps ->
  -- | Layer to run
  Layer m deps env ->
  -- | Action that needs the env
  (env -> m r) ->
  m r
withLayerAndInterceptor i deps (Layer l) useEnv = runResourceT $ do
  lenv <- LayerEnv <$> newMVar emptyTypeMap <*> pure i
  env <- l lenv deps
  lift $ useEnv env