Circular reference fixes (#3451)

fix #3447 
Fixed a few issues with circular references in the JSON Schema emitter
and emitter framework:

* The emitter framework wouldn't restore context correctly when directly
emitting a type reference to a type with a circular reference.
* The JSON Schema emitter did not handle circular references involving
non-JSON Schema types.
* The JSON Schema emitter would create an infinite loop when circular
references needed to be put into $defs.

---------

Co-authored-by: Timothee Guerin <timothee.guerin@outlook.com>

.chronus/changes/json-schema-fixes-2024-4-25-0-58-10.md

@@ -0,0 +1,7 @@
+---
+changeKind: fix
+packages:
+  - "@typespec/compiler"
+---
+
+Emitter framework: fix losing context when referencing circular types

.chronus/changes/json-schema-fixes-2024-4-25-0-58-28.md

@@ -0,0 +1,6 @@
+---
+changeKind: internal
+packages:
+  - "@typespec/json-schema"
+---
+

packages/compiler/src/emitter-framework/asset-emitter.ts

@@ -733,8 +733,8 @@ export function createAssetEmitter<T, TOptions extends object>(
    * Invoke the callback with the given context.
    */
   function withContext(newContext: EmitterState, cb: () => void) {
-    const oldContext = newContext.context;
-    const oldTypeStack = newContext.lexicalTypeStack;
+    const oldContext = context;
+    const oldTypeStack = lexicalTypeStack;
     context = newContext.context;
     lexicalTypeStack = newContext.lexicalTypeStack;
 

packages/compiler/test/emitter-framework/context.test.ts

@@ -504,4 +504,50 @@ describe("emitter-framework: emitter context", () => {
       });
     });
   });
+
+  describe("instantiation context", () => {
+    it("restores  context when after referencing a type with a circular reference", async () => {
+      class Emitter extends CodeTypeEmitter {
+        programContext(program: Program): Context {
+          return {
+            scope: this.emitter.createSourceFile("foo.txt").globalScope,
+          };
+        }
+        modelDeclarationContext(model: Model, name: string): Context {
+          return {
+            inModel: name,
+          };
+        }
+        modelDeclaration(model: Model, name: string): EmitterOutput<string> {
+          super.modelDeclaration(model, name);
+          return this.emitter.result.declaration(name, "Declaration for model " + name);
+        }
+        modelPropertyLiteralContext(property: ModelProperty): Context {
+          return { inProp: property.name };
+        }
+        modelPropertyLiteral(property: ModelProperty): EmitterOutput<string> {
+          const beforeContext = this.emitter.getContext();
+          const res = super.modelPropertyLiteral(property);
+          assert.deepStrictEqual(beforeContext, this.emitter.getContext());
+          return res;
+        }
+      }
+
+      await emitTypeSpec(
+        Emitter,
+        `
+        model A {
+          a: B;
+        }
+  
+        model B {
+          b: B;
+        }
+        
+        `,
+        {},
+        false
+      );
+    });
+  });
 });

packages/json-schema/src/json-schema-emitter.ts

@@ -96,7 +96,6 @@ export class JsonSchemaEmitter extends TypeEmitter<Record<string, any>, JSONSche
     }
 
     this.#applyConstraints(model, schema);
-
     return this.#createDeclaration(model, name, schema);
   }
 
@@ -371,7 +370,7 @@ export class JsonSchemaEmitter extends TypeEmitter<Record<string, any>, JSONSche
         // into the defs of the root schema which references this schema.
         return { $ref: "#/$defs/" + targetDeclaration.name };
       } else {
-        // referencing a schema that is in a source file, but doesn't have an id.
+        // referencing a schema that is in a different source file, but doesn't have an id.
         // nb: this may be dead code.
 
         // when either targetSfScope or currentSfScope are undefined, we have a common scope
@@ -385,7 +384,17 @@ export class JsonSchemaEmitter extends TypeEmitter<Record<string, any>, JSONSche
       }
     }
 
-    throw new Error("$ref to $defs not yet supported.");
+    if (!currentSfScope && !targetSfScope) {
+      // referencing ourself, and we don't have an id (otherwise we'd have
+      // returned that above), so just return a ref.
+      // This should be accurate because the only case when this can happen is if
+      // the target declaration is not a root schema, and so it will be present in
+      // the defs of some root schema, and there is only one level of defs.
+
+      return { $ref: "#/$defs/" + targetDeclaration.name };
+    }
+
+    throw new Error("JSON Pointer refs to arbitrary schemas is not supported");
   }
 
   scalarInstantiation(
@@ -719,13 +728,19 @@ export class JsonSchemaEmitter extends TypeEmitter<Record<string, any>, JSONSche
     const decls = sourceFile.globalScope.declarations;
     compilerAssert(decls.length === 1, "Multiple decls in single schema per file mode");
     const content: Record<string, any> = { ...decls[0].value };
+    const bundledDecls = new Set<Declaration<object>>();
 
     if (sourceFile.meta.bundledRefs.length > 0) {
       // bundle any refs, including refs of refs
       content.$defs = {};
       const refsToBundle: Declaration<object>[] = [...sourceFile.meta.bundledRefs];
       while (refsToBundle.length > 0) {
         const decl = refsToBundle.shift()!;
+        if (bundledDecls.has(decl)) {
+          // already $def'd, no need to def it again.
+          continue;
+        }
+        bundledDecls.add(decl);
         content.$defs[decl.name] = decl.value;
 
         // all scopes are source file scopes in this emitter

packages/json-schema/test/references.test.ts

@@ -196,4 +196,28 @@ describe("referencing non-JSON Schema types", () => {
     assert(amSchemas["B.json"] !== undefined);
     assert(Object.keys(amSchemas).length === 2);
   });
+
+  it("handles circular refs among non-json schema types", async () => {
+    const schemas = await emitSchema(
+      `
+      @jsonSchema
+      model R {
+        test: A;
+      }
+      
+      model A {
+        a: B;
+      }
+      
+      model B {
+        a: A;
+      }
+    `,
+      {},
+      { emitNamespace: false, emitTypes: ["R"] }
+    );
+    assert(schemas["R.json"] !== undefined);
+    assert(Object.keys(schemas).length === 1);
+    assert(Object.keys(schemas["R.json"].$defs).length === 2);
+  });
 });