@apollo/federation-internals 2.1.0-alpha.1 → 2.1.0-alpha.4

package/src/operations.ts

@@ -39,14 +39,16 @@ import {
   variableDefinitionsFromAST,
   CompositeType,
   typenameFieldName,
-  NamedType,
   sameDirectiveApplications,
   isConditionalDirective,
   isDirectiveApplicationsSubset,
+  isAbstractType,
+  DeferDirectiveArgs,
+  Variable,
 } from "./definitions";
 import { ERRORS } from "./error";
-import { sameType } from "./types";
-import { assert, mapEntries, MapWithCachedArrays, MultiMap } from "./utils";
+import { isDirectSubtype, sameType } from "./types";
+import { assert, mapEntries, MapWithCachedArrays, MultiMap, SetMultiMap } from "./utils";
 import { argumentsEquals, argumentsFromAST, isValidValue, valueToAST, valueToString } from "./values";
 
 function validate(condition: any, message: () => string, sourceAST?: ASTNode): asserts condition {
@@ -178,27 +180,46 @@ export class Field<TArgs extends {[key: string]: any} = {[key: string]: any}> ex
   updateForAddingTo(selectionSet: SelectionSet): Field<TArgs> {
     const selectionParent = selectionSet.parentType;
     const fieldParent = this.definition.parent;
-    if (selectionParent.name !== fieldParent.name) {
-      if (this.name === typenameFieldName) {
-        return this.withUpdatedDefinition(selectionParent.typenameField()!);
-      }
+    if (selectionParent === fieldParent) {
+      return this;
+    }
 
-      // We accept adding a selection of an interface field to a selection of one of its subtype. But otherwise, it's invalid.
-      // Do note that the field might come from a supergraph while the selection is on a subgraph, so we avoid relying on isDirectSubtype (because
-      // isDirectSubtype relies on the subtype knowing which interface it implements, but the one of the subgraph might not declare implementing
-      // the supergraph interface, even if it does in the subgraph).
-      validate(
+    if (this.name === typenameFieldName) {
+      return this.withUpdatedDefinition(selectionParent.typenameField()!);
+    }
+
+    // We accept adding a selection of an interface field to a selection of one of its subtype. But otherwise, it's invalid.
+    // Do note that the field might come from a supergraph while the selection is on a subgraph, so we avoid relying on isDirectSubtype (because
+    // isDirectSubtype relies on the subtype knowing which interface it implements, but the one of the subgraph might not declare implementing
+    // the supergraph interface, even if it does in the subgraph).
+    validate(
+      selectionParent.name == fieldParent.name
+      || (
         !isUnionType(selectionParent)
         && (
           (isInterfaceType(fieldParent) && fieldParent.allImplementations().some(i => i.name == selectionParent.name))
           || (isObjectType(fieldParent) && fieldParent.name == selectionParent.name)
-        ),
-        () => `Cannot add selection of field "${this.definition.coordinate}" to selection set of parent type "${selectionSet.parentType}"`
-      );
-      const fieldDef = selectionParent.field(this.name);
-      validate(fieldDef, () => `Cannot add selection of field "${this.definition.coordinate}" to selection set of parent type "${selectionParent} (that does not declare that type)"`);
-      return this.withUpdatedDefinition(fieldDef);
-    }
+        )
+      ),
+      () => `Cannot add selection of field "${this.definition.coordinate}" to selection set of parent type "${selectionSet.parentType}"`
+    );
+    const fieldDef = selectionParent.field(this.name);
+    validate(fieldDef, () => `Cannot add selection of field "${this.definition.coordinate}" to selection set of parent type "${selectionParent}" (that does not declare that field)`);
+    return this.withUpdatedDefinition(fieldDef);
+  }
+
+  hasDefer(): boolean {
+    // @defer cannot be on field at the moment
+    return false;
+  }
+
+  deferDirectiveArgs(): undefined {
+    // @defer cannot be on field at the moment (but exists so we can call this method on any `OperationElement` conveniently)
+    return undefined;
+  }
+
+  withoutDefer(): Field<TArgs> {
+    // @defer cannot be on field at the moment
     return this;
   }
 
@@ -243,8 +264,15 @@ export class FragmentElement extends AbstractOperationElement<FragmentElement> {
     return this.sourceType;
   }
 
+  castedType(): CompositeType {
+    return this.typeCondition ? this.typeCondition : this.sourceType;
+  }
+
   withUpdatedSourceType(newSourceType: CompositeType): FragmentElement {
-    const newFragment = new FragmentElement(newSourceType, this.typeCondition);
+    // Note that we pass the type-condition name instead of the type itself, to ensure that if `newSourceType` was from a different
+    // schema (typically, the supergraph) than `this.sourceType` (typically, a subgraph), then the new condition uses the
+    // definition of the proper schema (the supergraph in such cases, instead of the subgraph).
+    const newFragment = new FragmentElement(newSourceType, this.typeCondition?.name);
     for (const directive of this.appliedDirectives) {
       newFragment.applyDirective(directive.definition!, directive.arguments());
     }
@@ -266,6 +294,106 @@ export class FragmentElement extends AbstractOperationElement<FragmentElement> {
     return this;
   }
 
+  hasDefer(): boolean {
+    return this.hasAppliedDirective('defer');
+  }
+
+  hasStream(): boolean {
+    return this.hasAppliedDirective('stream');
+  }
+
+  deferDirectiveArgs(): DeferDirectiveArgs | undefined {
+    // Note: @defer is not repeatable, so the return array below is either empty, or has a single value.
+    return this.appliedDirectivesOf(this.schema().deferDirective())[0]?.arguments();
+  }
+
+  /**
+   * Returns this fragment element but with any @defer directive on it removed.
+   *
+   * This method will return `undefined` if, upon removing @defer, the fragment has no conditions nor
+   * any remaining applied directives (meaning that it carries no information whatsoever and can be
+   * ignored).
+   */
+  withoutDefer(): FragmentElement | undefined {
+    const deferName = this.schema().deferDirective().name;
+    const updatedDirectives = this.appliedDirectives.filter((d) => d.name !== deferName);
+    if (!this.typeCondition && updatedDirectives.length === 0) {
+      return undefined;
+    }
+
+    if (updatedDirectives.length === this.appliedDirectives.length) {
+      return this;
+    }
+
+    const updated = new FragmentElement(this.sourceType, this.typeCondition);
+    updatedDirectives.forEach((d) => updated.applyDirective(d.definition!, d.arguments()));
+    return updated;
+  }
+
+  /**
+   * Returns this fragment element, but it is has a @defer directive, the element is returned with
+   * the @defer "normalized".
+   *
+   * See `Operation.withNormalizedDefer` for details on our so-called @defer normalization.
+   */
+  withNormalizedDefer(normalizer: DeferNormalizer): FragmentElement | undefined {
+    const deferArgs = this.deferDirectiveArgs();
+    if (!deferArgs) {
+      return this;
+    }
+
+    let newDeferArgs: DeferDirectiveArgs | undefined = undefined;
+    let conditionVariable: Variable | undefined = undefined;
+    if (deferArgs.if !== undefined) {
+      if (typeof deferArgs.if === 'boolean') {
+        if (deferArgs.if) {
+          // Harcoded `if: true`, remove the `if`
+          newDeferArgs = {
+            ...deferArgs,
+            if: undefined,
+          }
+        } else {
+          // Harcoded `if: false`, remove the @defer altogether
+          return this.withoutDefer();
+        }
+      } else {
+        // `if` on a variable
+        conditionVariable = deferArgs.if;
+      }
+    }
+
+    let label = deferArgs.label;
+    if (!label) {
+      label = normalizer.newLabel();
+      if (newDeferArgs) {
+        newDeferArgs.label = label;
+      } else {
+        newDeferArgs = {
+          ...deferArgs,
+          label,
+        }
+      }
+    }
+
+    // Now that we are sure to have a label, if we had a (non-trivial) condition,
+    // associate it to that label.
+    if (conditionVariable) {
+      normalizer.registerCondition(label, conditionVariable);
+    }
+
+    if (!newDeferArgs) {
+      return this;
+    }
+
+    const updated = new FragmentElement(this.sourceType, this.typeCondition);
+    const deferDirective = this.schema().deferDirective();
+    // Re-apply all the non-defer directives
+    this.appliedDirectives.filter((d) => d.name !== deferDirective.name).forEach((d) => updated.applyDirective(d.definition!, d.arguments()));
+    // And then re-apply the @defer with the new label.
+    updated.applyDirective(this.schema().deferDirective(), newDeferArgs);
+    return updated;
+  }
+
   equals(that: OperationElement): boolean {
     if (this === that) {
       return true;
@@ -360,30 +488,37 @@ export class Operation {
   }
 
   optimize(fragments?: NamedFragments, minUsagesToOptimize: number = 2): Operation {
-    if (!fragments) {
+    assert(minUsagesToOptimize >= 1, `Expected 'minUsagesToOptimize' to be at least 1, but got ${minUsagesToOptimize}`)
+    if (!fragments || fragments.isEmpty()) {
       return this;
     }
+
     let optimizedSelection = this.selectionSet.optimize(fragments);
     if (optimizedSelection === this.selectionSet) {
       return this;
     }
 
-    // Optimizing fragments away, and then de-optimizing if it's used less than we want, feels a bit wasteful,
-    // but it's simple and probably don't matter too much in practice (we only call this optimization on the
-    // final compted query plan, so not a very hot path; plus in most case we won't even reach that point
-    // either because there is no fragment, or none will have been optimized away and we'll exit above). We
-    // can optimize later if this show up in profiling though.
-    if (minUsagesToOptimize > 1) {
-      const usages = new Map<string, number>();
-      optimizedSelection.collectUsedFragmentNames(usages);
-      for (const fragment of fragments.names()) {
-        if (!usages.has(fragment)) {
-          usages.set(fragment, 0);
-        }
+    const usages = new Map<string, number>();
+    optimizedSelection.collectUsedFragmentNames(usages);
+    for (const fragment of fragments.names()) {
+      if (!usages.has(fragment)) {
+        usages.set(fragment, 0);
       }
-      const toDeoptimize = mapEntries(usages).filter(([_, count]) => count < minUsagesToOptimize).map(([name]) => name);
-      optimizedSelection = optimizedSelection.expandFragments(toDeoptimize);
     }
+
+    // We re-expand any fragments that is used less than our minimum. Optimizing all fragments to potentially
+    // re-expand some is not entirely optimal, but it's simple and probably don't matter too much in practice
+    // (we only call this optimization on the final computed query plan, so not a very hot path; plus in most
+    // cases we won't even reach that point either because there is no fragment, or none will have been
+    // optimized away so we'll exit above). We can optimize later if this show up in profiling though.
+    //
+    // Also note `toDeoptimize` will always contains the unused fragments, which will allow `expandFragments`
+    // to remove them from the listed fragments in `optimizedSelection` (here again, this could make use call
+    // `expandFragments` on _only_ unused fragments and that case could be dealt with more efficiently, but
+    // probably not noticeable in practice so ...).
+    const toDeoptimize = mapEntries(usages).filter(([_, count]) => count < minUsagesToOptimize).map(([name]) => name);
+    optimizedSelection = optimizedSelection.expandFragments(toDeoptimize);
+
     return new Operation(this.rootKind, optimizedSelection, this.variableDefinitions, this.name);
   }
 
@@ -401,6 +536,61 @@ export class Operation {
     );
   }
 
+  /**
+   * Returns this operation but potentially modified so all/some of the @defer applications have been removed.
+   *
+   * @param labelsToRemove - If provided, then only the `@defer` applications with labels in the provided
+   * set will be remove. Other `@defer` applications will be untouched. If `undefined`, then all `@defer`
+   * applications are removed.
+   */
+  withoutDefer(labelsToRemove?: Set<string>): Operation {
+    // If we have named fragments, we should be looking inside those and either expand those having @defer or,
+    // probably better, replace them with a verison without @defer. But as we currently only call this method
+    // after `expandAllFragments`, we'll implement this when/if we need it.
+    assert(!this.selectionSet.fragments || this.selectionSet.fragments.isEmpty(), 'Removing @defer currently only work on "expanded" selections (no named fragments)');
+    const updated = this.selectionSet.withoutDefer(labelsToRemove);
+    return updated == this.selectionSet
+      ? this
+      : new Operation(this.rootKind, updated, this.variableDefinitions, this.name);
+  }
+
+  /**
+   * Returns this operation but modified to "normalize" all the @defer applications.
+   *
+   * "Normalized" in this context means that all the `@defer` application in the
+   * resulting operation will:
+   *  - have a (unique) label. Which imply that this method generates label for
+   *    any `@defer` not having a label.
+   *  - have a non-trivial `if` condition, if any. By non-trivial, we mean that
+   *    the condition will be a variable and not an hard-coded `true` or `false`.
+   *    To do this, this method will remove the condition of any `@defer` that
+   *    has `if: true`, and will completely remove any `@defer` application that
+   *    has `if: false`.
+   */
+  withNormalizedDefer(): {
+    operation: Operation,
+    hasDefers: boolean,
+    assignedDeferLabels: Set<string>,
+    deferConditions: SetMultiMap<string, string>,
+  } {
+    // Similar comment than in `withoutDefer`
+    assert(!this.selectionSet.fragments || this.selectionSet.fragments.isEmpty(), 'Assigning @defer lables currently only work on "expanded" selections (no named fragments)');
+
+    const normalizer = new DeferNormalizer();
+    const { hasDefers, hasNonLabelledOrConditionalDefers } = normalizer.init(this.selectionSet);
+    let updatedOperation: Operation = this;
+    if (hasNonLabelledOrConditionalDefers) {
+      const updated = this.selectionSet.withNormalizedDefer(normalizer);
+      updatedOperation = new Operation(this.rootKind, updated, this.variableDefinitions, this.name);
+    }
+    return {
+      operation: updatedOperation,
+      hasDefers,
+      assignedDeferLabels: normalizer.assignedLabels,
+      deferConditions: normalizer.deferConditions,
+    };
+  }
+
   toString(expandFragments: boolean = false, prettyPrint: boolean = true): string {
     return this.selectionSet.toOperationString(this.rootKind, this.variableDefinitions, this.name, expandFragments, prettyPrint);
   }
@@ -434,6 +624,10 @@ export class NamedFragmentDefinition extends DirectiveTargetElement<NamedFragmen
     super(schema);
   }
 
+  withUpdatedSelectionSet(newSelectionSet: SelectionSet): NamedFragmentDefinition {
+    return new NamedFragmentDefinition(this.schema(), this.name, this.typeCondition, newSelectionSet);
+  }
+
   variables(): Variables {
     return mergeVariables(this.variablesInAppliedDirectives(), this.selectionSet.usedVariables());
   }
@@ -460,6 +654,19 @@ export class NamedFragmentDefinition extends DirectiveTargetElement<NamedFragmen
     };
   }
 
+  /**
+   * Whether this fragment may apply at the provided type, that is if its type condition matches the type
+   * or is a supertype of it.
+   *
+   * @param type - the type at which we're looking at applying the fragment
+   */
+  canApplyAtType(type: CompositeType): boolean {
+    return (
+      sameType(this.typeCondition, type)
+      || (isAbstractType(this.typeCondition) && !isUnionType(type) && isDirectSubtype(this.typeCondition, type))
+    );
+  }
+
   toString(indent?: string): string {
     return (indent ?? '') + `fragment ${this.name} on ${this.typeCondition}${this.appliedDirectivesToString()} ${this.selectionSet.toString(false, true, indent)}`;
   }
@@ -497,8 +704,8 @@ export class NamedFragments {
     }
   }
 
-  onType(type: NamedType): NamedFragmentDefinition[] {
-    return this.fragments.values().filter(f => f.typeCondition.name === type.name);
+  maybeApplyingAtType(type: CompositeType): NamedFragmentDefinition[] {
+    return this.fragments.values().filter(f => f.canApplyAtType(type));
   }
 
   without(names: string[]): NamedFragments {
@@ -525,6 +732,10 @@ export class NamedFragments {
     return this.fragments.get(name);
   }
 
+  has(name: string): boolean {
+    return this.fragments.has(name);
+  }
+
   definitions(): readonly NamedFragmentDefinition[] {
     return this.fragments.values();
   }
@@ -593,6 +804,67 @@ abstract class Freezable<T> {
   abstract clone(): T;
 }
 
+/**
+ * Utility class used to handle "normalizing" the @defer in an operation.
+ *
+ * See `Operation.withNormalizedDefer` for details on what we mean by normalizing in
+ * this context.
+ */
+class DeferNormalizer {
+  private index = 0;
+  readonly assignedLabels = new Set<string>();
+  readonly deferConditions = new SetMultiMap<string, string>();
+  private readonly usedLabels = new Set<string>();
+
+  /**
+   * Initializes the "labeller" with all the labels used in the provided selections set.
+   *
+   * @return - whether `selectionSet` has any non-labeled @defer.
+   */
+  init(selectionSet: SelectionSet): { hasDefers: boolean, hasNonLabelledOrConditionalDefers: boolean }  {
+    let hasNonLabelledOrConditionalDefers = false;
+    let hasDefers = false;
+    const stack: Selection[] = selectionSet.selections().concat();
+    while (stack.length > 0) {
+      const selection = stack.pop()!;
+      if (selection.kind === 'FragmentSelection') {
+        const deferArgs = selection.element().deferDirectiveArgs();
+        if (deferArgs) {
+          hasDefers = true;
+          if (deferArgs.label) {
+            this.usedLabels.add(deferArgs.label);
+          } else {
+            hasNonLabelledOrConditionalDefers = true;
+          }
+        }
+      }
+      if (selection.selectionSet) {
+        selection.selectionSet.selections().forEach((s) => stack.push(s));
+      }
+    }
+    return { hasDefers, hasNonLabelledOrConditionalDefers };
+  }
+
+  private nextLabel(): string {
+    return `qp__${this.index++}`;
+  }
+
+  newLabel(): string {
+    let candidate = this.nextLabel();
+    // It's unlikely that auto-generated label would conflict an existing one, but
+    // not taking any chances.
+    while (this.usedLabels.has(candidate)) {
+      candidate = this.nextLabel();
+    }
+    this.assignedLabels.add(candidate);
+    return candidate;
+  }
+
+  registerCondition(label: string, condition: Variable): void {
+    this.deferConditions.add(condition.name, label);
+  }
+}
+
 export class SelectionSet extends Freezable<SelectionSet> {
   // The argument is either the responseName (for fields), or the type name (for fragments), with the empty string being used as a special
   // case for a fragment with no type condition.
@@ -648,10 +920,6 @@ export class SelectionSet extends Freezable<SelectionSet> {
   }
 
   collectUsedFragmentNames(collector: Map<string, number>) {
-    if (!this.fragments) {
-      return;
-    }
-
     for (const byResponseName of this._selections.values()) {
       for (const selection of byResponseName) {
         selection.collectUsedFragmentNames(collector);
@@ -680,10 +948,6 @@ export class SelectionSet extends Freezable<SelectionSet> {
   }
 
   expandFragments(names?: string[], updateSelectionSetFragments: boolean = true): SelectionSet {
-    if (!this.fragments) {
-      return this;
-    }
-
     if (names && names.length === 0) {
       return this;
     }
@@ -693,11 +957,60 @@ export class SelectionSet extends Freezable<SelectionSet> {
       : this.fragments;
     const withExpanded = new SelectionSet(this.parentType, newFragments);
     for (const selection of this.selections()) {
-      withExpanded.add(selection.expandFragments(names, updateSelectionSetFragments));
+      const expanded = selection.expandFragments(names, updateSelectionSetFragments);
+      if (Array.isArray(expanded)) {
+        withExpanded.addAll(expanded);
+      } else {
+        withExpanded.add(expanded as Selection);
+      }
     }
     return withExpanded;
   }
 
+  /**
+   * Returns the result of mapping the provided `mapper` to all the selection of this selection set.
+   *
+   * This method assumes that the `mapper` may often return it's argument directly, meaning that only
+   * a small subset of selection actually need any modifications, and will avoid re-creating new
+   * objects when that is the case. This does mean that the resulting selection set may be `this`
+   * directly, or may alias some of the sub-selection in `this`.
+   */
+  private lazyMap(mapper: (selection: Selection) => Selection | SelectionSet | undefined): SelectionSet {
+    let updatedSelections: Selection[] | undefined = undefined;
+    const selections = this.selections();
+    for (let i = 0; i < selections.length; i++) {
+      const selection = selections[i];
+      const updated = mapper(selection);
+      if (updated !== selection && !updatedSelections) {
+        updatedSelections = [];
+        for (let j = 0; j < i; j++) {
+          updatedSelections.push(selections[j]);
+        }
+      }
+      if (!!updated && updatedSelections) {
+        if (updated instanceof SelectionSet) {
+          updated.selections().forEach((s) => updatedSelections!.push(s));
+        } else {
+          updatedSelections.push(updated);
+        }
+      }
+    }
+    if (!updatedSelections) {
+      return this;
+    }
+    return new SelectionSet(this.parentType, this.fragments).addAll(updatedSelections)
+  }
+
+  withoutDefer(labelsToRemove?: Set<string>): SelectionSet {
+    assert(!this.fragments, 'Not yet supported');
+    return this.lazyMap((selection) => selection.withoutDefer(labelsToRemove));
+  }
+
+  withNormalizedDefer(normalizer: DeferNormalizer): SelectionSet {
+    assert(!this.fragments, 'Not yet supported');
+    return this.lazyMap((selection) => selection.withNormalizedDefer(normalizer));
+  }
+
   /**
    * Returns the selection select from filtering out any selection that does not match the provided predicate.
    *
@@ -705,14 +1018,12 @@ export class SelectionSet extends Freezable<SelectionSet> {
    * call `optimize` on the result if you want to re-apply some fragments.
    */
   filter(predicate: (selection: Selection) => boolean): SelectionSet {
-    const filtered = new SelectionSet(this.parentType, this.fragments);
-    for (const selection of this.selections()) {
-      const filteredSelection = selection.filter(predicate);
-      if (filteredSelection) {
-        filtered.add(filteredSelection);
-      }
-    }
-    return filtered;
+    return this.lazyMap((selection) => selection.filter(predicate));
+  }
+
+  withoutEmptyBranches(): SelectionSet | undefined {
+    const updated = this.filter((selection) => selection.selectionSet?.isEmpty() !== true);
+    return updated.isEmpty() ? undefined : updated;
   }
 
   protected freezeInternals(): void {
@@ -1080,8 +1391,44 @@ export class FieldSelection extends Freezable<FieldSelection> {
     }
   }
 
-  optimize(fragments: NamedFragments): FieldSelection {
+  optimize(fragments: NamedFragments): Selection {
     const optimizedSelection = this.selectionSet ? this.selectionSet.optimize(fragments) : undefined;
+    const fieldBaseType = baseType(this.field.definition.type!);
+    if (isCompositeType(fieldBaseType) && optimizedSelection) {
+      for (const candidate of fragments.maybeApplyingAtType(fieldBaseType)) {
+        // TODO: Checking `equals` here is very simple, but somewhat restrictive in theory. That is, if a query
+        // is:
+        //   {
+        //     t {
+        //       a
+        //       b
+        //       c
+        //     }
+        //   }
+        // and we have:
+        //   fragment X on T {
+        //     t {
+        //       a
+        //       b
+        //     }
+        //   }
+        // then the current code will not use the fragment because `c` is not in the fragment, but in relatity,
+        // we could use it and make the result be:
+        //   {
+        //     ...X
+        //     t {
+        //       c
+        //     }
+        //   }
+        // To do that, we can change that `equals` to `contains`, but then we should also "extract" the remainder
+        // of `optimizedSelection` that isn't covered by the fragment, and that is the part slighly more involved.
+        if (optimizedSelection.equals(candidate.selectionSet)) {
+          const fragmentSelection = new FragmentSpreadSelection(fieldBaseType, fragments, candidate.name);
+          return new FieldSelection(this.field, selectionSetOf(fieldBaseType, fragmentSelection));
+        }
+      }
+    }
+
     return this.selectionSet === optimizedSelection
       ? this
       : new FieldSelection(this.field, optimizedSelection);
@@ -1094,7 +1441,11 @@ export class FieldSelection extends Freezable<FieldSelection> {
     if (!this.selectionSet) {
       return this;
     }
-    return new FieldSelection(this.field, this.selectionSet.filter(predicate));
+
+    const updatedSelectionSet = this.selectionSet.filter(predicate);
+    return this.selectionSet === updatedSelectionSet
+      ? this
+      : new FieldSelection(this.field, updatedSelectionSet);
   }
 
   protected freezeInternals(): void {
@@ -1203,6 +1554,20 @@ export class FieldSelection extends Freezable<FieldSelection> {
     return this.selectionSet?.fragments;
   }
 
+  withoutDefer(labelsToRemove?: Set<string>): FieldSelection {
+    const updatedSubSelections = this.selectionSet?.withoutDefer(labelsToRemove);
+    return updatedSubSelections === this.selectionSet
+      ? this
+      : new FieldSelection(this.field, updatedSubSelections);
+  }
+
+  withNormalizedDefer(normalizer: DeferNormalizer): FieldSelection {
+    const updatedSubSelections = this.selectionSet?.withNormalizedDefer(normalizer);
+    return updatedSubSelections === this.selectionSet
+      ? this
+      : new FieldSelection(this.field, updatedSubSelections);
+  }
+
   clone(): FieldSelection {
     if (!this.selectionSet) {
       return this;
@@ -1230,33 +1595,69 @@ export abstract class FragmentSelection extends Freezable<FragmentSelection> {
 
   abstract optimize(fragments: NamedFragments): FragmentSelection;
 
-  abstract expandFragments(names?: string[]): FragmentSelection;
+  abstract expandFragments(names?: string[]): Selection | readonly Selection[];
 
   abstract toSelectionNode(): SelectionNode;
 
   abstract validate(): void;
 
+  abstract withoutDefer(labelsToRemove?: Set<string>): FragmentSelection | SelectionSet;
+
+  abstract withNormalizedDefer(normalizer: DeferNormalizer): FragmentSelection | SelectionSet;
+
   protected us(): FragmentSelection {
     return this;
   }
 
+  protected validateDeferAndStream() {
+    if (this.element().hasDefer() || this.element().hasStream()) {
+      const schemaDef = this.element().schema().schemaDefinition;
+      const parentType = this.element().parentType;
+      validate(
+        schemaDef.rootType('mutation') !== parentType && schemaDef.rootType('subscription') !== parentType,
+        () => `The @defer and @stream directives cannot be used on ${schemaDef.roots().filter((t) => t.type === parentType).pop()?.rootKind} root type "${parentType}"`,
+      );
+    }
+  }
+
   usedVariables(): Variables {
     return mergeVariables(this.element().variables(), this.selectionSet.usedVariables());
   }
 
   updateForAddingTo(selectionSet: SelectionSet): FragmentSelection {
     const updatedFragment = this.element().updateForAddingTo(selectionSet);
-    return this.element() === updatedFragment
-      ? this.cloneIfFrozen()
-      : new InlineFragmentSelection(updatedFragment, this.selectionSet.cloneIfFrozen());
+    if (this.element() === updatedFragment) {
+      return this.cloneIfFrozen();
+    }
+
+    // Like for fields, we create a new selection that not only uses the updated fragment, but also ensures
+    // the underlying selection set uses the updated type as parent type.
+    const updatedCastedType = updatedFragment.castedType();
+    let updatedSelectionSet : SelectionSet | undefined;
+    if (this.selectionSet.parentType !== updatedCastedType) {
+      updatedSelectionSet = new SelectionSet(updatedCastedType);
+      // Note that re-adding every selection ensures that anything frozen will be cloned as needed, on top of handling any knock-down
+      // effect of the type change.
+      for (const selection of this.selectionSet.selections()) {
+        updatedSelectionSet.add(selection);
+      }
+    } else {
+      updatedSelectionSet = this.selectionSet?.cloneIfFrozen();
+    }
+
+    return new InlineFragmentSelection(updatedFragment, updatedSelectionSet);
   }
 
-  filter(predicate: (selection: Selection) => boolean): InlineFragmentSelection | undefined {
+  filter(predicate: (selection: Selection) => boolean): FragmentSelection | undefined {
     if (!predicate(this)) {
       return undefined;
     }
     // Note that we essentially expand all fragments as part of this.
-    return new InlineFragmentSelection(this.element(), this.selectionSet.filter(predicate));
+    const selectionSet = this.selectionSet;
+    const updatedSelectionSet = selectionSet.filter(predicate);
+    return updatedSelectionSet === selectionSet
+      ? this
+      : new InlineFragmentSelection(this.element(), updatedSelectionSet);
   }
 
   protected freezeInternals() {
@@ -1302,6 +1703,7 @@ class InlineFragmentSelection extends FragmentSelection {
   }
 
   validate() {
+    this.validateDeferAndStream();
     // Note that validation is kind of redundant since `this.selectionSet.validate()` will check that it isn't empty. But doing it
     // allow to provide much better error messages.
     validate(
@@ -1311,7 +1713,6 @@ class InlineFragmentSelection extends FragmentSelection {
     this.selectionSet.validate();
   }
 
-
   get selectionSet(): SelectionSet {
     return this._selectionSet;
   }
@@ -1343,13 +1744,21 @@ class InlineFragmentSelection extends FragmentSelection {
   }
 
   optimize(fragments: NamedFragments): FragmentSelection {
-    const optimizedSelection = this.selectionSet.optimize(fragments);
+    let optimizedSelection = this.selectionSet.optimize(fragments);
     const typeCondition = this.element().typeCondition;
     if (typeCondition) {
-      for (const candidate of fragments.onType(typeCondition)) {
-        if (candidate.selectionSet.equals(optimizedSelection)) {
-          fragments.addIfNotExist(candidate);
-          return new FragmentSpreadSelection(this.element().parentType, fragments, candidate.name);
+      for (const candidate of fragments.maybeApplyingAtType(typeCondition)) {
+        // See comment in `FieldSelection.optimize` about the `equals`: this fully apply here too.
+        if (optimizedSelection.equals(candidate.selectionSet)) {
+          const spread = new FragmentSpreadSelection(this.element().parentType, fragments, candidate.name);
+          // We use the fragment when the fragments condition is either the same, or a supertype of our current condition.
+          // If it's the same type, then we don't really want to preserve the current condition, it is included in the
+          // spread and we can return it directive. But if the fragment condition is a superset, then we should preserve
+          // our current condition since it restricts the selection more than the fragment actual does.
+          if (sameType(typeCondition, candidate.typeCondition)) {
+            return spread;
+          }
+          optimizedSelection = selectionSetOf(spread.element().parentType, spread);
         }
       }
     }
@@ -1369,6 +1778,32 @@ class InlineFragmentSelection extends FragmentSelection {
     this.selectionSet.collectUsedFragmentNames(collector);
   }
 
+  withoutDefer(labelsToRemove?: Set<string>): FragmentSelection | SelectionSet {
+    const updatedSubSelections = this.selectionSet.withoutDefer(labelsToRemove);
+    const deferArgs = this.fragmentElement.deferDirectiveArgs();
+    const hasDeferToRemove = deferArgs && (!labelsToRemove || (deferArgs.label && labelsToRemove.has(deferArgs.label)));
+    if (updatedSubSelections === this.selectionSet && !hasDeferToRemove) {
+      return this;
+    }
+
+    const newFragment = hasDeferToRemove ? this.fragmentElement.withoutDefer() : this.fragmentElement;
+    if (!newFragment) {
+      return updatedSubSelections;
+    }
+    return new InlineFragmentSelection(newFragment, updatedSubSelections);
+  }
+
+  withNormalizedDefer(normalizer: DeferNormalizer): InlineFragmentSelection | SelectionSet {
+    const newFragment = this.fragmentElement.withNormalizedDefer(normalizer);
+    const updatedSubSelections = this.selectionSet.withNormalizedDefer(normalizer);
+    if (!newFragment) {
+      return updatedSubSelections;
+    }
+    return newFragment === this.fragmentElement && updatedSubSelections === this.selectionSet
+      ? this
+      : new InlineFragmentSelection(newFragment, updatedSubSelections);
+  }
+
   toString(expandFragments: boolean = true, indent?: string): string {
     return (indent ?? '') + this.fragmentElement + ' ' + this.selectionSet.toString(expandFragments, true, indent);
   }
@@ -1412,7 +1847,9 @@ class FragmentSpreadSelection extends FragmentSelection {
   }
 
   validate(): void {
-    // We don't do anything because fragment definition are validated when created.
+    this.validateDeferAndStream();
+
+    // We don't do anything else because fragment definition are validated when created.
   }
 
   toSelectionNode(): FragmentSpreadNode {
@@ -1440,11 +1877,15 @@ class FragmentSpreadSelection extends FragmentSelection {
     return this;
   }
 
-  expandFragments(names?: string[], updateSelectionSetFragments: boolean = true): FragmentSelection {
+  expandFragments(names?: string[], updateSelectionSetFragments: boolean = true): FragmentSelection | readonly Selection[] {
     if (names && !names.includes(this.namedFragment.name)) {
       return this;
     }
-    return new InlineFragmentSelection(this._element, this.selectionSet.expandFragments(names, updateSelectionSetFragments));
+
+    const expandedSubSelections = this.selectionSet.expandFragments(names, updateSelectionSetFragments);
+    return sameType(this._element.parentType, this.namedFragment.typeCondition)
+      ? expandedSubSelections.selections()
+      : new InlineFragmentSelection(this._element, expandedSubSelections);
   }
 
   collectUsedFragmentNames(collector: Map<string, number>): void {
@@ -1453,6 +1894,14 @@ class FragmentSpreadSelection extends FragmentSelection {
     collector.set(this.namedFragment.name, usageCount === undefined ? 1 : usageCount + 1);
   }
 
+  withoutDefer(_labelsToRemove?: Set<string>): FragmentSelection {
+    assert(false, 'Unsupported, see `Operation.withoutDefer`');
+  }
+
+  withNormalizedDefer(_normalizezr: DeferNormalizer): FragmentSelection {
+    assert(false, 'Unsupported, see `Operation.withAllDeferLabelled`');
+  }
+
   private spreadDirectives(): Directive<FragmentElement>[] {
     return this._element.appliedDirectives.slice(this.namedFragment.appliedDirectives.length);
   }