WIP: Pass Trace transparently through withError

arrow-libs/core/arrow-core/api/arrow-core.api

@@ -881,6 +881,7 @@ public final class arrow/core/raise/DefaultRaise : arrow/core/raise/Raise {
 	public fun bindAll-1TN0_VU (Ljava/util/Set;)Ljava/util/Set;
 	public fun bindAll-vcjLgH4 (Ljava/util/List;)Ljava/util/List;
 	public final fun complete ()Z
+	public fun getUnderlyingRaise ()Larrow/core/raise/Raise;
 	public fun invoke (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun invoke (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public final fun isTraced ()Z
@@ -909,6 +910,7 @@ public final class arrow/core/raise/IorRaise : arrow/core/raise/Raise {
 	public final fun bindAllIor (Ljava/util/Set;)Ljava/util/Set;
 	public final fun combine (Ljava/lang/Object;)Ljava/lang/Object;
 	public final fun getCombineError ()Lkotlin/jvm/functions/Function2;
+	public fun getUnderlyingRaise ()Larrow/core/raise/Raise;
 	public fun invoke (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun invoke (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun raise (Ljava/lang/Object;)Ljava/lang/Void;
@@ -923,6 +925,7 @@ public abstract interface class arrow/core/raise/Raise {
 	public abstract fun bindAll (Ljava/util/Map;)Ljava/util/Map;
 	public abstract fun bindAll-1TN0_VU (Ljava/util/Set;)Ljava/util/Set;
 	public abstract fun bindAll-vcjLgH4 (Ljava/util/List;)Ljava/util/List;
+	public abstract fun getUnderlyingRaise ()Larrow/core/raise/Raise;
 	public abstract fun invoke (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public abstract fun invoke (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public abstract fun raise (Ljava/lang/Object;)Ljava/lang/Void;
@@ -936,6 +939,7 @@ public final class arrow/core/raise/Raise$DefaultImpls {
 	public static fun bindAll (Larrow/core/raise/Raise;Ljava/util/Map;)Ljava/util/Map;
 	public static fun bindAll-1TN0_VU (Larrow/core/raise/Raise;Ljava/util/Set;)Ljava/util/Set;
 	public static fun bindAll-vcjLgH4 (Larrow/core/raise/Raise;Ljava/util/List;)Ljava/util/List;
+	public static fun getUnderlyingRaise (Larrow/core/raise/Raise;)Larrow/core/raise/Raise;
 	public static fun invoke (Larrow/core/raise/Raise;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static fun invoke (Larrow/core/raise/Raise;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
@@ -954,6 +958,7 @@ public class arrow/core/raise/RaiseAccumulate : arrow/core/raise/Raise {
 	public fun bindAll-vcjLgH4 (Ljava/util/List;)Ljava/util/List;
 	public final fun bindNel (Larrow/core/Either;)Ljava/lang/Object;
 	public final fun getRaise ()Larrow/core/raise/Raise;
+	public fun getUnderlyingRaise ()Larrow/core/raise/Raise;
 	public fun invoke (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun invoke (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public final fun mapOrAccumulate (Ljava/lang/Iterable;Lkotlin/jvm/functions/Function2;)Ljava/util/List;
@@ -1019,11 +1024,14 @@ public final class arrow/core/raise/RaiseKt {
 	public static final fun merge (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun nullable (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun option (Lkotlin/jvm/functions/Function1;)Larrow/core/Option;
+	public static final fun raiseWithTrace-LzgXH30 (Larrow/core/raise/Raise;Larrow/core/raise/Traced;Ljava/lang/Object;)Ljava/lang/Void;
 	public static final fun raisedOrRethrow (Ljava/util/concurrent/CancellationException;Larrow/core/raise/DefaultRaise;)Ljava/lang/Object;
+	public static final fun realUnderlying (Larrow/core/raise/Raise;)Larrow/core/raise/DefaultRaise;
 	public static final fun recover (Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun recover (Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun recover (Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Lkotlin/jvm/functions/Function1;
 	public static final fun recover (Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;)Lkotlin/jvm/functions/Function2;
+	public static final fun recoverTraced (Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;
 	public static final fun result (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun singleton (Lkotlin/jvm/functions/Function0;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun toEither (Lkotlin/jvm/functions/Function1;)Larrow/core/Either;
@@ -1037,7 +1045,6 @@ public final class arrow/core/raise/RaiseKt {
 	public static final fun toResult (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun toResult (Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun traced (Larrow/core/raise/Raise;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;
-	public static final fun withCause (Larrow/core/raise/Traced;Larrow/core/raise/Traced;)Larrow/core/raise/Traced;
 	public static final fun withError (Larrow/core/raise/Raise;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun zipOrAccumulate (Larrow/core/raise/Raise;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function9;)Ljava/lang/Object;
 	public static final fun zipOrAccumulate (Larrow/core/raise/Raise;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function8;)Ljava/lang/Object;
@@ -1071,6 +1078,7 @@ public final class arrow/core/raise/ResultRaise : arrow/core/raise/Raise {
 	public final fun bindAllResult (Ljava/util/List;)Ljava/util/List;
 	public final fun bindAllResult (Ljava/util/Map;)Ljava/util/Map;
 	public final fun bindAllResult (Ljava/util/Set;)Ljava/util/Set;
+	public fun getUnderlyingRaise ()Larrow/core/raise/Raise;
 	public fun invoke (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun invoke (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public synthetic fun raise (Ljava/lang/Object;)Ljava/lang/Void;
@@ -1099,6 +1107,7 @@ public final class arrow/core/raise/SingletonRaise : arrow/core/raise/Raise {
 	public final fun bindAllOption (Ljava/util/Set;)Ljava/util/Set;
 	public final fun ensure (Z)V
 	public final fun ensureNotNull (Ljava/lang/Object;)Ljava/lang/Object;
+	public fun getUnderlyingRaise ()Larrow/core/raise/Raise;
 	public final fun ignoreErrors (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun invoke (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun invoke (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
@@ -1108,18 +1117,18 @@ public final class arrow/core/raise/SingletonRaise : arrow/core/raise/Raise {
 }
 
 public final class arrow/core/raise/Trace {
-	public static final synthetic fun box-impl (Ljava/util/concurrent/CancellationException;)Larrow/core/raise/Trace;
-	public static fun constructor-impl (Ljava/util/concurrent/CancellationException;)Ljava/util/concurrent/CancellationException;
+	public static final synthetic fun box-impl (Larrow/core/raise/Traced;)Larrow/core/raise/Trace;
+	public static fun constructor-impl (Larrow/core/raise/Traced;)Larrow/core/raise/Traced;
 	public fun equals (Ljava/lang/Object;)Z
-	public static fun equals-impl (Ljava/util/concurrent/CancellationException;Ljava/lang/Object;)Z
-	public static final fun equals-impl0 (Ljava/util/concurrent/CancellationException;Ljava/util/concurrent/CancellationException;)Z
+	public static fun equals-impl (Larrow/core/raise/Traced;Ljava/lang/Object;)Z
+	public static final fun equals-impl0 (Larrow/core/raise/Traced;Larrow/core/raise/Traced;)Z
 	public fun hashCode ()I
-	public static fun hashCode-impl (Ljava/util/concurrent/CancellationException;)I
-	public static final fun printStackTrace-impl (Ljava/util/concurrent/CancellationException;)V
-	public static final fun stackTraceToString-impl (Ljava/util/concurrent/CancellationException;)Ljava/lang/String;
-	public static final fun suppressedExceptions-impl (Ljava/util/concurrent/CancellationException;)Ljava/util/List;
+	public static fun hashCode-impl (Larrow/core/raise/Traced;)I
+	public static final fun printStackTrace-impl (Larrow/core/raise/Traced;)V
+	public static final fun stackTraceToString-impl (Larrow/core/raise/Traced;)Ljava/lang/String;
+	public static final fun suppressedExceptions-impl (Larrow/core/raise/Traced;)Ljava/util/List;
 	public fun toString ()Ljava/lang/String;
-	public static fun toString-impl (Ljava/util/concurrent/CancellationException;)Ljava/lang/String;
-	public final synthetic fun unbox-impl ()Ljava/util/concurrent/CancellationException;
+	public static fun toString-impl (Larrow/core/raise/Traced;)Ljava/lang/String;
+	public final synthetic fun unbox-impl ()Larrow/core/raise/Traced;
 }
 

arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/raise/Fold.kt

@@ -152,6 +152,27 @@ public inline fun <Error, A, B> fold(
   }
 }
 
+@OptIn(DelicateRaiseApi::class)
+@ExperimentalTraceApi
+@RaiseDSL
+public inline fun <Error, A> recoverTraced(
+  @BuilderInference block: Raise<Error>.() -> A,
+  recover: (trace: Trace, error: Error) -> A
+): A {
+  contract {
+    callsInPlace(block, AT_MOST_ONCE)
+    callsInPlace(recover, AT_MOST_ONCE)
+  }
+  val nested = DefaultRaise(true)
+  return try {
+    block(nested).also { nested.complete() }
+  } catch (e: Traced) {
+    nested.complete()
+    val r: Error = e.raisedOrRethrow(nested)
+    recover(Trace(e), r)
+  }
+}
+
 /**
  * Inspect a [Trace] value of [Error].
  *
@@ -201,27 +222,21 @@ public inline fun <Error, A> Raise<Error>.traced(
     callsInPlace(block, AT_MOST_ONCE)
     callsInPlace(trace, AT_MOST_ONCE)
   }
-  val nested = DefaultRaise(true)
-  return try {
-    block(nested).also { nested.complete() }
-  } catch (e: Traced) {
-    nested.complete()
-    val r: Error = e.raisedOrRethrow(nested)
-    trace(Trace(e), r)
-    // If our outer Raise happens to be traced
-    // Then we want the stack trace to match the inner one
-    try {
-      raise(r)
-    } catch (rethrown: Traced) {
-      throw rethrown.withCause(e)
-    }
+  return recoverTraced(block) { t, r ->
+    trace(t, r)
+    raiseWithTrace(t, r)
   }
 }
 
-@PublishedApi
 @DelicateRaiseApi
-internal fun Traced.withCause(cause: Traced): Traced =
-  Traced(raised, raise, cause)
+@ExperimentalTraceApi
+@RaiseDSL
+public fun <Error> Raise<Error>.raiseWithTrace(trace: Trace, r: Error): Nothing =
+  try {
+    raise(r)
+  } catch (rethrown: Traced) {
+    throw Traced(rethrown.raised, rethrown.raise, trace)
+  }
 
 /** Returns the raised value, rethrows the CancellationException if not our scope */
 @PublishedApi
@@ -240,13 +255,22 @@ internal class DefaultRaise(@PublishedApi internal val isTraced: Boolean) : Rais
 
   @PublishedApi
   internal fun complete(): Boolean = isActive.getAndSet(false)
-  @OptIn(DelicateRaiseApi::class)
+  @OptIn(DelicateRaiseApi::class, ExperimentalTraceApi::class)
   override fun raise(r: Any?): Nothing = when {
     isActive.value -> throw if (isTraced) Traced(r, this) else NoTrace(r, this)
     else -> throw RaiseLeakedException()
   }
 }
 
+@DelicateRaiseApi
+@PublishedApi
+internal tailrec fun Raise<*>.realUnderlying(): DefaultRaise? =
+  when (val underlying = underlyingRaise) {
+    is DefaultRaise -> underlying
+    this -> null
+    else -> underlying.realUnderlying()
+  }
+
 @MustBeDocumented
 @Retention(AnnotationRetention.BINARY)
 @RequiresOptIn(RaiseCancellationExceptionCaptured, RequiresOptIn.Level.WARNING)
@@ -267,7 +291,8 @@ public sealed class RaiseCancellationException(
 internal expect class NoTrace(raised: Any?, raise: Raise<Any?>) : RaiseCancellationException
 
 @DelicateRaiseApi
-internal class Traced(raised: Any?, raise: Raise<Any?>, override val cause: Traced? = null): RaiseCancellationException(raised, raise)
+@ExperimentalTraceApi
+internal class Traced(raised: Any?, raise: Raise<Any?>, val originalTrace: Trace? = null): RaiseCancellationException(raised, raise)
 
 private class RaiseLeakedException : IllegalStateException(
   """

arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/raise/Raise.kt

@@ -144,6 +144,7 @@ public annotation class RaiseDSL
  * <!--- TEST lines.isEmpty() -->
  */
 public interface Raise<in Error> {
+  public val underlyingRaise: Raise<*> get() = this
 
   /**
    * Raises a _logical failure_ of type [Error].
@@ -653,6 +654,7 @@ public inline fun <Error, B : Any> Raise<Error>.ensureNotNull(value: B?, raise:
  * <!--- KNIT example-raise-dsl-11.kt -->
  * <!--- TEST lines.isEmpty() -->
  */
+@OptIn(DelicateRaiseApi::class, ExperimentalTraceApi::class)
 @RaiseDSL
 public inline fun <Error, OtherError, A> Raise<Error>.withError(
   transform: (OtherError) -> Error,
@@ -661,7 +663,9 @@ public inline fun <Error, OtherError, A> Raise<Error>.withError(
   contract {
     callsInPlace(transform, AT_MOST_ONCE)
   }
-  return recover(block) { raise(transform(it)) }
+  val outer = realUnderlying()
+  return if (outer?.isTraced == true) recoverTraced(block) { t, r -> raiseWithTrace(t, transform(r)) }
+  else recover(block) { raise(transform(it)) }
 }
 
 /**

arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/raise/Trace.kt

@@ -14,15 +14,17 @@ public annotation class ExperimentalTraceApi
 
 /** Tracing result. Allows to inspect the traces from where raise was called. */
 @ExperimentalTraceApi
+@OptIn(DelicateRaiseApi::class)
 @JvmInline
-public value class Trace(private val exception: CancellationException) {
+public value class Trace
+@PublishedApi internal constructor(private val exception: Traced) {
   /**
    * Returns the stacktrace as a [String]
    *
    * Note, the first line in the stacktrace will be the `RaiseCancellationException`.
    * The users call to `raise` can found in the_second line of the stacktrace.
    */
-  public fun stackTraceToString(): String = (exception.cause ?: exception).stackTraceToString()
+  public fun stackTraceToString(): String = exception.originalTrace?.stackTraceToString() ?: exception.stackTraceToString()
 
   /**
    * Prints the stacktrace.
@@ -31,7 +33,7 @@ public value class Trace(private val exception: CancellationException) {
    * The users call to `raise` can found in the_second line of the stacktrace.
    */
   public fun printStackTrace(): Unit =
-    (exception.cause ?: exception).printStackTrace()
+    exception.originalTrace?.printStackTrace() ?: exception.printStackTrace()
 
   /**
    * Returns the suppressed exceptions that occurred during cancellation of the surrounding coroutines,
@@ -41,5 +43,5 @@ public value class Trace(private val exception: CancellationException) {
    * if the finalizer then results in a `Throwable` it will be added as a `suppressedException` to the [CancellationException].
    */
   public fun suppressedExceptions(): List<Throwable> =
-    exception.cause?.suppressedExceptions.orEmpty() + exception.suppressedExceptions
+    exception.originalTrace?.suppressedExceptions().orEmpty() + exception.suppressedExceptions
 }

arrow-libs/core/arrow-core/src/commonTest/kotlin/arrow/core/raise/TraceSpec.kt

@@ -11,7 +11,7 @@ import kotlin.test.Test
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.test.runTest
 
-@OptIn(ExperimentalTraceApi::class)
+@OptIn(ExperimentalTraceApi::class, DelicateRaiseApi::class)
 class TraceSpec {
   @Test fun traceIsEmptyWhenNoErrors() = runTest {
     checkAll(Arb.int()) { i ->
@@ -54,4 +54,22 @@ class TraceSpec {
       }) { _, unit -> unit shouldBe Unit }
     }
   }
+
+  @Test
+  fun withErrorMaintainsTrace() = runTest {
+    val inner = CompletableDeferred<String>()
+    merge {
+      traced({
+        withError({ str: String -> str.length }) {
+          val e = shouldThrow<Traced> {
+            raise("")
+          }
+          inner.complete(e.stackTraceToString())
+          throw e
+        }
+      }) { traced, _ ->
+        inner.await() shouldBe traced.stackTraceToString()
+      }
+    } shouldBe 0
+  }
 }