Derived context values (#164)

* [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

* Document derived context

* Cleanup

Author: captbaritone

src/Errors.ts

@@ -607,3 +607,19 @@ export function noTypesDefined() {
 export function tsConfigNotFound(cwd: string) {
   return `Grats: Could not find \`tsconfig.json\` searching in ${cwd}.\n\nSee https://www.typescriptlang.org/download/ for instructors on how to add TypeScript to your project. Then run \`npx tsc --init\` to create a \`tsconfig.json\` file.`;
 }
+
+export function cyclicDerivedContext() {
+  return `Cyclic dependency detected in derived context. This derived context value depends upon itself.`;
+}
+
+export function invalidDerivedContextArgType() {
+  return "Invalid type for derived context function argument. Derived context functions may only accept other `@gqlContext` types as arguments.";
+}
+
+export function missingReturnTypeForDerivedResolver() {
+  return 'Expected derived resolver to have an explicit return type. This is needed to allow Grats to "see" which type to treat as a derived context type.';
+}
+
+export function derivedResolverInvalidReturnType() {
+  return "Expected derived resolver function's return type to be a type reference. Grats uses this type reference to determine which type to treat as a derived context type.";
+}

src/Extractor.ts

@@ -27,7 +27,11 @@ import {
 } from "./utils/DiagnosticError";
 import { err, ok } from "./utils/Result";
 import * as ts from "typescript";
-import { NameDefinition, UNRESOLVED_REFERENCE_NAME } from "./TypeContext";
+import {
+  DeclarationDefinition,
+  NameDefinition,
+  UNRESOLVED_REFERENCE_NAME,
+} from "./TypeContext";
 import * as E from "./Errors";
 import { traverseJSDocTags } from "./utils/JSDoc";
 import { GraphQLConstructor } from "./GraphQLConstructor";
@@ -86,6 +90,10 @@ export type ExtractionSnapshot = {
   readonly definitions: DefinitionNode[];
   readonly unresolvedNames: Map<ts.EntityName, NameNode>;
   readonly nameDefinitions: Map<ts.DeclarationStatement, NameDefinition>;
+  readonly implicitNameDefinitions: Map<
+    DeclarationDefinition,
+    ts.TypeReferenceNode
+  >;
   readonly typesWithTypename: Set<string>;
   readonly interfaceDeclarations: Array<ts.InterfaceDeclaration>;
 };
@@ -117,6 +125,8 @@ class Extractor {
   // Snapshot data
   unresolvedNames: Map<ts.EntityName, NameNode> = new Map();
   nameDefinitions: Map<ts.DeclarationStatement, NameDefinition> = new Map();
+  implicitNameDefinitions: Map<DeclarationDefinition, ts.TypeReferenceNode> =
+    new Map();
   typesWithTypename: Set<string> = new Set();
   interfaceDeclarations: Array<ts.InterfaceDeclaration> = [];
 
@@ -188,8 +198,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;
         }
@@ -270,6 +284,7 @@ class Extractor {
       definitions: this.definitions,
       unresolvedNames: this.unresolvedNames,
       nameDefinitions: this.nameDefinitions,
+      implicitNameDefinitions: this.implicitNameDefinitions,
       typesWithTypename: this.typesWithTypename,
       interfaceDeclarations: this.interfaceDeclarations,
     });
@@ -329,6 +344,38 @@ class Extractor {
       }
     }
   }
+  recordDerivedContext(node: ts.FunctionDeclaration, tag: ts.JSDocTag) {
+    const returnType = node.type;
+    if (returnType == null) {
+      return this.report(node, E.missingReturnTypeForDerivedResolver());
+    }
+    if (!ts.isTypeReferenceNode(returnType)) {
+      return this.report(returnType, E.missingReturnTypeForDerivedResolver());
+    }
+
+    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)) {

src/TypeContext.ts

@@ -11,13 +11,24 @@ import {
   DiagnosticResult,
   tsErr,
   gqlRelated,
+  DiagnosticsResult,
+  FixableDiagnosticWithLocation,
 } 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 DerivedResolverDefinition = {
+  name: NameNode;
+  path: string;
+  exportName: string | null;
+  args: ResolverArgument[];
+  kind: "DERIVED_CONTEXT";
+};
+
 export type NameDefinition = {
   name: NameNode;
   kind:
@@ -31,6 +42,8 @@ export type NameDefinition = {
     | "INFO";
 };
 
+export type DeclarationDefinition = NameDefinition | DerivedResolverDefinition;
+
 type TsIdentifier = number;
 
 /**
@@ -48,22 +61,50 @@ type TsIdentifier = number;
 export class TypeContext {
   checker: ts.TypeChecker;
 
-  _declarationToName: Map<ts.Declaration, NameDefinition> = new Map();
+  _declarationToDefinition: Map<ts.Declaration, DeclarationDefinition> =
+    new Map();
   _unresolvedNodes: Map<TsIdentifier, ts.EntityName> = new Map();
   _idToDeclaration: Map<TsIdentifier, ts.Declaration> = new Map();
 
   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._recordDeclaration(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._declarationToDefinition.get(declaration);
+      if (existing != null) {
+        errors.push(
+          tsErr(
+            declaration,
+            "Multiple derived contexts defined for given type",
+            [
+              gqlRelated(definition.name, "One was defined here"),
+              gqlRelated(existing.name, "Another here"),
+            ],
+          ),
+        );
+        continue;
+      }
+      self._recordDeclaration(declaration, definition);
+    }
+
+    if (errors.length > 0) {
+      return err(errors);
     }
-    return self;
+    return ok(self);
   }
 
   constructor(checker: ts.TypeChecker) {
@@ -72,22 +113,21 @@ export class TypeContext {
 
   // Record that a GraphQL construct of type `kind` with the name `name` is
   // declared at `node`.
-  private _recordTypeName(
+  private _recordDeclaration(
     node: ts.Declaration,
-    name: NameNode,
-    kind: NameDefinition["kind"],
+    definition: DeclarationDefinition,
   ) {
-    this._idToDeclaration.set(name.tsIdentifier, node);
-    this._declarationToName.set(node, { name, kind });
+    this._idToDeclaration.set(definition.name.tsIdentifier, node);
+    this._declarationToDefinition.set(node, definition);
   }
 
   // Record that a type references `node`
   private _markUnresolvedType(node: ts.EntityName, name: NameNode) {
     this._unresolvedNodes.set(name.tsIdentifier, node);
   }
 
-  allNameDefinitions(): Iterable<NameDefinition> {
-    return this._declarationToName.values();
+  allDefinitions(): Iterable<DeclarationDefinition> {
+    return this._declarationToDefinition.values();
   }
 
   findSymbolDeclaration(startSymbol: ts.Symbol): ts.Declaration | null {
@@ -135,7 +175,9 @@ export class TypeContext {
       );
     }
 
-    const nameDefinition = this._declarationToName.get(declarationResult.value);
+    const nameDefinition = this._declarationToDefinition.get(
+      declarationResult.value,
+    );
     if (nameDefinition == null) {
       return err(gqlErr(unresolved, E.unresolvedTypeReference()));
     }
@@ -156,12 +198,12 @@ export class TypeContext {
     if (referenceNode == null) return false;
     const declaration = this.maybeTsDeclarationForTsName(referenceNode);
     if (declaration == null) return false;
-    return this._declarationToName.has(declaration);
+    return this._declarationToDefinition.has(declaration);
   }
 
   gqlNameDefinitionForGqlName(
     nameNode: NameNode,
-  ): DiagnosticResult<NameDefinition> {
+  ): DiagnosticResult<DeclarationDefinition> {
     const referenceNode = this.getEntityName(nameNode);
     if (referenceNode == null) {
       throw new Error("Expected to find reference node for name node.");
@@ -171,7 +213,7 @@ export class TypeContext {
     if (declaration == null) {
       return err(gqlErr(nameNode, E.unresolvedTypeReference()));
     }
-    const definition = this._declarationToName.get(declaration);
+    const definition = this._declarationToDefinition.get(declaration);
     if (definition == null) {
       return err(gqlErr(nameNode, E.unresolvedTypeReference()));
     }
@@ -192,7 +234,9 @@ export class TypeContext {
       );
     }
 
-    const nameDefinition = this._declarationToName.get(declarationResult.value);
+    const nameDefinition = this._declarationToDefinition.get(
+      declarationResult.value,
+    );
     if (nameDefinition == null) {
       return err(tsErr(node, E.unresolvedTypeReference()));
     }

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

@@ -90,7 +90,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(
@@ -177,6 +181,7 @@ function combineSnapshots(snapshots: ExtractionSnapshot[]): ExtractionSnapshot {
   const result: ExtractionSnapshot = {
     definitions: [],
     nameDefinitions: new Map(),
+    implicitNameDefinitions: new Map(),
     unresolvedNames: new Map(),
     typesWithTypename: new Set(),
     interfaceDeclarations: [],
@@ -195,6 +200,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);
     }