[WIP] Sketch of derived context

Summary: A sketch of derived contexts as described in #159

Not sure this is how the implemenetaiton should work. Was just focusing on getting things working end to end.

If we go this way, I'd want to focus a bit more on internal architecture as well as error handling.

Test Plan:

ghstack-source-id: 0513dda
Pull Request resolved: #161

Author: captbaritone

src/Extractor.ts

@@ -82,6 +82,7 @@ export type ExtractionSnapshot = {
   readonly definitions: DefinitionNode[];
   readonly unresolvedNames: Map<ts.EntityName, NameNode>;
   readonly nameDefinitions: Map<ts.DeclarationStatement, NameDefinition>;
+  readonly implicitNameDefinitions: Map<NameDefinition, ts.TypeReferenceNode>;
   readonly typesWithTypename: Set<string>;
   readonly interfaceDeclarations: Array<ts.InterfaceDeclaration>;
 };
@@ -113,6 +114,8 @@ class Extractor {
   // Snapshot data
   unresolvedNames: Map<ts.EntityName, NameNode> = new Map();
   nameDefinitions: Map<ts.DeclarationStatement, NameDefinition> = new Map();
+  implicitNameDefinitions: Map<NameDefinition, ts.TypeReferenceNode> =
+    new Map();
   typesWithTypename: Set<string> = new Set();
   interfaceDeclarations: Array<ts.InterfaceDeclaration> = [];
 
@@ -132,6 +135,7 @@ class Extractor {
     name: NameNode,
     kind: NameDefinition["kind"],
   ): void {
+    // @ts-ignore FIXME
     this.nameDefinitions.set(node, { name, kind });
   }
 
@@ -218,8 +222,12 @@ class Extractor {
           if (!ts.isDeclarationStatement(node)) {
             this.report(tag, E.contextTagOnNonDeclaration());
           } else {
-            const name = this.gql.name(tag, "CONTEXT_DUMMY_NAME");
-            this.recordTypeName(node, name, "CONTEXT");
+            if (ts.isFunctionDeclaration(node)) {
+              this.recordDerivedContext(node, tag);
+            } else {
+              const name = this.gql.name(tag, "CONTEXT_DUMMY_NAME");
+              this.recordTypeName(node, name, "CONTEXT");
+            }
           }
           break;
         }
@@ -293,11 +301,45 @@ class Extractor {
       definitions: this.definitions,
       unresolvedNames: this.unresolvedNames,
       nameDefinitions: this.nameDefinitions,
+      implicitNameDefinitions: this.implicitNameDefinitions,
       typesWithTypename: this.typesWithTypename,
       interfaceDeclarations: this.interfaceDeclarations,
     });
   }
 
+  recordDerivedContext(node: ts.FunctionDeclaration, tag: ts.JSDocTag) {
+    const returnType = node.type;
+    if (returnType == null) {
+      throw new Error("Function declaration must have a return type");
+    }
+    if (!ts.isTypeReferenceNode(returnType)) {
+      throw new Error("Function declaration must return an explicit type");
+    }
+
+    const funcName = this.namedFunctionExportName(node);
+
+    if (!ts.isSourceFile(node.parent)) {
+      return this.report(node, E.functionFieldNotTopLevel());
+    }
+
+    const tsModulePath = relativePath(node.getSourceFile().fileName);
+
+    const paramResults = this.resolverParams(node.parameters);
+    if (paramResults == null) return null;
+
+    const name = this.gql.name(tag, "CONTEXT_DUMMY_NAME");
+    this.implicitNameDefinitions.set(
+      {
+        kind: "DERIVED_CONTEXT",
+        name,
+        path: tsModulePath,
+        exportName: funcName?.text ?? null,
+        args: paramResults.resolverParams,
+      },
+      returnType,
+    );
+  }
+
   extractType(node: ts.Node, tag: ts.JSDocTag) {
     if (ts.isClassDeclaration(node)) {
       this.typeClassDeclaration(node, tag);

src/TypeContext.ts

@@ -11,26 +11,40 @@ import {
   DiagnosticResult,
   tsErr,
   gqlRelated,
+  DiagnosticsResult,
+  FixableDiagnosticWithLocation,
+  tsRelated,
 } from "./utils/DiagnosticError";
 import { err, ok } from "./utils/Result";
 import * as E from "./Errors";
 import { ExtractionSnapshot } from "./Extractor";
+import { ResolverArgument } from "./resolverSignature";
 
 export const UNRESOLVED_REFERENCE_NAME = `__UNRESOLVED_REFERENCE__`;
 
-export type NameDefinition = {
+export type DerivedResolverDefinition = {
   name: NameNode;
-  kind:
-    | "TYPE"
-    | "INTERFACE"
-    | "UNION"
-    | "SCALAR"
-    | "INPUT_OBJECT"
-    | "ENUM"
-    | "CONTEXT"
-    | "INFO";
+  path: string;
+  exportName: string | null;
+  args: ResolverArgument[];
+  kind: "DERIVED_CONTEXT";
 };
 
+export type NameDefinition =
+  | {
+      name: NameNode;
+      kind:
+        | "TYPE"
+        | "INTERFACE"
+        | "UNION"
+        | "SCALAR"
+        | "INPUT_OBJECT"
+        | "ENUM"
+        | "CONTEXT"
+        | "INFO";
+    }
+  | DerivedResolverDefinition;
+
 type TsIdentifier = number;
 
 /**
@@ -55,15 +69,40 @@ export class TypeContext {
   static fromSnapshot(
     checker: ts.TypeChecker,
     snapshot: ExtractionSnapshot,
-  ): TypeContext {
+  ): DiagnosticsResult<TypeContext> {
+    const errors: FixableDiagnosticWithLocation[] = [];
     const self = new TypeContext(checker);
     for (const [node, typeName] of snapshot.unresolvedNames) {
       self._markUnresolvedType(node, typeName);
     }
     for (const [node, definition] of snapshot.nameDefinitions) {
-      self._recordTypeName(node, definition.name, definition.kind);
+      self._recordTypeName(node, definition);
+    }
+    for (const [definition, reference] of snapshot.implicitNameDefinitions) {
+      const declaration = self.maybeTsDeclarationForTsName(reference.typeName);
+      if (declaration == null) {
+        errors.push(tsErr(reference.typeName, E.unresolvedTypeReference()));
+        continue;
+      }
+      const existing = self._declarationToName.get(declaration);
+      if (existing != null) {
+        errors.push(
+          // TODO: Better error messages here
+          tsErr(declaration, "Duplicate derived contexts for given type", [
+            tsRelated(reference, "One was defined here"),
+            gqlRelated(existing.name, "Other here"),
+          ]),
+        );
+        continue;
+      }
+
+      self._recordTypeName(declaration, definition);
+    }
+
+    if (errors.length > 0) {
+      return err(errors);
     }
-    return self;
+    return ok(self);
   }
 
   constructor(checker: ts.TypeChecker) {
@@ -72,13 +111,9 @@ export class TypeContext {
 
   // Record that a GraphQL construct of type `kind` with the name `name` is
   // declared at `node`.
-  private _recordTypeName(
-    node: ts.Declaration,
-    name: NameNode,
-    kind: NameDefinition["kind"],
-  ) {
-    this._idToDeclaration.set(name.tsIdentifier, node);
-    this._declarationToName.set(node, { name, kind });
+  private _recordTypeName(node: ts.Declaration, definition: NameDefinition) {
+    this._idToDeclaration.set(definition.name.tsIdentifier, node);
+    this._declarationToName.set(node, definition);
   }
 
   // Record that a type references `node`

src/codegen/TSAstBuilder.ts

@@ -9,6 +9,7 @@ const F = ts.factory;
  * A helper class to build up a TypeScript document AST.
  */
 export default class TSAstBuilder {
+  _globalNames: Map<string, number> = new Map();
   _imports: ts.Statement[] = [];
   imports: Map<string, { name: string; as?: string }[]> = new Map();
   _helpers: ts.Statement[] = [];
@@ -209,7 +210,21 @@ export default class TSAstBuilder {
       sourceFile,
     );
   }
+
+  // Given a desired name in the module scope, return a name that is unique. If
+  // the name is already taken, a suffix will be added to the name to make it
+  // unique.
+  //
+  // NOTE: This is not truly unique, as it only checks the names that have been
+  // generated through this method. In the future we could add more robust
+  // scope/name tracking.
+  getUniqueName(name: string): string {
+    const count = this._globalNames.get(name) ?? 0;
+    this._globalNames.set(name, count + 1);
+    return count === 0 ? name : `${name}_${count}`;
+  }
 }
+
 function replaceExt(filePath: string, newSuffix: string): string {
   const ext = path.extname(filePath);
   return filePath.slice(0, -ext.length) + newSuffix;

src/codegen/resolverCodegen.ts

@@ -20,6 +20,7 @@ const F = ts.factory;
  */
 export default class ResolverCodegen {
   _helpers: Set<string> = new Set();
+  _derivedContextNames: Map<string, string> = new Map();
   constructor(public ts: TSAstBuilder, public _resolvers: Metadata) {}
   resolveMethod(
     fieldName: string,
@@ -178,11 +179,36 @@ export default class ResolverCodegen {
           F.createIdentifier("args"),
           F.createIdentifier(arg.name),
         );
+      case "derivedContext": {
+        const localName = this.getDerivedContextName(arg.path, arg.exportName);
+        this.ts.importUserConstruct(arg.path, arg.exportName, localName);
+        return F.createCallExpression(
+          F.createIdentifier(localName),
+          undefined,
+          arg.args.map((arg) => this.resolverParam(arg)),
+        );
+      }
+
       default:
         // @ts-expect-error
         throw new Error(`Unexpected resolver kind ${arg.kind}`);
     }
   }
+
+  // Derived contexts are not anchored to anything that we know to be
+  // globally unique, like GraphQL type names, so must ensure this name is
+  // unique within our module. However, we want to avoid generating a new
+  // name for the same derived context more than once.
+  getDerivedContextName(path: string, exportName: string | null): string {
+    const key = `${path}:${exportName ?? ""}`;
+    let name = this._derivedContextNames.get(key);
+    if (name == null) {
+      name = this.ts.getUniqueName(exportName ?? "deriveContext");
+      this._derivedContextNames.set(key, name);
+    }
+    return name;
+  }
+
   // If a field is smantically non-null, we need to wrap the resolver in a
   // runtime check to ensure that the resolver does not return null.
   maybeApplySemanticNullRuntimeCheck(

src/lib.ts

@@ -89,7 +89,11 @@ export function extractSchemaAndDoc(
       const { typesWithTypename } = snapshot;
       const config = options.raw.grats;
       const checker = program.getTypeChecker();
-      const ctx = TypeContext.fromSnapshot(checker, snapshot);
+      const ctxResult = TypeContext.fromSnapshot(checker, snapshot);
+      if (ctxResult.kind === "ERROR") {
+        return ctxResult;
+      }
+      const ctx = ctxResult.value;
 
       // Collect validation errors
       const validationResult = concatResults(
@@ -173,6 +177,7 @@ function combineSnapshots(snapshots: ExtractionSnapshot[]): ExtractionSnapshot {
   const result: ExtractionSnapshot = {
     definitions: [],
     nameDefinitions: new Map(),
+    implicitNameDefinitions: new Map(),
     unresolvedNames: new Map(),
     typesWithTypename: new Set(),
     interfaceDeclarations: [],
@@ -191,6 +196,10 @@ function combineSnapshots(snapshots: ExtractionSnapshot[]): ExtractionSnapshot {
       result.unresolvedNames.set(node, typeName);
     }
 
+    for (const [node, definition] of snapshot.implicitNameDefinitions) {
+      result.implicitNameDefinitions.set(node, definition);
+    }
+
     for (const typeName of snapshot.typesWithTypename) {
       result.typesWithTypename.add(typeName);
     }

src/metadata.ts

@@ -82,11 +82,14 @@ export type StaticMethodResolver = {
   arguments: ResolverArgument[] | null;
 };
 
+export type ContextArgs = ContextArgument | DerivedContextArgument;
+
 /** An argument expected by a resolver function or method */
 export type ResolverArgument =
   | SourceArgument
   | ArgumentsObjectArgument
   | ContextArgument
+  | DerivedContextArgument
   | InformationArgument
   | NamedArgument;
 
@@ -105,6 +108,14 @@ export type ContextArgument = {
   kind: "context";
 };
 
+/** A context value which is expressed as a function of the global context */
+export type DerivedContextArgument = {
+  kind: "derivedContext";
+  path: string; // Path to the module
+  exportName: string | null; // Export name. If omitted, the class is the default export
+  args: Array<ContextArgs>;
+};
+
 /** The GraphQL info object */
 export type InformationArgument = {
   kind: "information";

src/resolverSignature.ts

@@ -60,6 +60,14 @@ export type ContextResolverArgument = {
   node: ts.Node;
 };
 
+export type DerivedContextResolverArgument = {
+  kind: "derivedContext";
+  path: string;
+  exportName: string | null;
+  args: Array<DerivedContextResolverArgument | ContextResolverArgument>;
+  node: ts.Node;
+};
+
 export type InformationResolverArgument = {
   kind: "information";
   node: ts.Node;
@@ -82,6 +90,7 @@ export type ResolverArgument =
   | SourceResolverArgument
   | ArgumentsObjectResolverArgument
   | ContextResolverArgument
+  | DerivedContextResolverArgument
   | InformationResolverArgument
   | NamedResolverArgument
   | UnresolvedResolverArgument;