[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: 602bcb6
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,41 @@ 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 name = this.gql.name(tag, "CONTEXT_DUMMY_NAME");
+    this.implicitNameDefinitions.set(
+      {
+        kind: "DERIVED_CONTEXT",
+        name,
+        path: tsModulePath,
+        exportName: funcName?.text ?? null,
+      },
+      returnType,
+    );
+  }
+
   extractType(node: ts.Node, tag: ts.JSDocTag) {
     if (ts.isClassDeclaration(node)) {
       this.typeClassDeclaration(node, tag);

src/TypeContext.ts

@@ -18,18 +18,25 @@ import { ExtractionSnapshot } from "./Extractor";
 
 export const UNRESOLVED_REFERENCE_NAME = `__UNRESOLVED_REFERENCE__`;
 
-export type NameDefinition = {
-  name: NameNode;
-  kind:
-    | "TYPE"
-    | "INTERFACE"
-    | "UNION"
-    | "SCALAR"
-    | "INPUT_OBJECT"
-    | "ENUM"
-    | "CONTEXT"
-    | "INFO";
-};
+export type NameDefinition =
+  | {
+      name: NameNode;
+      kind:
+        | "TYPE"
+        | "INTERFACE"
+        | "UNION"
+        | "SCALAR"
+        | "INPUT_OBJECT"
+        | "ENUM"
+        | "CONTEXT"
+        | "INFO";
+    }
+  | {
+      name: NameNode;
+      path: string;
+      exportName: string | null;
+      kind: "DERIVED_CONTEXT";
+    };
 
 type TsIdentifier = number;
 
@@ -61,7 +68,16 @@ export class TypeContext {
       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) {
+        throw new Error(
+          "Expected to find declaration for implicit name definition.",
+        );
+      }
+      self._recordTypeName(declaration, definition);
     }
     return self;
   }
@@ -72,13 +88,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/resolverCodegen.ts

@@ -178,6 +178,16 @@ export default class ResolverCodegen {
           F.createIdentifier("args"),
           F.createIdentifier(arg.name),
         );
+      case "derivedContext": {
+        const localName = "derivedContext";
+        this.ts.importUserConstruct(arg.path, arg.exportName, localName);
+        return F.createCallExpression(
+          F.createIdentifier(localName),
+          undefined,
+          [this.resolverParam(arg.input)],
+        );
+      }
+
       default:
         // @ts-expect-error
         throw new Error(`Unexpected resolver kind ${arg.kind}`);

src/lib.ts

@@ -173,6 +173,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 +192,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

@@ -87,6 +87,7 @@ export type ResolverArgument =
   | SourceArgument
   | ArgumentsObjectArgument
   | ContextArgument
+  | DerivedContextArgument
   | InformationArgument
   | NamedArgument;
 
@@ -105,6 +106,15 @@ 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
+  input: ContextArgument | DerivedContextArgument;
+  // TODO: Add a parent which could be ContextArgument or another DerivedContextArgument
+};
+
 /** 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;
+  // TODO: Support custom inputs
+  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;

src/tests/fixtures/derived_context/simpleDerivedContext.ts

@@ -0,0 +1,21 @@
+/** @gqlContext */
+type RootContext = {
+  userName: string;
+};
+
+type DerivedContext = {
+  greeting: string;
+};
+
+/** @gqlContext */
+export function createDerivedContext(ctx: RootContext): DerivedContext {
+  return { greeting: `Hello, ${ctx.userName}!` };
+}
+
+/** @gqlType */
+type Query = unknown;
+
+/** @gqlField */
+export function greeting(_: Query, ctx: DerivedContext): string {
+  return ctx.greeting;
+}

src/tests/fixtures/derived_context/simpleDerivedContext.ts.expected

@@ -78,23 +78,33 @@ function transformArgs(
   if (args == null) {
     return null;
   }
-  return args.map((arg): ResolverArgument => {
-    switch (arg.kind) {
-      case "argumentsObject":
-        return { kind: "argumentsObject" };
-      case "named":
-        return { kind: "named", name: arg.name };
-      case "source":
-        return { kind: "source" };
-      case "information":
-        return { kind: "information" };
-      case "context":
-        return { kind: "context" };
-      case "unresolved":
-        throw new Error("Unresolved argument in resolver");
-      default:
-        // @ts-expect-error
-        throw new Error(`Unknown argument kind: ${arg.kind}`);
-    }
-  });
+  return args.map(transformArg);
+}
+
+function transformArg(arg: DirectiveResolverArgument): ResolverArgument {
+  switch (arg.kind) {
+    case "argumentsObject":
+      return { kind: "argumentsObject" };
+    case "named":
+      return { kind: "named", name: arg.name };
+    case "source":
+      return { kind: "source" };
+    case "information":
+      return { kind: "information" };
+    case "context":
+      return { kind: "context" };
+    case "derivedContext":
+      return {
+        kind: "derivedContext",
+        path: arg.path,
+        exportName: arg.exportName,
+        // TODO: Support custom inputs
+        input: { kind: "context" },
+      };
+    case "unresolved":
+      throw new Error("Unresolved argument in resolver");
+    default:
+      // @ts-expect-error
+      throw new Error(`Unknown argument kind: ${arg.kind}`);
+  }
 }

src/transforms/makeResolverSignature.ts

@@ -108,6 +108,7 @@ class ResolverParamsResolver {
       case "argumentsObject":
       case "information":
       case "context":
+      case "derivedContext":
       case "source":
         return param;
       case "unresolved": {
@@ -124,6 +125,13 @@ class ResolverParamsResolver {
             return param;
           }
           switch (resolved.value.kind) {
+            case "DERIVED_CONTEXT":
+              return {
+                kind: "derivedContext",
+                node: param.node,
+                path: resolved.value.path,
+                exportName: resolved.value.exportName,
+              };
             case "CONTEXT":
               return { kind: "context", node: param.node };
             case "INFO":