module OpenTelemetry.Contrib.SpanTraversals (
  alterSpansUpwards,
  IterationInstruction (..),
) where

import Control.Monad.IO.Class
import Data.IORef
import OpenTelemetry.Internal.Trace.Types


data IterationInstruction a = Continue a | Halt


{- | Alter traces upwards from the provides span to the highest available mutable span. Only mutable spans may be altered.

 The step value indicates whether the desired topmost span has been reached or not. This function will continue to iterate
 upwards until either a span that cannot be mutated has been reached, or there are no more parent spans remaining.
-}
alterSpansUpwards :: (MonadIO m) => Span -> st -> (st -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan)) -> m st
alterSpansUpwards :: forall (m :: * -> *) st.
MonadIO m =>
Span
-> st
-> (st
    -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan))
-> m st
alterSpansUpwards (Span IORef ImmutableSpan
immutableSpanRef) st
st st -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan)
f = IO st -> m st
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO st -> m st) -> IO st -> m st
forall a b. (a -> b) -> a -> b
$ do
  (IterationInstruction st
step, ImmutableSpan
a') <- IORef ImmutableSpan
-> (ImmutableSpan
    -> (ImmutableSpan, (IterationInstruction st, ImmutableSpan)))
-> IO (IterationInstruction st, ImmutableSpan)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef ImmutableSpan
immutableSpanRef (\ImmutableSpan
a -> let (IterationInstruction st
step, ImmutableSpan
a') = st -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan)
f st
st ImmutableSpan
a in (ImmutableSpan
a', (IterationInstruction st
step, ImmutableSpan
a')))
  case IterationInstruction st
step of
    Continue st
st' -> case ImmutableSpan -> Maybe Span
spanParent ImmutableSpan
a' of
      Maybe Span
Nothing -> st -> IO st
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st'
      Just Span
s -> Span
-> st
-> (st
    -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan))
-> IO st
forall (m :: * -> *) st.
MonadIO m =>
Span
-> st
-> (st
    -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan))
-> m st
alterSpansUpwards Span
s st
st' st -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan)
f
    IterationInstruction st
Halt -> st -> IO st
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st
alterSpansUpwards (FrozenSpan SpanContext
_) st
st st -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan)
_ = st -> m st
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st
alterSpansUpwards (Dropped SpanContext
_) st
st st -> ImmutableSpan -> (IterationInstruction st, ImmutableSpan)
_ = st -> m st
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st