fix: extract references from .gql.ts files

src/definition-provider/reference-provider-for-schema.ts

@@ -1,39 +1,130 @@
+import { globStream } from 'fast-glob'
 import * as gql from 'graphql'
-import { Position, Range } from '../diff'
+import { Location, Position } from '../diff'
 import { getGQLNodeRange, getGQLNodeRangeWithoutDescription, makeQueryParsable } from '../gql'
 import { isPositionWithInRange } from '../position/is-position-within-range'
+import { parseGraphQLDocumentFromTS, readTSFile } from '../ts'
+
+interface SelectedField {
+  parent: string
+  name: string
+}
 
 function isInRange(node: gql.ASTNode, position: Position, offset?: Position) {
   const nodeRange = getGQLNodeRange(node, offset)
   return isPositionWithInRange(position, nodeRange, true)
 }
 
-export function provideReferenceForSchema(source: string, position: Position) {
+function referencesByType(
+  document: gql.DocumentNode,
+  sourcePath: string,
+  type: string | undefined,
+) {
+  if (!type) return []
+  const locations: Location[] = []
+  gql.visit(document, {
+    NamedType(node) {
+      if (node.name.value !== type) return
+      locations.push(new Location(sourcePath, getGQLNodeRangeWithoutDescription(node)))
+    },
+  })
+  return locations
+}
+
+function processFile(file: string, field: SelectedField, schema: gql.GraphQLSchema) {
+  const locations: Location[] = []
+  try {
+    const sourceFile = readTSFile(file)
+    const { document, offset } = parseGraphQLDocumentFromTS(sourceFile)
+    if (!document) return locations
+    const typeInfo = new gql.TypeInfo(schema)
+    gql.visit(
+      document,
+      gql.visitWithTypeInfo(typeInfo, {
+        Field(node) {
+          if (node.name.value !== field.name) return
+          const parent = typeInfo.getParentType()
+          if (!parent || parent.name !== field.parent) return
+          locations.push(new Location(file, getGQLNodeRange(node.name, offset)))
+        },
+      }),
+    )
+  } catch (e) {
+    console.log(`Failed to process file ${file}`, e.message)
+  }
+  return locations
+}
+
+async function referencesByField(
+  schema: gql.GraphQLSchema,
+  field: SelectedField | undefined,
+  pattern: string,
+) {
+  if (!field || !pattern) return []
+  const stream = globStream(pattern)
+  let locations: Location[] = []
+  for await (const file of stream) {
+    locations = locations.concat(processFile(file as string, field, schema))
+  }
+  return locations
+}
+
+export async function provideReferenceForSchema(
+  source: string,
+  sourcePath: string,
+  position: Position,
+  pattern: string,
+) {
   try {
     const fixed = makeQueryParsable(source)
     const document = gql.parse(fixed)
-    let selectedName: string | undefined
-    const processNode = (node: gql.NameNode | gql.NamedTypeNode) => {
+    let type: string | undefined
+    let selectedField: SelectedField | undefined
+    const processNode = (node: gql.NameNode) => {
       if (!isInRange(node, position)) return
-      selectedName = node.kind === gql.Kind.NAME ? node.value : node.name.value
+      type = node.value
+      return gql.BREAK
+    }
+    const processField = (node: gql.TypeDefinitionNode) => {
+      switch (node.kind) {
+        case gql.Kind.OBJECT_TYPE_DEFINITION:
+        case gql.Kind.INPUT_OBJECT_TYPE_DEFINITION:
+        case gql.Kind.INTERFACE_TYPE_DEFINITION:
+          for (const field of node.fields ?? []) {
+            if (isInRange(field.name, position)) {
+              selectedField = { parent: node.name.value, name: field.name.value }
+              return gql.BREAK
+            }
+          }
+      }
     }
     gql.visit(document, {
       Name(node) {
         return processNode(node)
       },
-      NamedType(node) {
-        return processNode(node)
+      EnumTypeDefinition(node) {
+        for (const value of node.values ?? []) {
+          if (isInRange(value.name, position)) {
+            selectedField = { parent: node.name.value, name: value.name.value }
+            return
+          }
+        }
       },
-    })
-    if (!selectedName) return []
-    const ranges: Range[] = []
-    gql.visit(document, {
-      NamedType(node) {
-        if (node.name.value !== selectedName) return
-        ranges.push(getGQLNodeRangeWithoutDescription(node))
+      ObjectTypeDefinition(node) {
+        return processField(node)
+      },
+      InterfaceTypeDefinition(node) {
+        return processField(node)
+      },
+      InputObjectTypeDefinition(node) {
+        return processField(node)
       },
     })
-    return ranges
+    if (type) {
+      return referencesByType(document, sourcePath, type)
+    }
+    const schema = gql.buildSchema(source)
+    return await referencesByField(schema, selectedField, pattern)
   } catch (e) {
     console.error(e)
     return []

src/diff/token.ts

@@ -22,6 +22,7 @@ export class Range {
   constructor(
     public start: Position,
     public end: Position,
+    public path?: string,
   ) {}
 
   setStart(start: Position) {
@@ -39,6 +40,17 @@ export class Range {
   }
 }
 
+export class Location {
+  constructor(
+    public path: string,
+    public range: Range,
+  ) {}
+
+  clone() {
+    return new Location(this.path, this.range.clone())
+  }
+}
+
 export class Token {
   constructor(
     public index: number,

src/ts/parse-graphql-document-from-ts.ts

@@ -12,7 +12,7 @@ export interface ParseOptions {
   fix?: boolean
 }
 
-export interface GraphQLContext {
+export interface GraphQLDocumentContext {
   sourceFile: ts.SourceFile
   variable?: ts.VariableDeclaration
   source?: string
@@ -23,7 +23,7 @@ export interface GraphQLContext {
 export function parseGraphQLDocumentFromTS(
   sourceFile: ts.SourceFile,
   options?: ParseOptions,
-): GraphQLContext {
+): GraphQLDocumentContext {
   const variable = getGraphQLQueryVariable(sourceFile)
   if (!variable) return { sourceFile }