From 5cefd879d42b42c12f9819e2bf83cb89557bb3e6 Mon Sep 17 00:00:00 2001
From: Tim <hello@timsmart.co>
Date: Sun, 4 Feb 2024 00:30:19 +1300
Subject: [PATCH] ensure TagClass stack is set correctly (#2029)

---
 packages/effect/src/Context.ts          | 10 ++-------
 packages/effect/src/internal/context.ts | 21 +++++++++++++++++--
 packages/effect/test/Context.test.ts    | 27 +++++++++++++++++++++++--
 3 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/packages/effect/src/Context.ts b/packages/effect/src/Context.ts
index aba2226494..a02b6671cc 100644
--- a/packages/effect/src/Context.ts
+++ b/packages/effect/src/Context.ts
@@ -107,7 +107,7 @@ export declare namespace Tag {
  * @since 2.0.0
  * @category constructors
  */
-export const Tag: <Identifier, Service = Identifier>(key: string) => Tag<Identifier, Service> = internal.makeTag
+export const Tag: <Identifier, Service = Identifier>(key: string) => Tag<Identifier, Service> = internal.makeGenericTag
 
 const TypeId: unique symbol = internal.TypeId as TypeId
 
@@ -387,10 +387,4 @@ export const omit: <Services, S extends Array<ValidTagsById<Services>>>(
  * @since 2.0.0
  * @category constructors
  */
-export const TagClass = <const Id extends string>(id: Id) => <Self, Shape>(): TagClass<Self, Id, Shape> => {
-  function TagClass() {}
-  const original = Tag(id)
-  Object.assign(TagClass, original)
-  Object.setPrototypeOf(TagClass, Object.getPrototypeOf(original))
-  return TagClass as any
-}
+export const TagClass: <const Id extends string>(id: Id) => <Self, Shape>() => TagClass<Self, Id, Shape> = internal.Tag
diff --git a/packages/effect/src/internal/context.ts b/packages/effect/src/internal/context.ts
index 3ceb675281..39ea659f88 100644
--- a/packages/effect/src/internal/context.ts
+++ b/packages/effect/src/internal/context.ts
@@ -56,7 +56,7 @@ export const TagProto: any = {
 }
 
 /** @internal */
-export const makeTag = <Identifier, Service = Identifier>(key: string): C.Tag<Identifier, Service> => {
+export const makeGenericTag = <Identifier, Service = Identifier>(key: string): C.Tag<Identifier, Service> => {
   const limit = Error.stackTraceLimit
   Error.stackTraceLimit = 2
   const creationError = new Error()
@@ -71,6 +71,24 @@ export const makeTag = <Identifier, Service = Identifier>(key: string): C.Tag<Id
   return tag
 }
 
+/** @internal */
+export const Tag = <const Id extends string>(id: Id) => <Self, Shape>(): C.TagClass<Self, Id, Shape> => {
+  const limit = Error.stackTraceLimit
+  Error.stackTraceLimit = 2
+  const creationError = new Error()
+  Error.stackTraceLimit = limit
+
+  function TagClass() {}
+  Object.setPrototypeOf(TagClass, TagProto)
+  TagClass.key = id
+  Object.defineProperty(TagClass, "stack", {
+    get() {
+      return creationError.stack
+    }
+  })
+  return TagClass as any
+}
+
 /** @internal */
 export const TypeId: C.TypeId = Symbol.for("effect/Context") as C.TypeId
 
@@ -180,7 +198,6 @@ export const unsafeGet = dual<
   <Services, S, I>(self: C.Context<Services>, tag: C.Tag<I, S>) => S
 >(2, (self, tag) => {
   if (!self.unsafeMap.has(tag.key)) {
-    console.log(self.unsafeMap, tag.key)
     throw serviceNotFoundError(tag as any)
   }
   return self.unsafeMap.get(tag.key)! as any
diff --git a/packages/effect/test/Context.test.ts b/packages/effect/test/Context.test.ts
index bcb6b9b319..f9f29ce23c 100644
--- a/packages/effect/test/Context.test.ts
+++ b/packages/effect/test/Context.test.ts
@@ -19,11 +19,20 @@ interface C {
 }
 const C = Context.Tag<C>("C")
 
+class D extends Context.TagClass("D")<D, { readonly d: number }>() {}
+
 describe("Context", () => {
   it("Tag.toJson()", () => {
     const json: any = A.toJSON()
     expect(json["_id"]).toEqual("Tag")
-    expect(json["identifier"]).toEqual(undefined)
+    expect(json["key"]).toEqual("A")
+    expect(typeof json["stack"]).toEqual("string")
+  })
+
+  it("TagClass.toJson()", () => {
+    const json: any = D.toJSON()
+    expect(json["_id"]).toEqual("Tag")
+    expect(json["key"]).toEqual("D")
     expect(typeof json["stack"]).toEqual("string")
   })
 
@@ -51,7 +60,8 @@ describe("Context", () => {
   it("adds and retrieve services", () => {
     const Services = pipe(
       Context.make(A, { a: 0 }),
-      Context.add(B, { b: 1 })
+      Context.add(B, { b: 1 }),
+      Context.add(D, { d: 2 })
     )
 
     expect(Context.get(Services, A)).toEqual({ a: 0 })
@@ -61,6 +71,11 @@ describe("Context", () => {
       Context.getOption(B)
     )).toEqual(O.some({ b: 1 }))
 
+    expect(pipe(
+      Services,
+      Context.get(D)
+    )).toEqual({ d: 2 })
+
     expect(pipe(
       Services,
       Context.getOption(C)
@@ -211,6 +226,14 @@ describe("Context", () => {
           new RegExp(/Error: Service not found: C \(defined at (.*)Context.test.ts:20:19\)/)
         )
       }
+      try {
+        Context.get(Context.empty(), D as never)
+      } catch (e) {
+        assert.match(
+          String(e),
+          new RegExp(/Error: Service not found: D \(defined at (.*)Context.test.ts:22:37\)/)
+        )
+      }
     }
   })