feat: preserve original ref info (#11)

* feat: preserve original ref info

* fix: update MirroredRegularNode

* chore: incorporate PR feedback

* feat: preserve originalFragment for combined nodes, too

* feat: preserve originalFragment for combined nodes, too

* feat: preserve originalFragment for mirrored nodes

* style: spacing

Co-authored-by: Jakub Rożek <[email protected]>

Author: Marcell Toth

src/__tests__/tree.spec.ts

@@ -3,6 +3,8 @@ import * as fs from 'fs';
 import type { JSONSchema4 } from 'json-schema';
 import * as path from 'path';
 
+import { isRegularNode } from '../guards';
+import type { RegularNode } from '../nodes';
 import { SchemaTree } from '../tree';
 import { printTree } from './utils/printTree';
 
@@ -102,6 +104,27 @@ describe('SchemaTree', () => {
         `);
       });
 
+      it('preserves the original $ref info', () => {
+        const schema: JSONSchema4 = {
+          type: 'object',
+          properties: {
+            foo: {
+              $ref: '#/properties/bar',
+            },
+            bar: {
+              type: 'boolean',
+            },
+          },
+        };
+
+        const tree = new SchemaTree(schema);
+        tree.populate();
+
+        const topLevelObject = tree.root.children[0] as RegularNode;
+        const fooObj = topLevelObject.children!.find(child => child.path[child.path.length - 1] === 'foo')!;
+        expect(isRegularNode(fooObj) && fooObj.originalFragment.$ref).toBe('#/properties/bar');
+      });
+
       it('given an array with $reffed items, should resolve', () => {
         const schema: JSONSchema4 = {
           type: 'object',

src/mergers/mergeOneOrAnyOf.ts

@@ -28,14 +28,13 @@ export function mergeOneOrAnyOf(
       const prunedSchema = { ...fragment };
       delete prunedSchema[combiner];
 
-      const resolvedItem =
-        typeof item.$ref === 'string' && walkingOptions.resolveRef !== null
-          ? walkingOptions.resolveRef(null, item.$ref)
-          : item;
-
       if (Object.keys(prunedSchema).length === 0) {
-        merged.push(resolvedItem);
+        merged.push(item);
       } else {
+        const resolvedItem =
+          typeof item.$ref === 'string' && walkingOptions.resolveRef !== null
+            ? walkingOptions.resolveRef(null, item.$ref)
+            : item;
         const mergedSchema = {
           allOf: [prunedSchema, resolvedItem],
         };

src/nodes/RegularNode.ts

@@ -31,8 +31,9 @@ export class RegularNode extends BaseNode {
   public readonly meta: Readonly<Partial<Dictionary<unknown, SchemaMeta>>>;
   public readonly annotations: Readonly<Partial<Dictionary<unknown, SchemaAnnotations>>>;
   public readonly validations: Readonly<Dictionary<unknown>>;
+  public readonly originalFragment: SchemaFragment;
 
-  constructor(public readonly fragment: SchemaFragment) {
+  constructor(public readonly fragment: SchemaFragment, context?: { originalFragment?: SchemaFragment }) {
     super(fragment);
 
     this.types = getTypes(fragment);
@@ -48,6 +49,7 @@ export class RegularNode extends BaseNode {
     this.meta = getMeta(fragment);
     this.annotations = getAnnotations(fragment);
     this.validations = getValidations(fragment, this.types);
+    this.originalFragment = context?.originalFragment ?? fragment;
 
     this.children = void 0;
   }

src/nodes/mirrored/MirroredRegularNode.ts

@@ -1,6 +1,7 @@
 import type { Dictionary } from '@stoplight/types';
 
 import { isRegularNode } from '../../guards';
+import type { SchemaFragment } from '../../types';
 import { isNonNullable } from '../../utils';
 import { BaseNode } from '../BaseNode';
 import type { ReferenceNode } from '../ReferenceNode';
@@ -22,14 +23,16 @@ export class MirroredRegularNode extends BaseNode implements RegularNode {
   public readonly meta!: Readonly<Partial<Dictionary<unknown, SchemaMeta>>>;
   public readonly annotations!: Readonly<Partial<Dictionary<unknown, SchemaAnnotations>>>;
   public readonly validations!: Readonly<Dictionary<unknown>>;
+  public readonly originalFragment!: SchemaFragment;
 
   public readonly simple!: boolean;
   public readonly unknown!: boolean;
 
   private readonly cache: WeakMap<RegularNode | ReferenceNode, MirroredRegularNode | MirroredReferenceNode>;
 
-  constructor(public readonly mirroredNode: RegularNode) {
+  constructor(public readonly mirroredNode: RegularNode, context?: { originalFragment?: SchemaFragment }) {
     super(mirroredNode.fragment);
+    this.originalFragment = context?.originalFragment ?? mirroredNode.originalFragment;
 
     this.cache = new WeakMap();
 

src/walker/walker.ts

@@ -237,11 +237,14 @@ export class Walker extends EventEmitter<WalkerEmitter> {
     this.schemaNode = schemaNode;
   }
 
-  protected retrieveFromFragment(fragment: SchemaFragment): MirroredSchemaNode | void {
+  protected retrieveFromFragment(
+    fragment: SchemaFragment,
+    originalFragment: SchemaFragment,
+  ): MirroredSchemaNode | void {
     const processedSchemaNode = this.processedFragments.get(fragment);
     if (processedSchemaNode !== void 0) {
       if (isRegularNode(processedSchemaNode)) {
-        return new MirroredRegularNode(processedSchemaNode);
+        return new MirroredRegularNode(processedSchemaNode, { originalFragment });
       }
 
       if (isReferenceNode(processedSchemaNode)) {
@@ -254,10 +257,10 @@ export class Walker extends EventEmitter<WalkerEmitter> {
   }
 
   protected processFragment(): SchemaNode {
-    const { walkingOptions, path } = this;
+    const { walkingOptions, path, fragment: originalFragment } = this;
     let { fragment } = this;
 
-    let retrieved = isNonNullable(fragment) ? this.retrieveFromFragment(fragment) : null;
+    let retrieved = isNonNullable(fragment) ? this.retrieveFromFragment(fragment, originalFragment) : null;
 
     if (retrieved) {
       return retrieved;
@@ -291,25 +294,23 @@ export class Walker extends EventEmitter<WalkerEmitter> {
       try {
         const merged = mergeOneOrAnyOf(fragment, path, walkingOptions);
         if (merged.length === 1) {
-          return new RegularNode(merged[0]);
+          return new RegularNode(merged[0], { originalFragment });
         } else {
           const combiner = SchemaCombinerName.OneOf in fragment ? SchemaCombinerName.OneOf : SchemaCombinerName.AnyOf;
-          return new RegularNode({
-            [combiner]: merged,
-          });
+          return new RegularNode({ [combiner]: merged }, { originalFragment });
         }
       } catch (ex) {
         super.emit('error', createMagicError(new MergingError(ex?.message ?? 'Unknown merging error')));
         // no the end of the world - we will render raw unprocessed fragment
       }
     }
 
-    retrieved = isNonNullable(fragment) ? this.retrieveFromFragment(fragment) : null;
+    retrieved = isNonNullable(fragment) ? this.retrieveFromFragment(fragment, originalFragment) : null;
 
     if (retrieved) {
       return retrieved;
     }
 
-    return new RegularNode(fragment);
+    return new RegularNode(fragment, { originalFragment });
   }
 }