From e0b9b09e70c386b2da17d1f0a15b0511861c89e8 Mon Sep 17 00:00:00 2001
From: Michael Arnaldi <michael.arnaldi@effectful.co>
Date: Sat, 30 Nov 2024 21:45:27 +0100
Subject: [PATCH] Implement Effect.fn to define traced functions (#3938)

Co-authored-by: Tim <hello@timsmart.co>
---
 .changeset/spicy-adults-hug.md         |  17 ++
 eslint.config.mjs                      |   1 +
 packages/effect/src/Effect.ts          | 327 ++++++++++++++++++++++++-
 packages/effect/src/Utils.ts           |   8 +
 packages/effect/test/Effect/fn.test.ts |  62 +++++
 5 files changed, 414 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/spicy-adults-hug.md
 create mode 100644 packages/effect/test/Effect/fn.test.ts

diff --git a/.changeset/spicy-adults-hug.md b/.changeset/spicy-adults-hug.md
new file mode 100644
index 00000000000..6b2d12515ea
--- /dev/null
+++ b/.changeset/spicy-adults-hug.md
@@ -0,0 +1,17 @@
+---
+"effect": minor
+---
+
+Implement Effect.fn to define traced functions.
+
+```ts
+import { Effect } from "effect"
+
+const logExample = Effect.fn("example")(function* <N extends number>(n: N) {
+  yield* Effect.annotateCurrentSpan("n", n)
+  yield* Effect.logInfo(`got: ${n}`)
+  yield* Effect.fail(new Error())
+}, Effect.delay("1 second"))
+
+Effect.runFork(logExample(100).pipe(Effect.catchAllCause(Effect.logError)))
+```
diff --git a/eslint.config.mjs b/eslint.config.mjs
index befca0edc42..5cb490c6d24 100644
--- a/eslint.config.mjs
+++ b/eslint.config.mjs
@@ -70,6 +70,7 @@ export default [
       ],
 
       "no-unused-vars": "off",
+      "require-yield": "off",
       "prefer-rest-params": "off",
       "prefer-spread": "off",
       "import/first": "error",
diff --git a/packages/effect/src/Effect.ts b/packages/effect/src/Effect.ts
index 0f658e3f540..b431b8073ab 100644
--- a/packages/effect/src/Effect.ts
+++ b/packages/effect/src/Effect.ts
@@ -25,6 +25,7 @@ import { dual } from "./Function.js"
 import type * as HashMap from "./HashMap.js"
 import type * as HashSet from "./HashSet.js"
 import type { TypeLambda } from "./HKT.js"
+import * as internalCause from "./internal/cause.js"
 import * as _console from "./internal/console.js"
 import { TagProto } from "./internal/context.js"
 import * as effect from "./internal/core-effect.js"
@@ -59,7 +60,7 @@ import type * as Supervisor from "./Supervisor.js"
 import type * as Tracer from "./Tracer.js"
 import type { Concurrency, Contravariant, Covariant, NoExcessProperties, NoInfer, NotFunction } from "./Types.js"
 import type * as Unify from "./Unify.js"
-import type { YieldWrap } from "./Utils.js"
+import { internalCall, isGeneratorFunction, type YieldWrap } from "./Utils.js"
 
 /**
  * @since 2.0.0
@@ -9776,3 +9777,327 @@ export declare namespace Service {
   export type MakeAccessors<Make> = Make extends { readonly accessors: true } ? true
     : false
 }
+
+/**
+ * @since 3.11.0
+ * @category models
+ */
+export namespace fn {
+  /**
+   * @since 3.11.0
+   * @category models
+   */
+  export type FnEffect<AEff, Eff extends YieldWrap<Effect<any, any, any>>> = Effect<
+    AEff,
+    [Eff] extends [never] ? never : [Eff] extends [YieldWrap<Effect<infer _A, infer E, infer _R>>] ? E : never,
+    [Eff] extends [never] ? never : [Eff] extends [YieldWrap<Effect<infer _A, infer _E, infer R>>] ? R : never
+  >
+
+  /**
+   * @since 3.11.0
+   * @category models
+   */
+  export type Gen = {
+    <Eff extends YieldWrap<Effect<any, any, any>>, AEff, Args extends Array<any>>(
+      body: (...args: Args) => Generator<Eff, AEff, never>
+    ): (...args: Args) => fn.FnEffect<AEff, Eff>
+    <Eff extends YieldWrap<Effect<any, any, any>>, AEff, Args extends Array<any>, A extends Effect<any, any, any>>(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A
+    ): (...args: Args) => A
+    <Eff extends YieldWrap<Effect<any, any, any>>, AEff, Args extends Array<any>, A, B extends Effect<any, any, any>>(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B
+    ): (...args: Args) => B
+    <
+      Eff extends YieldWrap<Effect<any, any, any>>,
+      AEff,
+      Args extends Array<any>,
+      A,
+      B,
+      C extends Effect<any, any, any>
+    >(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B,
+      c: (_: B) => C
+    ): (...args: Args) => C
+    <
+      Eff extends YieldWrap<Effect<any, any, any>>,
+      AEff,
+      Args extends Array<any>,
+      A,
+      B,
+      C,
+      D extends Effect<any, any, any>
+    >(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B,
+      c: (_: B) => C,
+      d: (_: C) => D
+    ): (...args: Args) => D
+    <
+      Eff extends YieldWrap<Effect<any, any, any>>,
+      AEff,
+      Args extends Array<any>,
+      A,
+      B,
+      C,
+      D,
+      E extends Effect<any, any, any>
+    >(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B,
+      c: (_: B) => C,
+      d: (_: C) => D,
+      e: (_: D) => E
+    ): (...args: Args) => E
+    <
+      Eff extends YieldWrap<Effect<any, any, any>>,
+      AEff,
+      Args extends Array<any>,
+      A,
+      B,
+      C,
+      D,
+      E,
+      F extends Effect<any, any, any>
+    >(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B,
+      c: (_: B) => C,
+      d: (_: C) => D,
+      e: (_: D) => E,
+      f: (_: E) => F
+    ): (...args: Args) => F
+    <
+      Eff extends YieldWrap<Effect<any, any, any>>,
+      AEff,
+      Args extends Array<any>,
+      A,
+      B,
+      C,
+      D,
+      E,
+      F,
+      G extends Effect<any, any, any>
+    >(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B,
+      c: (_: B) => C,
+      d: (_: C) => D,
+      e: (_: D) => E,
+      f: (_: E) => F,
+      g: (_: F) => G
+    ): (...args: Args) => G
+    <
+      Eff extends YieldWrap<Effect<any, any, any>>,
+      AEff,
+      Args extends Array<any>,
+      A,
+      B,
+      C,
+      D,
+      E,
+      F,
+      G,
+      H extends Effect<any, any, any>
+    >(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B,
+      c: (_: B) => C,
+      d: (_: C) => D,
+      e: (_: D) => E,
+      f: (_: E) => F,
+      g: (_: F) => G,
+      h: (_: G) => H
+    ): (...args: Args) => H
+    <
+      Eff extends YieldWrap<Effect<any, any, any>>,
+      AEff,
+      Args extends Array<any>,
+      A,
+      B,
+      C,
+      D,
+      E,
+      F,
+      G,
+      H,
+      I extends Effect<any, any, any>
+    >(
+      body: (...args: Args) => Generator<Eff, AEff, never>,
+      a: (_: fn.FnEffect<AEff, Eff>) => A,
+      b: (_: A) => B,
+      c: (_: B) => C,
+      d: (_: C) => D,
+      e: (_: D) => E,
+      f: (_: E) => F,
+      g: (_: F) => G,
+      h: (_: G) => H,
+      i: (_: H) => I
+    ): (...args: Args) => I
+  }
+
+  /**
+   * @since 3.11.0
+   * @category models
+   */
+  export type NonGen = {
+    <Eff extends Effect<any, any, any>, Args extends Array<any>>(
+      body: (...args: Args) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, C, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => C,
+      c: (_: C) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, C, D, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => C,
+      c: (_: C) => D,
+      d: (_: D) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, C, D, E, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => C,
+      c: (_: C) => D,
+      d: (_: D) => E,
+      e: (_: E) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, C, D, E, F, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => C,
+      c: (_: C) => D,
+      d: (_: D) => E,
+      e: (_: E) => F,
+      f: (_: E) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, C, D, E, F, G, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => C,
+      c: (_: C) => D,
+      d: (_: D) => E,
+      e: (_: E) => F,
+      f: (_: E) => G,
+      g: (_: G) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, C, D, E, F, G, H, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => C,
+      c: (_: C) => D,
+      d: (_: D) => E,
+      e: (_: E) => F,
+      f: (_: E) => G,
+      g: (_: G) => H,
+      h: (_: H) => Eff
+    ): (...args: Args) => Eff
+    <Eff extends Effect<any, any, any>, A, B, C, D, E, F, G, H, I, Args extends Array<any>>(
+      body: (...args: Args) => A,
+      a: (_: A) => B,
+      b: (_: B) => C,
+      c: (_: C) => D,
+      d: (_: D) => E,
+      e: (_: E) => F,
+      f: (_: E) => G,
+      g: (_: G) => H,
+      h: (_: H) => I,
+      i: (_: H) => Eff
+    ): (...args: Args) => Eff
+  }
+}
+
+/**
+ * Creates a function that returns an Effect which is automatically traced with a span pointing to the call site.
+ *
+ * The function can be created both using a generator function that can yield effects or using a normal function.
+ *
+ * @since 3.11.0
+ * @category function
+ *
+ * @example
+ * import { Effect } from "effect"
+ *
+ * const logExample = Effect.fn("logExample")(
+ *   function*<N extends number>(n: N) {
+ *     yield* Effect.annotateCurrentSpan("n", n)
+ *     yield* Effect.logInfo(`got: ${n}`)
+ *     yield* Effect.fail(new Error())
+ *   },
+ *   Effect.delay("1 second")
+ * )
+ *
+ * Effect.runFork(
+ *   // this location is printed on the stack trace of the following `Effect.logError`
+ *   logExample(100).pipe(
+ *     Effect.catchAllCause(Effect.logError)
+ *   )
+ * )
+ */
+export const fn: (
+  name: string,
+  options?: Tracer.SpanOptions
+) => fn.Gen & fn.NonGen = (name, options) => (body: Function, ...pipeables: Array<any>) => {
+  return function(this: any, ...args: Array<any>) {
+    const limit = Error.stackTraceLimit
+    Error.stackTraceLimit = 2
+    const error = new Error()
+    Error.stackTraceLimit = limit
+    let cache: false | string = false
+    const captureStackTrace = () => {
+      if (cache !== false) {
+        return cache
+      }
+      if (error.stack) {
+        const stack = error.stack.trim().split("\n")
+        cache = stack.slice(2).join("\n").trim()
+        return cache
+      }
+    }
+    let effect: Effect<any, any, any>
+    let fnError: any = undefined
+    try {
+      effect = isGeneratorFunction(body)
+        ? gen(() => internalCall(() => body.apply(this, args)))
+        : body.apply(this, args)
+    } catch (error) {
+      fnError = error
+      effect = die(error)
+    }
+    try {
+      for (const x of pipeables) {
+        effect = x(effect)
+      }
+    } catch (error) {
+      effect = fnError
+        ? failCause(internalCause.sequential(
+          internalCause.die(fnError),
+          internalCause.die(error)
+        ))
+        : die(error)
+    }
+    const opts: any = (options && "captureStackTrace" in options) ? options : { captureStackTrace, ...options }
+    return withSpan(effect, name, opts)
+  }
+}
diff --git a/packages/effect/src/Utils.ts b/packages/effect/src/Utils.ts
index 4591b903df1..6d2d3464f58 100644
--- a/packages/effect/src/Utils.ts
+++ b/packages/effect/src/Utils.ts
@@ -791,3 +791,11 @@ const tracingFunction = (name: string) => {
  * @category tracing
  */
 export const internalCall = tracingFunction("effect_internal_function")
+
+const genConstructor = (function*() {}).constructor
+
+/**
+ * @since 3.11.0
+ */
+export const isGeneratorFunction = (u: unknown): u is (...args: Array<any>) => Generator<any, any, any> =>
+  isObject(u) && u.constructor === genConstructor
diff --git a/packages/effect/test/Effect/fn.test.ts b/packages/effect/test/Effect/fn.test.ts
new file mode 100644
index 00000000000..7a1725ea909
--- /dev/null
+++ b/packages/effect/test/Effect/fn.test.ts
@@ -0,0 +1,62 @@
+import { Cause, Effect } from "effect"
+import { assert, describe, it } from "effect/test/utils/extend"
+
+describe("Effect.fn", () => {
+  it.effect("catches defects in the function", () =>
+    Effect.gen(function*() {
+      let caught: Cause.Cause<never> | undefined
+      const fn = Effect.fn("test")(
+        (): Effect.Effect<void> => {
+          throw new Error("test")
+        },
+        Effect.tapErrorCause((cause) => {
+          caught = cause
+          return Effect.void
+        })
+      )
+      const cause = yield* fn().pipe(
+        Effect.sandbox,
+        Effect.flip
+      )
+      assert(Cause.isDieType(cause))
+      assert.isTrue(cause.defect instanceof Error && cause.defect.message === "test")
+      assert.strictEqual(caught, cause)
+    }))
+
+  it.effect("catches defects in pipeline", () =>
+    Effect.gen(function*() {
+      const fn = Effect.fn("test")(
+        () => Effect.void,
+        (_): Effect.Effect<void> => {
+          throw new Error("test")
+        }
+      )
+      const cause = yield* fn().pipe(
+        Effect.sandbox,
+        Effect.flip
+      )
+      assert(Cause.isDieType(cause))
+      assert.isTrue(cause.defect instanceof Error && cause.defect.message === "test")
+    }))
+
+  it.effect("catches defects in both fn & pipeline", () =>
+    Effect.gen(function*() {
+      const fn = Effect.fn("test")(
+        (): Effect.Effect<void> => {
+          throw new Error("test")
+        },
+        (_): Effect.Effect<void> => {
+          throw new Error("test2")
+        }
+      )
+      const cause = yield* fn().pipe(
+        Effect.sandbox,
+        Effect.flip
+      )
+      assert(Cause.isSequentialType(cause))
+      assert(Cause.isDieType(cause.left))
+      assert(Cause.isDieType(cause.right))
+      assert.isTrue(cause.left.defect instanceof Error && cause.left.defect.message === "test")
+      assert.isTrue(cause.right.defect instanceof Error && cause.right.defect.message === "test2")
+    }))
+})