Be stricter about types and usage of getDep/getDepOptions

Author: benjie

grafast/dataplan-pg/src/steps/pgPolymorphic.ts

@@ -82,12 +82,13 @@ export class PgPolymorphicStep<
     }) as any;
   }
 
-  itemPlan(): TItemStep {
-    return this.getDep<any>(this.itemStepId);
+  private itemPlan(): TItemStep {
+    return this.getDepOptions<TItemStep>(this.itemStepId).step;
   }
 
-  typeSpecifierPlan(): TTypeSpecifierStep {
-    return this.getDep<TTypeSpecifierStep>(this.typeSpecifierStepId);
+  private typeSpecifierPlan(): TTypeSpecifierStep {
+    return this.getDepOptions<TTypeSpecifierStep>(this.typeSpecifierStepId)
+      .step;
   }
 
   planForType(type: GraphQLObjectType): ExecutableStep {
@@ -101,7 +102,9 @@ export class PgPolymorphicStep<
         ).join("', '")}'`,
       );
     }
-    return spec.plan(this.typeSpecifierPlan(), this.itemPlan());
+    const $typeSpecifier = this.typeSpecifierPlan();
+    const $item = this.itemPlan();
+    return spec.plan($typeSpecifier, $item);
   }
 
   private getTypeNameFromSpecifier(specifier: TTypeSpecifier) {

grafast/dataplan-pg/src/steps/pgSelect.ts

@@ -1644,7 +1644,7 @@ export class PgSelectRowsStep<
   }
 
   public getClassStep(): PgSelectStep<TResource> {
-    return this.getDep<PgSelectStep<TResource>>(0);
+    return this.getDepOptions<PgSelectStep<TResource>>(0).step;
   }
 
   listItem(itemPlan: Step) {

grafast/dataplan-pg/src/steps/pgUnionAll.ts

@@ -1096,7 +1096,7 @@ export class PgUnionAllRowsStep<
     this.addDependency($pgUnionAll);
   }
   public getClassStep(): PgUnionAllStep<TAttributes, TTypeNames> {
-    return this.getDep<PgUnionAllStep<TAttributes, TTypeNames>>(0);
+    return this.getDepOptions<PgUnionAllStep<TAttributes, TTypeNames>>(0).step;
   }
 
   listItem(itemPlan: Step) {

grafast/grafast/src/engine/OutputPlan.ts

@@ -466,7 +466,7 @@ export class OutputPlan<TType extends OutputPlanType = OutputPlanType> {
         | undefined = ($root.unbatchedExecute! as any)[expressionSymbol];
       if (expressionDetails !== undefined) {
         // @ts-ignore
-        const $parent: Step = $root.getDep(0);
+        const { step: $parent } = $root.getDepOptions(0);
         this.layerPlan.operationPlan.stepTracker.setOutputPlanRootStep(
           this,
           $parent,

grafast/grafast/src/interfaces.ts

@@ -788,7 +788,7 @@ export interface LocationDetails {
 }
 
 export type UnwrapPlanTuple</* const */ TIn extends readonly Step[]> = {
-  [Index in keyof TIn]: TIn[Index] extends Step<infer U> ? U : never;
+  [Index in keyof TIn]: DataFromStep<TIn[Index]>;
 };
 
 export type NotVariableValueNode = Exclude<ValueNode, VariableNode>;
@@ -809,14 +809,21 @@ export type Maybe<T> = T | null | undefined;
 
 export * from "./planJSONInterfaces.js";
 
-export interface AddDependencyOptions {
-  step: Step;
+export interface AddDependencyOptions<TStep extends Step = Step> {
+  step: TStep;
   skipDeduplication?: boolean;
   /** @defaultValue `FLAG_NULL` */
   acceptFlags?: ExecutionEntryFlags;
-  onReject?: null | Error | undefined;
+  onReject?: Maybe<Error>;
   nonUnaryMessage?: ($dependent: Step, $dependency: Step) => string;
 }
+
+export interface DependencyOptions<TStep extends Step = Step> {
+  step: TStep;
+  acceptFlags: ExecutionEntryFlags;
+  onReject: Maybe<Error>;
+}
+
 /**
  * @internal
  */

grafast/grafast/src/step.ts

@@ -20,6 +20,7 @@ import { getDebug } from "./global.js";
 import { inspect } from "./inspect.js";
 import type {
   AddDependencyOptions,
+  DependencyOptions,
   ExecutionDetails,
   ExecutionEntryFlags,
   ExecutionResults,
@@ -36,7 +37,7 @@ import {
   ALL_FLAGS,
   DEFAULT_FORBIDDEN_FLAGS,
 } from "./interfaces.js";
-import type { __ItemStep } from "./steps/index.js";
+import type { __FlagStep, __ItemStep } from "./steps/index.js";
 import { stepADependsOnStepB, stepAMayDependOnStepB } from "./utils.js";
 
 /**
@@ -167,8 +168,11 @@ export /* abstract */ class Step<TData = any> {
   /**
    * What execution entry flags we can't handle for the given indexed dependency
    * (default = this.defaultForbiddenFlags)
+   *
+   * @internal
    */
   protected readonly dependencyForbiddenFlags: ReadonlyArray<ExecutionEntryFlags>;
+  /** @internal */
   protected readonly dependencyOnReject: ReadonlyArray<
     Error | null | undefined
   >;
@@ -347,25 +351,49 @@ export /* abstract */ class Step<TData = any> {
     return this.layerPlan.getStep(id, this);
   }
 
-  protected getDepOptions(depId: number): AddDependencyOptions {
-    const step = this.dependencies[depId];
+  protected getDepOptions<TStep extends Step = Step>(
+    depId: number,
+  ): DependencyOptions<TStep> {
+    const step = this.dependencies[depId] as TStep;
     const forbiddenFlags = this.dependencyForbiddenFlags[depId];
     const onReject = this.dependencyOnReject[depId];
     const acceptFlags = ALL_FLAGS & ~forbiddenFlags;
     return { step, acceptFlags, onReject };
   }
 
-  protected getDep<T extends Step = Step>(_depId: number): T {
+  protected getDep<TStep extends Step = Step>(
+    _depId: number,
+  ): TStep | __FlagStep<TStep>;
+  protected getDep<TStep extends Step = Step>(
+    _depId: number,
+    throwOnFlagged: true,
+  ): TStep;
+  protected getDep<TStep extends Step = Step>(
+    _depId: number,
+    _throwOnFlagged = false,
+  ): TStep | __FlagStep<TStep> {
     // This gets replaced when `__FlagStep` is loaded. Were we on ESM we could
     // just put the code here, but since we're not we have to avoid the
     // circular dependency.
     throw new Error(`Grafast failed to load correctly`);
   }
 
-  protected maybeGetDep<T extends Step = Step>(
+  protected maybeGetDep<TStep extends Step = Step>(
+    depId: number | null | undefined,
+  ): TStep | __FlagStep<TStep> | null;
+  protected maybeGetDep<TStep extends Step = Step>(
+    depId: number | null | undefined,
+    throwOnFlagged: true,
+  ): TStep | null;
+  protected maybeGetDep<TStep extends Step = Step>(
     depId: number | null | undefined,
-  ): T | null {
-    return depId == null ? null : this.getDep<T>(depId);
+    throwOnFlagged = false,
+  ): TStep | __FlagStep<TStep> | null {
+    return depId == null
+      ? null
+      : throwOnFlagged
+      ? this.getDep<TStep>(depId, true)
+      : this.getDep<TStep>(depId);
   }
 
   protected getDepOrConstant<TData = any>(

grafast/grafast/src/steps/__flag.ts

@@ -3,6 +3,7 @@ import { $$inhibit, flagError, SafeError } from "../error.js";
 import { inspect } from "../inspect.js";
 import type {
   AddDependencyOptions,
+  DataFromStep,
   ExecutionDetails,
   ExecutionEntryFlags,
   GrafastResultsList,
@@ -89,7 +90,7 @@ function resolveTrapValue(tv: TrapValue): ResolvedTrapValue {
   }
 }
 
-export class __FlagStep<TData> extends Step<TData> {
+export class __FlagStep<TStep extends Step> extends Step<DataFromStep<TStep>> {
   static $$export = {
     moduleName: "grafast",
     exportName: "__FlagStep",
@@ -102,7 +103,7 @@ export class __FlagStep<TData> extends Step<TData> {
   private valueForInhibited: ResolvedTrapValue;
   private valueForError: ResolvedTrapValue;
   private canBeInlined: boolean;
-  constructor(step: Step, options: FlagStepOptions) {
+  constructor(step: TStep, options: FlagStepOptions) {
     super();
     const {
       acceptFlags = DEFAULT_ACCEPT_FLAGS,
@@ -207,8 +208,8 @@ export class __FlagStep<TData> extends Step<TData> {
   }
 
   public execute(
-    _details: ExecutionDetails<[data: TData, cond?: boolean]>,
-  ): GrafastResultsList<TData> {
+    _details: ExecutionDetails<[data: DataFromStep<TStep>, cond?: boolean]>,
+  ): GrafastResultsList<DataFromStep<TStep>> {
     throw new Error(`${this} not finalized?`);
   }
 
@@ -222,7 +223,7 @@ export class __FlagStep<TData> extends Step<TData> {
   }
 
   private fancyExecute(
-    details: ExecutionDetails<[data: TData, cond?: boolean]>,
+    details: ExecutionDetails<[data: DataFromStep<TStep>, cond?: boolean]>,
   ): any {
     const dataEv = details.values[0]!;
     const condEv =
@@ -268,7 +269,7 @@ export class __FlagStep<TData> extends Step<TData> {
 
   // Checks already performed via addDependency, just pass everything through. Should have been inlined!
   private passThroughExecute(
-    details: ExecutionDetails<[data: TData, cond?: boolean]>,
+    details: ExecutionDetails<[data: DataFromStep<TStep>, cond?: boolean]>,
   ): any {
     const ev = details.values[0];
     if (ev.isBatch) {
@@ -284,11 +285,11 @@ export class __FlagStep<TData> extends Step<TData> {
  * Example use case: get user by id, but id is null: no need to fetch the user
  * since we know they won't exist.
  */
-export function inhibitOnNull<T>(
-  $step: Step<T>,
+export function inhibitOnNull<TStep extends Step>(
+  $step: TStep,
   options?: { if?: FlagStepOptions["if"] },
 ) {
-  return new __FlagStep<T>($step, {
+  return new __FlagStep<TStep>($step, {
     ...options,
     acceptFlags: DEFAULT_ACCEPT_FLAGS & ~FLAG_NULL,
   });
@@ -299,39 +300,52 @@ export function inhibitOnNull<T>(
  * that represents a Post instead: throw error to tell user they've sent invalid
  * data.
  */
-export function assertNotNull<T>(
-  $step: Step<T>,
+export function assertNotNull<TStep extends Step>(
+  $step: TStep,
   message: string,
   options?: { if?: FlagStepOptions["if"] },
 ) {
-  return new __FlagStep<T>($step, {
+  return new __FlagStep<TStep>($step, {
     ...options,
     acceptFlags: DEFAULT_ACCEPT_FLAGS & ~FLAG_NULL,
     onReject: new SafeError(message),
   });
 }
 
-export function trap<T>(
-  $step: Step<T>,
+export function trap<TStep extends Step>(
+  $step: TStep,
   acceptFlags: ExecutionEntryFlags,
   options?: {
     valueForInhibited?: FlagStepOptions["valueForInhibited"];
     valueForError?: FlagStepOptions["valueForError"];
     if?: FlagStepOptions["if"];
   },
 ) {
-  return new __FlagStep<T>($step, {
+  return new __FlagStep<TStep>($step, {
     ...options,
     acceptFlags: (acceptFlags & TRAPPABLE_FLAGS) | FLAG_NULL,
   });
 }
 
 // Have to overwrite the getDep method due to circular dependency
-(Step.prototype as any).getDep = function (this: Step, depId: number) {
+(Step.prototype as any).getDep = function (
+  this: Step,
+  depId: number,
+  throwOnFlagged = false,
+) {
   const { step, acceptFlags, onReject } = this.getDepOptions(depId);
   if (acceptFlags === DEFAULT_ACCEPT_FLAGS && onReject == null) {
     return step;
   } else {
+    if (throwOnFlagged) {
+      throw new Error(
+        `When retrieving dependency ${step} of ${this}, the dependency is flagged as ${digestAcceptFlags(
+          acceptFlags,
+        )}/onReject=${String(
+          onReject,
+        )}. Please use \`this.getDepOptions(depId)\` instead, and handle the flags`,
+      );
+    }
     // Return a __FlagStep around options.step so that all the options are preserved.
     return new __FlagStep(step, { acceptFlags, onReject });
   }

grafast/grafast/src/steps/__inputDynamicScalar.ts

@@ -132,7 +132,7 @@ export class __InputDynamicScalarStep<
   /** @internal */
   eval(): TLeaf {
     const variableValues = this.variableNames.map((variableName, i) =>
-      this.getDep<__TrackedValueStep>(i).eval(),
+      this.getDep<__TrackedValueStep>(i, true).eval(),
     );
     return this.valueFromValues(variableValues);
   }

grafast/grafast/src/steps/__trackedValue.ts

@@ -166,7 +166,7 @@ export class __TrackedValueStep<
   }
 
   private getValuePlan() {
-    return this.getDep<__ValueStep<TData> | AccessStep<TData>>(0);
+    return this.getDep<__ValueStep<TData> | AccessStep<TData>>(0, true);
   }
 
   /**

grafast/grafast/src/steps/access.ts

@@ -160,12 +160,12 @@ export class AccessStep<TData> extends UnbatchedStep<TData> {
   }
 
   toStringMeta(): string {
-    return `${chalk.bold.yellow(String(this.getDep(0).id))}.${this.path
-      .map((p) => String(p))
-      .join(".")}`;
+    return `${chalk.bold.yellow(
+      String(this.getDepOptions(0).step.id),
+    )}.${this.path.map((p) => String(p)).join(".")}`;
   }
 
-  getParentStep(): Step {
+  getParentStep() {
     return this.getDep(0);
   }