Emitter framework: Allow reference context to be passed to `emitTypeR…

…eference` (#2712) Current way of setting reference context is flawed in that you cannot have a different context to different reference. It is also run before the element itself instead of before emitting the reference so this mean the context would apply to the referrer as well. This allows the `emitTypeReference` to take a parameter to patch the reference context Next step would be to reconsider the `*ReferenceContext` and `*Context` methods to make more sense.
microsoft · Dec 1, 2023 · 4e63cab · 4e63cab
1 parent 0b33308
commit 4e63cab
Show file tree

Hide file tree

Showing 4 changed files with 304 additions and 89 deletions.
diff --git a/common/changes/@typespec/compiler/emitter-framework-ref-context_2023-11-30-18-46.json b/common/changes/@typespec/compiler/emitter-framework-ref-context_2023-11-30-18-46.json
@@ -0,0 +1,10 @@
+{
+  "changes": [
+    {
+      "packageName": "@typespec/compiler",
+      "comment": "[API] Emitter framework: `emitTypeReference` function takes an optional reference context that can be used to patch the context for the target.",
+      "type": "none"
+    }
+  ],
+  "packageName": "@typespec/compiler"
+}
diff --git a/packages/compiler/src/emitter-framework/asset-emitter.ts b/packages/compiler/src/emitter-framework/asset-emitter.ts
@@ -22,6 +22,7 @@ import {
   EmitEntity,
   EmitterResult,
   EmitterState,
+  EmitTypeReferenceOptions,
   LexicalTypeStackEntry,
   NamespaceScope,
   NoEmit,
@@ -205,107 +206,122 @@ export function createAssetEmitter<T, TOptions extends object>(
       return sourceFile;
     },
 
-    emitTypeReference(target): EmitEntity<T> {
-      if (target.kind === "ModelProperty") {
-        return invokeTypeEmitter("modelPropertyReference", target);
-      } else if (target.kind === "EnumMember") {
-        return invokeTypeEmitter("enumMemberReference", target);
-      }
-
-      const oldIncomingReferenceContext = incomingReferenceContext;
-      const oldIncomingReferenceContextTarget = incomingReferenceContextTarget;
+    emitTypeReference(target, options?: EmitTypeReferenceOptions): EmitEntity<T> {
+      return withPatchedReferenceContext(options?.referenceContext, () => {
+        const oldIncomingReferenceContext = incomingReferenceContext;
+        const oldIncomingReferenceContextTarget = incomingReferenceContextTarget;
 
-      incomingReferenceContext = context.referenceContext ?? null;
-      incomingReferenceContextTarget = incomingReferenceContext ? target : null;
+        incomingReferenceContext = context.referenceContext ?? null;
+        incomingReferenceContextTarget = incomingReferenceContext ? target : null;
 
-      const entity = this.emitType(target);
+        let result;
+        if (target.kind === "ModelProperty") {
+          result = invokeTypeEmitter("modelPropertyReference", target);
+        } else if (target.kind === "EnumMember") {
+          result = invokeTypeEmitter("enumMemberReference", target);
+        }
 
-      incomingReferenceContext = oldIncomingReferenceContext;
-      incomingReferenceContextTarget = oldIncomingReferenceContextTarget;
+        if (result) {
+          incomingReferenceContext = oldIncomingReferenceContext;
+          incomingReferenceContextTarget = oldIncomingReferenceContextTarget;
+          return result;
+        }
 
-      let placeholder: Placeholder<T> | null = null;
+        const entity = this.emitType(target);
 
-      if (entity.kind === "circular") {
-        let waiting = waitingCircularRefs.get(entity.emitEntityKey);
-        if (!waiting) {
-          waiting = [];
-          waitingCircularRefs.set(entity.emitEntityKey, waiting);
-        }
+        incomingReferenceContext = oldIncomingReferenceContext;
+        incomingReferenceContextTarget = oldIncomingReferenceContextTarget;
 
-        const typeChainSnapshot = referenceTypeChain;
-        waiting.push({
-          state: {
-            lexicalTypeStack,
-            context,
-          },
-          cb: (resolvedEntity) =>
-            invokeReference(
-              this,
-              resolvedEntity,
-              true,
-              resolveReferenceCycle(typeChainSnapshot, entity, typeToEmitEntity as any)
-            ),
-        });
+        let placeholder: Placeholder<T> | null = null;
 
-        placeholder = new Placeholder();
-        return this.result.rawCode(placeholder);
-      } else {
-        return invokeReference(this, entity, false);
-      }
+        if (entity.kind === "circular") {
+          let waiting = waitingCircularRefs.get(entity.emitEntityKey);
+          if (!waiting) {
+            waiting = [];
+            waitingCircularRefs.set(entity.emitEntityKey, waiting);
+          }
 
-      function invokeReference(
-        assetEmitter: AssetEmitter<T, TOptions>,
-        entity: EmitEntity<T>,
-        circular: boolean,
-        cycle?: ReferenceCycle
-      ): EmitEntity<T> {
-        let ref;
-        const scope = currentScope();
+          const typeChainSnapshot = referenceTypeChain;
+          waiting.push({
+            state: {
+              lexicalTypeStack,
+              context,
+            },
+            cb: (resolvedEntity) =>
+              invokeReference(
+                this,
+                resolvedEntity,
+                true,
+                resolveReferenceCycle(typeChainSnapshot, entity, typeToEmitEntity as any)
+              ),
+          });
 
-        if (circular) {
-          ref = typeEmitter.circularReference(entity, scope, cycle!);
+          placeholder = new Placeholder();
+          return this.result.rawCode(placeholder);
         } else {
-          if (entity.kind !== "declaration") {
-            return entity;
-          }
-          compilerAssert(
-            scope,
-            "Emit context must have a scope set in order to create references to declarations."
-          );
-          const { pathUp, pathDown, commonScope } = resolveDeclarationReferenceScope(entity, scope);
-          ref = typeEmitter.reference(entity, pathUp, pathDown, commonScope);
+          return invokeReference(this, entity, false);
         }
 
-        if (!(ref instanceof EmitterResult)) {
-          ref = assetEmitter.result.rawCode(ref) as RawCode<T>;
-        }
+        function invokeReference(
+          assetEmitter: AssetEmitter<T, TOptions>,
+          entity: EmitEntity<T>,
+          circular: boolean,
+          cycle?: ReferenceCycle
+        ): EmitEntity<T> {
+          let ref;
+          const scope = currentScope();
+
+          if (circular) {
+            ref = typeEmitter.circularReference(entity, scope, cycle!);
+          } else {
+            if (entity.kind !== "declaration") {
+              return entity;
+            }
+            compilerAssert(
+              scope,
+              "Emit context must have a scope set in order to create references to declarations."
+            );
+            const { pathUp, pathDown, commonScope } = resolveDeclarationReferenceScope(
+              entity,
+              scope
+            );
+            ref = typeEmitter.reference(entity, pathUp, pathDown, commonScope);
+          }
 
-        if (placeholder) {
-          // this should never happen as this function shouldn't be called until
-          // the target declaration is finished being emitted.
-          compilerAssert(ref.kind !== "circular", "TypeEmitter `reference` returned circular emit");
-
-          // this could presumably be allowed if we want.
-          compilerAssert(
-            ref.kind === "none" || !(ref.value instanceof Placeholder),
-            "TypeEmitter's `reference` method cannot return a placeholder."
-          );
-
-          switch (ref.kind) {
-            case "code":
-            case "declaration":
-              placeholder.setValue(ref.value as T);
-              break;
-            case "none":
-              // this cast is incorrect, think about what should happen
-              // if reference returns noEmit...
-              placeholder.setValue("" as T);
-              break;
+          if (!(ref instanceof EmitterResult)) {
+            ref = assetEmitter.result.rawCode(ref) as RawCode<T>;
           }
-        }
 
-        return ref;
-      }
+          if (placeholder) {
+            // this should never happen as this function shouldn't be called until
+            // the target declaration is finished being emitted.
+            compilerAssert(
+              ref.kind !== "circular",
+              "TypeEmitter `reference` returned circular emit"
+            );
+
+            // this could presumably be allowed if we want.
+            compilerAssert(
+              ref.kind === "none" || !(ref.value instanceof Placeholder),
+              "TypeEmitter's `reference` method cannot return a placeholder."
+            );
+
+            switch (ref.kind) {
+              case "code":
+              case "declaration":
+                placeholder.setValue(ref.value as T);
+                break;
+              case "none":
+                // this cast is incorrect, think about what should happen
+                // if reference returns noEmit...
+                placeholder.setValue("" as T);
+                break;
+            }
+          }
+
+          return ref;
+        }
+      });
     },
 
     emitDeclarationName(type): string | undefined {
@@ -696,6 +712,29 @@ export function createAssetEmitter<T, TOptions extends object>(
     referenceTypeChain = oldRefTypeStack;
   }
 
+  function withPatchedReferenceContext<T>(
+    referenceContext: Record<string, any> | undefined,
+    cb: () => T
+  ): T {
+    if (referenceContext !== undefined) {
+      const oldContext = context;
+
+      context = stateInterner.intern({
+        lexicalContext: context.lexicalContext,
+        referenceContext: stateInterner.intern({
+          ...context.referenceContext,
+          ...referenceContext,
+        }),
+      });
+
+      const result = cb();
+      context = oldContext;
+      return result;
+    } else {
+      return cb();
+    }
+  }
+
   /**
    * Invoke the callback with the given context.
    */

diff --git a/packages/compiler/src/emitter-framework/types.ts b/packages/compiler/src/emitter-framework/types.ts
@@ -13,11 +13,16 @@ import {
 } from "../core/index.js";
 import { Placeholder } from "./placeholder.js";
 import { TypeEmitter } from "./type-emitter.js";
+
 type AssetEmitterOptions<TOptions extends object> = {
   noEmit: boolean;
   emitterOutputDir: string;
 } & TOptions;
 
+export interface EmitTypeReferenceOptions {
+  readonly referenceContext?: Record<string, any>;
+}
+
 export interface AssetEmitter<T, TOptions extends object = Record<string, unknown>> {
   /**
    * Get the current emitter context as set by the TypeEmitter's various
@@ -28,7 +33,7 @@ export interface AssetEmitter<T, TOptions extends object = Record<string, unknow
   getContext(): Context;
   getOptions(): AssetEmitterOptions<TOptions>;
   getProgram(): Program;
-  emitTypeReference(type: Type): EmitEntity<T>;
+  emitTypeReference(type: Type, context?: EmitTypeReferenceOptions): EmitEntity<T>;
   emitDeclarationName(type: TypeSpecDeclaration): string | undefined;
   emitType(type: Type, context?: Partial<ContextState>): EmitEntity<T>;
   emitProgram(options?: { emitGlobalNamespace?: boolean; emitTypeSpecNamespace?: boolean }): void;