replace legacy mechanism for looking up key name from definitions

src/annotator.ts

@@ -1,41 +1,78 @@
 import {isPlainObject} from 'lodash'
 import {DereferencedPaths} from './resolver'
-import {AnnotatedJSONSchema, JSONSchema, Parent, isAnnotated} from './types/JSONSchema'
+import {AnnotatedJSONSchema, JSONSchema, Parent, KeyNameFromDefinition, isAnnotated} from './types/JSONSchema'
 
-/**
- * Traverses over the schema, assigning to each
- * node metadata that will be used downstream.
- */
-export function annotate(
-  schema: JSONSchema,
-  dereferencedPaths: DereferencedPaths,
-  parent: JSONSchema | null = null,
-): AnnotatedJSONSchema {
-  if (!Array.isArray(schema) && !isPlainObject(schema)) {
-    return schema as AnnotatedJSONSchema
+const annotators = new Map<
+  string,
+  (schema: JSONSchema, parent: JSONSchema | null, dereferencedPaths: DereferencedPaths) => void
+>()
+
+annotators.set('Add Parent to each node', (schema, parent) => {
+  Object.defineProperty(schema, Parent, {
+    enumerable: false,
+    value: parent,
+    writable: false,
+  })
+})
+
+annotators.set('Add key name from the original definition, if there was one', (schema, _, dereferencedPaths) => {
+  const ref = dereferencedPaths.get(schema)
+  if (!ref) {
+    return
   }
 
-  // Handle cycles
-  if (isAnnotated(schema)) {
-    return schema
+  const keyName = getDefinitionKeyNameFromRef(ref)
+  if (!keyName) {
+    return
   }
 
-  // Add a reference to this schema's parent
-  Object.defineProperty(schema, Parent, {
+  Object.defineProperty(schema, KeyNameFromDefinition, {
     enumerable: false,
-    value: parent,
+    value: keyName,
     writable: false,
   })
+})
 
-  // Arrays
-  if (Array.isArray(schema)) {
-    schema.forEach(child => annotate(child, dereferencedPaths, schema))
+function getDefinitionKeyNameFromRef(ref: string): string | undefined {
+  const parts = ref.split('/')
+  if (parts[parts.length - 2] !== 'definitions') {
+    return
   }
+  return parts[parts.length - 1]
+}
 
-  // Objects
-  for (const key in schema) {
-    annotate(schema[key], dereferencedPaths, schema)
+/**
+ * Traverses over the schema, assigning to each
+ * node metadata that will be used downstream.
+ */
+export function annotate(schema: JSONSchema, dereferencedPaths: DereferencedPaths): AnnotatedJSONSchema {
+  function go(s: JSONSchema, parent: JSONSchema | null): void {
+    if (!Array.isArray(s) && !isPlainObject(s)) {
+      return
+    }
+
+    // Handle cycles
+    if (isAnnotated(s)) {
+      return
+    }
+
+    // Run annotators
+    annotators.forEach(f => {
+      f(s, parent, dereferencedPaths)
+    })
+
+    // Handle arrays
+    if (Array.isArray(s)) {
+      s.forEach(_ => go(_, s))
+    }
+
+    // Handle objects
+    for (const key in s) {
+      go(s[key], s)
+    }
   }
 
+  go(schema, null)
+
   return schema as AnnotatedJSONSchema
 }

src/parser.ts

@@ -1,5 +1,5 @@
 import {JSONSchema4Type, JSONSchema4TypeName} from 'json-schema'
-import {findKey, includes, isPlainObject, map, memoize, omit} from 'lodash'
+import {includes, map, omit} from 'lodash'
 import {format} from 'util'
 import {Options} from './'
 import {typesOfSchema} from './typesOfSchema'
@@ -17,10 +17,9 @@ import {
 } from './types/AST'
 import {
   EnumJSONSchema,
-  getRootSchema,
   isBoolean,
   isPrimitive,
-  JSONSchemaWithDefinitions,
+  KeyNameFromDefinition,
   NormalizedJSONSchema,
   Parent,
   SchemaSchema,
@@ -143,8 +142,7 @@ function parseNonLiteral(
   processed: Processed,
   usedNames: UsedNames,
 ): AST {
-  const definitions = getDefinitionsMemoized(getRootSchema(schema as any)) // TODO
-  const keyNameFromDefinition = findKey(definitions, _ => _ === schema)
+  const keyNameFromDefinition = schema[KeyNameFromDefinition]
 
   switch (type) {
     case 'ALL_OF':
@@ -487,47 +485,3 @@ via the \`definition\` "${key}".`
       })
   }
 }
-
-type Definitions = {[k: string]: NormalizedJSONSchema}
-
-function getDefinitions(
-  schema: NormalizedJSONSchema,
-  isSchema = true,
-  processed = new Set<NormalizedJSONSchema>(),
-): Definitions {
-  if (processed.has(schema)) {
-    return {}
-  }
-  processed.add(schema)
-  if (Array.isArray(schema)) {
-    return schema.reduce(
-      (prev, cur) => ({
-        ...prev,
-        ...getDefinitions(cur, false, processed),
-      }),
-      {},
-    )
-  }
-  if (isPlainObject(schema)) {
-    return {
-      ...(isSchema && hasDefinitions(schema) ? schema.$defs : {}),
-      ...Object.keys(schema).reduce<Definitions>(
-        (prev, cur) => ({
-          ...prev,
-          ...getDefinitions(schema[cur], false, processed),
-        }),
-        {},
-      ),
-    }
-  }
-  return {}
-}
-
-const getDefinitionsMemoized = memoize(getDefinitions)
-
-/**
- * TODO: Reduce rate of false positives
- */
-function hasDefinitions(schema: NormalizedJSONSchema): schema is JSONSchemaWithDefinitions {
-  return '$defs' in schema
-}

src/types/JSONSchema.ts

@@ -1,5 +1,5 @@
 import {JSONSchema4, JSONSchema4Type, JSONSchema4TypeName} from 'json-schema'
-import {isPlainObject, memoize} from 'lodash'
+import {isPlainObject} from 'lodash'
 
 export type SchemaType =
   | 'ALL_OF'
@@ -40,10 +40,15 @@ export interface JSONSchema extends JSONSchema4 {
   deprecated?: boolean
 }
 
+export const KeyNameFromDefinition = Symbol('KeyNameFromDefinition')
 export const Parent = Symbol('Parent')
 
 export interface AnnotatedJSONSchema extends JSONSchema {
   [Parent]: AnnotatedJSONSchema
+  /**
+   * The name of the definition from the original $ref that was dereferenced, if there was one.
+   */
+  [KeyNameFromDefinition]?: string
 
   additionalItems?: boolean | AnnotatedJSONSchema
   additionalProperties?: boolean | AnnotatedJSONSchema
@@ -112,24 +117,10 @@ export interface SchemaSchema extends NormalizedJSONSchema {
   required: string[]
 }
 
-export interface JSONSchemaWithDefinitions extends NormalizedJSONSchema {
-  $defs: {
-    [k: string]: NormalizedJSONSchema
-  }
-}
-
 export interface CustomTypeJSONSchema extends NormalizedJSONSchema {
   tsType: string
 }
 
-export const getRootSchema = memoize((schema: NormalizedJSONSchema): NormalizedJSONSchema => {
-  const parent = schema[Parent]
-  if (!parent) {
-    return schema
-  }
-  return getRootSchema(parent)
-})
-
 export function isBoolean(schema: AnnotatedJSONSchema | JSONSchemaType): schema is boolean {
   return schema === true || schema === false
 }