@apollo/query-graphs/src/querygraph.ts

Apollo Federation library to work with 'query graphs'

import {
  assert,
  MultiMap,
  InterfaceType,
  isInterfaceType,
  isFed1Supergraph,
  isObjectType,
  isUnionType,
  NamedType,
  ObjectType,
  Schema,
  SchemaRootKind,
  Type,
  UnionType,
  baseType,
  SelectionSet,
  isFederationSubgraphSchema,
  FieldDefinition,
  isCompositeType,
  parseFieldSetArgument,
  AbstractType,
  isAbstractType,
  possibleRuntimeTypes,
  MapWithCachedArrays,
  mapKeys,
  firstOf,
  FEDERATION_RESERVED_SUBGRAPH_NAME,
  federationMetadata,
  FederationMetadata,
  DirectiveDefinition,
  Directive,
  typenameFieldName,
  Field,
  selectionSetOfElement,
  SelectionSetUpdates,
  Supergraph,
  NamedSchemaElement,
  validateSupergraph,
  parseContext,
} from '@apollo/federation-internals';
import { inspect } from 'util';
import { DownCast, FieldCollection, subgraphEnteringTransition, SubgraphEnteringTransition, Transition, KeyResolution, RootTypeResolution, InterfaceObjectFakeDownCast } from './transition';
import { preComputeNonTrivialFollowupEdges } from './nonTrivialEdgePrecomputing';
import { NonLocalSelectionsMetadata } from './nonLocalSelectionsEstimation';

// We use our federation reserved subgraph name to avoid risk of conflict with other subgraph names (wouldn't be a huge
// deal, but safer that way). Using something short like `_` is also on purpose: it makes it stand out in debug messages
// without taking space.
export const FEDERATED_GRAPH_ROOT_SOURCE = FEDERATION_RESERVED_SUBGRAPH_NAME;
const FEDERATED_GRAPH_ROOT_SCHEMA = new Schema();

export function federatedGraphRootTypeName(rootKind: SchemaRootKind): string {
  return `[${rootKind}]`;
}

export function isFederatedGraphRootType(type: NamedType) {
  return type.name.startsWith('[') && type.name.endsWith(']');
}

/**
 * A vertex of a query graph, which points to a type (definition) in a particular graphQL schema (the `source` being
 * an identifier for that schema).
 *
 * @see QueryGraph
 */
export class Vertex {
  hasReachableCrossSubgraphEdges: boolean = false;
  // @provides works by creating duplicates of the vertex/type involved in the provides and adding the provided
  // edges only to those copy. This means that with @provides, you can have more than one vertex per-type-and-subgraph
  // in a query graph. Which is fined, but this `provideId` allows to distinguish if a vertex was created as part of
  // this @provides duplication or not. The value of this field has no other meaning than to be unique per-@provide,
  // and so all the vertex copied for a given @provides application will have the same `provideId`. Overall, this
  // mostly exists for debugging visualization.
  provideId: number | undefined;

  constructor(
    /** Index used for this vertex in the query graph it is part of. */
    readonly index: number,
    /** The graphQL type the vertex points to. */
    readonly type: NamedType,
    /**
     * An identifier of the underlying schema containing the `type` this vertex points to.
     * This is mainly used in "federated" query graphs, where the `source` is a subgraph name.
     */
    readonly source : string
  ) {}

  toString(): string {
    const label = `${this.type}(${this.source})`;
    return this.provideId ? `${label}-${this.provideId}` : label;
  }
}

/**
 * A "root" `Vertex`, that is a vertex that one of the root of a query graph.
 *
 * @see Vertex
 * @see QueryGraph.roots
 */
export class RootVertex extends Vertex {
  constructor(
    readonly rootKind: SchemaRootKind,
    index: number,
    type: NamedType,
    source : string
  ) {
    super(index, type, source);
  }

  toString(): string {
    return super.toString() + '*';
  }
}

function toRootVertex(vertex: Vertex, rootKind: SchemaRootKind): RootVertex {
  return new RootVertex(rootKind, vertex.index, vertex.type, vertex.source);
}

export function isRootVertex(vertex: Vertex): vertex is RootVertex {
  return vertex instanceof RootVertex;
}

export interface OverrideCondition {
  label: string;
  condition: boolean;
}

export function checkOverrideCondition(
  overrideCondition: OverrideCondition,
  conditionsToCheck: Map<string, boolean>
): boolean {
  const { label, condition } = overrideCondition;
  return conditionsToCheck.has(label) ? conditionsToCheck.get(label) === condition : false;
}

export type ContextCondition = {
  context: string;
  subgraphName: string;
  namedParameter: string;
  selection: string;
  typesWithContextSet: Set<string>;
  argType: Type,
  coordinate: string;
}

/**
 * An edge of a query graph.
 *
 * Query graphs are directed and an edge goes from its `head` vertex to its `tail` one.
 * Edges also have additional metadata: their `transition` and, optionally, `conditions`.
 */
export class Edge {
  private _conditions?: SelectionSet;
  
  public requiredContexts: ContextCondition[] = [];
  
  constructor(
    /**
     * Index used for this edge in the query graph it is part of (note that this index is "scoped" within
     * the head vertex, meaning that if 2 different vertices of the same query graph both have a single
     * out-edge, then both of those edges have index 0, and if a vertex has 3 out-edges, their index will
      * be 0, 1 and 2).
     */
    public readonly index: number,
    /**
     * The vertex from which the edge starts.
     */
    public readonly head: Vertex,
    /**
     * The vertex on which the edge ends.
     */
    public readonly tail: Vertex,
    /**
     * Indicates what kind of edges this is and what the edges does/represents.
     * For instance, if the edge represents a field, the `transition` will be a `FieldCollection` transition
     * and will link to the definition of the field it represents.
     *
     * @see Transition
     */
    public readonly transition: Transition,
    /**
     * Optional conditions on an edge.
     *
     * Conditions are a select of selections (in the graphQL sense) that the traversal of a query graph
     * needs to "collect" (traverse edges with transitions corresponding to those selections) in order
     * to be able to collect that edge.
     *
     * Conditions are primarily used for edges corresponding to @key, in which case they correspond
     * to the fields composing the key. In other words, for a key edges, conditions basically represents
     * the fact that you need the key to be able to use a key edge.
     *
     * Outside of keys, @requires also rely on conditions.
     */
    conditions?: SelectionSet,
    /**
     * Edges can require that an override condition (provided during query
     * planning) be met in order to be taken. This is used for progressive
     * @override, where (at least) 2 subgraphs can resolve the same field, but
     * one of them has an @override with a label. If the override condition
     * matches the query plan parameters, this edge can be taken.
     */
    public overrideCondition?: OverrideCondition,
    
    /**
     * Potentially multiple context conditions. When @fromContext is used on a argument definition, the edge representing
     * the argument's field needs to reflect that the condition must be satisfied in order for the edge to be taken
     */
    requiredContexts?: ContextCondition[],
  ) {
    this._conditions = conditions;
    if (requiredContexts) {
      this.requiredContexts = [...requiredContexts];
    }
  }

  get conditions(): SelectionSet | undefined {
    return this._conditions;
  }

  isEdgeForField(name: string): boolean {
    return this.transition.kind === 'FieldCollection' && this.transition.definition.name === name;
  }

  matchesSupergraphTransition(otherTransition: Transition): boolean {
    assert(otherTransition.collectOperationElements, () => `Supergraphs shouldn't have transition that don't collect elements; got ${otherTransition}"`);
    const transition = this.transition;
    switch (transition.kind) {
      case 'FieldCollection': return otherTransition.kind === 'FieldCollection' && transition.definition.name === otherTransition.definition.name;
      case 'DownCast': return otherTransition.kind === 'DownCast' && transition.castedType.name === otherTransition.castedType.name;
      case 'InterfaceObjectFakeDownCast': return otherTransition.kind === 'DownCast' && transition.castedTypeName === otherTransition.castedType.name;
      default: return false;
    }
  }

  changesSubgraph(): boolean {
    return this.head.source !== this.tail.source;
  }

  label(): string {
    if (this.transition instanceof SubgraphEnteringTransition && !this._conditions) {
      return "";
    }

    let conditionsString = (this._conditions ?? '').toString();
    if (this.overrideCondition) {
      if (conditionsString.length) conditionsString += ', ';
      conditionsString += `${this.overrideCondition.label} = ${this.overrideCondition.condition}`;
    }
    // we had at least some condition, add the turnstile and spacing
    if (conditionsString.length) conditionsString += ' ⊢ ';

    return conditionsString + this.transition.toString();
  }

  withNewHead(newHead: Vertex): Edge {
    return new Edge(
      this.index,
      newHead,
      this.tail,
      this.transition,
      this._conditions,
      this.overrideCondition,
      this.requiredContexts,
    );
  }

  addToConditions(newConditions: SelectionSet) {
    this._conditions = this._conditions
      ? new SelectionSetUpdates().add(this._conditions).add(newConditions).toSelectionSet(this._conditions.parentType)
      : newConditions;
  }
  
  addToContextConditions(contextConditions: ContextCondition[]) {
    this.requiredContexts.push(...contextConditions);
  }

  isKeyOrRootTypeEdgeToSelf(): boolean {
    return this.head === this.tail && (this.transition.kind === 'KeyResolution' || this.transition.kind === 'RootTypeResolution');
  }

  satisfiesOverrideConditions(conditionsToCheck: Map<string, boolean>) {
    if (!this.overrideCondition) return true;
    return checkOverrideCondition(
      this.overrideCondition,
      conditionsToCheck,
    );
  }

  toString(): string {
    return `${this.head} -> ${this.tail} (${this.label()})`;
  }
}

/**
 * An immutable directed graph data structure (built of vertices and edges) that is layered over one or multiple
 * graphQL schema, that aims to facilitate reasoning about queries expressed on the underlying schema.
 *
 * On top of its set of vertices and edges, a query graph exposes:
 *  - its set of "sources": pointers to the graphQL schema on which the query graph was built.
 *  - a set of distinguished vertices called "root" vertices. A query graph has at most one root
 *    vertex per `SchemaRootKind`, and those vertices are usually entry points for traversals of
 *    a query graph.
 *
 * In practice, the code builds 2 "kind" of query graphs:
 *  1. a supergraph query graph, which is built on top of a supergraph API schema (@see buildGraph()),
 *     and is built to reason about user queries (made on the supergraph API). This supergraph query
 *     graph is used to validate composition.
 *  2. a "federated" query graph, which is a single graph built on top of a) a number of subgraph
 *    API schema and b) the additional federation directives on those subgraphs (@see buildFederatedQueryGraph()).
 *    This query graph is used both for validating composition and for query planning.
 *
 * Note that this class handles both cases, but a "supergraph query graph" will, by construction, have
 * a few more invariants than a "federated query graph". Including (but not necessarily limited to):
 *  - edges in a super graph query graph will never have `conditions` or 'key' edges (edges with a `KeyResolution` edges).
 *  - supergraph query graphs will have a single value in `sources` (the supergraph schema).
 *
 * Also note that as query graphs are about reasoning about queries over schema, they only contain vertices
 * that points to "reachable" types (reachable from any kind of operations).
 */
export class QueryGraph {
  /**
   * Given an edge, returns the possible edges that can follow it "productively", that is without creating
   * a trivially inefficient path.
   *
   * More precisely, `nonTrivialFollowupEdges(e)` is equivalent calling `outEdges(e.tail)` and filtering
   * the edges that "never make sense" after `e`, which mainly amounts to avoiding chaining key edges
   * when we know there is guaranteed to be a better option. As an example, suppose we have 3 subgraphs
   * A, B and C which all defined a `@key(fields: "id")` on some entity type `T`. Then it is never
   * interesting to take that key edge from B -> C after A -> B because if we're in A and want to get
   * to C, we can always do A -> C (of course, this is only true because it's the "same" key).
   *
   * See `preComputeNonTrivialFollowupEdges` for more details on which exact edges are filtered.
   *
   * Lastly, note that the main reason for exposing this method is that its result is pre-computed.
   * Which in turn is done for performance reasons: having the same key defined in multiple subgraphs
   * is _the_ most common pattern, and while our later algorithms (composition validation and query
   * planning) would know to not select those trivially inefficient "detour", they might have to redo
   * those checks many times and pre-computing once it is significantly faster (and pretty easy).
   * Fwiw, when originally introduced, this optimization lowered composition validation on a big
   * composition (100+ subgraphs) from ~4 "minutes" to ~10 seconds.
   */
  readonly nonTrivialFollowupEdges: (edge: Edge) => readonly Edge[];
  /**
   * To speed up the estimation of counting non-local selections, we
   * precompute specific metadata. We only computed this for federated query
   * graphs used during query planning.
   */
  readonly nonLocalSelectionsMetadata: NonLocalSelectionsMetadata | null;

  /**
   * Creates a new query graph.
   *
   * This isn't meant to be be called directly outside of `GraphBuilder.build`.
   */
  constructor(
    /** A name to identify the graph. Mostly for pretty-printing/debugging purpose. */
    readonly name: string,
    /** The vertices of the query graph. The index of each vertex in the array will be the value of its `Vertex.index` value. */
    readonly vertices: Vertex[],
    /**
    * For each vertex, the edges that originate from that array. This array has the same length as `vertices` and `_outEdges[i]`
    * is an array of all the edges starting at vertices[i].
    */
    private readonly _outEdges: Edge[][],
    /**
     * A map that associate type names of the underlying schema on which this query graph was built to each of the vertex
     * (vertex index) that points to a type of that name. Note that in a "supergraph query graph", each type name will only
     * map to a single vertex,
     */
    private readonly typesToVertices: MultiMap<string, number>,
    /** The set of distinguished root vertices (pointers to vertices in `vertices`), keyed by their kind.  */
    private readonly rootVertices: MapWithCachedArrays<SchemaRootKind, RootVertex>,
    /**
     * The sources on which the query graph was built, that is a set (which can be of size 1) of graphQL schema keyed by
     * the name identifying them. Note that the `source` string in the `Vertex` of a query graph is guaranteed to be
     * valid key in this map.
     */
    readonly sources: ReadonlyMap<string, Schema>,
    
    readonly subgraphToArgs: Map<string, string[]>,
    
    readonly subgraphToArgIndices: Map<string, Map<string, string>>,

    readonly schema: Schema,

    isFederatedAndForQueryPlanning?: boolean,
  ) {
    this.nonTrivialFollowupEdges = preComputeNonTrivialFollowupEdges(this);
    this.nonLocalSelectionsMetadata = isFederatedAndForQueryPlanning
     ? new NonLocalSelectionsMetadata(this)
     : null;
  }

  /** The number of vertices in this query graph. */
  verticesCount(): number {
    return this.vertices.length;
  }

  /** The number of edges in this query graph. */
  edgesCount(): number {
    // We could count edges as we add them and pass it to the ctor. For now though, it's not meant to be
    // on a hot path, so recomputing is probably fine.
    return this._outEdges.reduce((acc, v) => acc + v.length, 0);
  }

  /**
   * The set of `SchemaRootKind` for which this query graph has a root vertex (for
   * which `root(SchemaRootKind)` will _not_ return `undefined`).
   */
  rootKinds(): readonly SchemaRootKind[] {
    return this.rootVertices.keys();
  }

  /**
   * The set of root vertices in this query graph.
   */
  roots(): readonly RootVertex[] {
    return this.rootVertices.values();
  }

  /**
   * The root vertex corresponding to the provided root kind, if this query graph has one.
   *
   * @param kind - the root kind for which to get the root vertex.
   * @returns the root vertex for `kind` if this query graph has one.
   */
  root(kind: SchemaRootKind): RootVertex | undefined {
    return this.rootVertices.get(kind);
  }

  /**
   * The edges out of the provided vertex.
   *
   * @param vertex - the vertex for which to return out edges. This method _assumes_ that
   *   the provided vertex is a vertex of this query graph (and its behavior is undefined
   *   if it isn't).
   * @param includeKeyAndRootTypeEdgesToSelf - whether key/root type edges that stay on the same
   *  vertex should be included. This default to `false` are those are rarely useful. More
   *   precisely, the only current use of them is for @defer where they may be needed to re-enter
   *   the current subgraph in a deferred section.
   * @returns the list of all the edges out of this vertex.
   */
  outEdges(vertex: Vertex, includeKeyAndRootTypeEdgesToSelf: boolean = false): readonly Edge[] {
    const allEdges = this._outEdges[vertex.index];
    return includeKeyAndRootTypeEdgesToSelf ? allEdges : allEdges.filter((e) => !e.isKeyOrRootTypeEdgeToSelf())
  }

  /**
   * The number of edges out of the provided vertex.
   *
   * This is a shortcut for `this.outEdges(vertex, true).length`, and the reason it considers
   * edge-to-self by default while `this.outEdges` doesn't is that this method is generally
   * used to size other arrays indexed by edges index, and so we want to consider all edges
   * in general.
   */
  outEdgesCount(vertex: Vertex): number {
    return this._outEdges[vertex.index].length;
  }

  /**
   * The edge out of the provided vertex at the provided (edge) index.
   *
   * @param vertex - the vertex for which to return the out edge. This method _assumes_ that
   *   the provided vertex is a vertex of this query graph (and its behavior is undefined
   *   if it isn't).
   * @param edgeIndex - the edge index (relative to `vertex`, see `Edge.index`) to retrieve.
   * @returns the `edgeIndex`^th edge out of `vertex`, if such edges exists.
   */
  outEdge(vertex: Vertex, edgeIndex: number): Edge | undefined {
    return this._outEdges[vertex.index][edgeIndex];
  }

  allVertices(): Iterable<Vertex> {
    return this.vertices;
  }

  *allEdges(): Iterable<Edge> {
    for (const vertexOutEdges of this._outEdges) {
      for (const outEdge of vertexOutEdges) {
        yield outEdge;
      }
    }
  }

  /**
   * Whether the provided vertex is a terminal one (has no out edges).
   *
   * @param vertex - the vertex to check.
   * @returns whether the provided vertex is terminal.
   */
  isTerminal(vertex: Vertex): boolean {
    return this.outEdgesCount(vertex) === 0;
  }

  /**
   * The set of vertices whose associated type (see `Vertex.type`) is of type `typeName`.
   */
  verticesForType(typeName: string): Vertex[] {
    const indexes = this.typesToVertices.get(typeName);
    return indexes == undefined ? [] : indexes.map(i => this.vertices[i]);
  }
}

/**
 * A utility class that allows to associate state to the vertices and/or edges of a query graph.
 *
 * @param VertexState - the type of the state associated to vertices.
 * @param EdgeState - the type of the state associated to edges. Defaults to `undefined`, which
 *   means that state is only associated to vertices.
 */
export class QueryGraphState<VertexState, EdgeState = undefined> {
  // Store some "user" state for each vertex (accessed by index)
  private readonly verticesStates: Map<number, VertexState> = new Map();
  private readonly adjacenciesStates: Map<number, Map<number, EdgeState>> = new Map();

  /**
   * Associates the provided state to the provided vertex.
   *
   * @param vertex - the vertex to which state should be associated. This method _assumes_
   *    that the provided vertex is a vertex of the query graph against which this
   *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
   * @param state - the state/value to associate to `vertex`.
   */
  setVertexState(vertex: Vertex, state: VertexState) {
    this.verticesStates.set(vertex.index, state);
  }

  /**
   * Removes the state associated to the provided vertex (if any is).
   *
   * @param vertex - the vertex for which state should be removed. This method _assumes_
   *    that the provided vertex is a vertex of the query graph against which this
   *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
   */
  removeVertexState(vertex: Vertex) {
    this.verticesStates.delete(vertex.index);
  }

  /**
   * Retrieves the state associated to the provided vertex (if any is).
   *
   * @param vertex - the vertex for which state should be retrieved. This method _assumes_
   *    that the provided vertex is a vertex of the query graph against which this
   *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
   * @return the state associated to `vertex`, if any.
   */
  getVertexState(vertex: Vertex): VertexState | undefined {
    return this.verticesStates.get(vertex.index);
  }

  /**
   * Associates the provided state to the provided edge.
   *
   * @param edge - the edge to which state should be associated. This method _assumes_
   *    that the provided edge is an edge of the query graph against which this
   *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
   * @param state - the state/value to associate to `edge`.
   */
  setEdgeState(edge: Edge, state: EdgeState) {
    let edgeMap = this.adjacenciesStates.get(edge.head.index)
    if (!edgeMap) {
      edgeMap = new Map();
      this.adjacenciesStates.set(edge.head.index, edgeMap);
    }
    edgeMap.set(edge.index, state);
  }

  /**
   * Removes the state associated to the provided edge (if any is).
   *
   * @param edge - the edge for which state should be removed. This method _assumes_
   *    that the provided edge is an edge of the query graph against which this
   *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
   */
  removeEdgeState(edge: Edge) {
    const edgeMap = this.adjacenciesStates.get(edge.head.index);
    if (edgeMap) {
      edgeMap.delete(edge.index);
      if (edgeMap.size === 0) {
        this.adjacenciesStates.delete(edge.head.index);
      }
    }
  }

  /**
   * Retrieves the state associated to the provided edge (if any is).
   *
   * @param edge - the edge for which state should be retrieved. This method _assumes_
   *    that the provided vertex is an edge of the query graph against which this
   *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
   * @return the state associated to `edge`, if any.
   */
  getEdgeState(edge: Edge): EdgeState | undefined {
    return this.adjacenciesStates.get(edge.head.index)?.get(edge.index);
  }

  toDebugString(
    vertexMapper: (s: VertexState) => string,
    edgeMapper: (e: EdgeState) => string
  ): string {
    const vs = Array.from(this.verticesStates.entries()).sort(([a], [b]) => a - b).map(([idx, state]) =>
      ` ${idx}: ${!state ? "<null>" : vertexMapper(state)}`
    ).join("\n");
    const es = Array.from(this.adjacenciesStates.entries()).sort(([a], [b]) => a - b).map(([vIdx, adj]) =>
      Array.from(adj.entries()).sort(([a], [b]) => a - b).map(([eIdx, state]) =>
        ` ${vIdx}[${eIdx}]: ${!state ? "<null>" : edgeMapper(state)}`
      ).join("\n")
    ).join("\n");
    return `vertices = {${vs}\n}, edges = {${es}\n}`;
  }
}

/**
 * Builds the query graph corresponding to the provided schema.
 *
 * Note that this method and mainly exported for the sake of testing but should rarely, if
 * ever, be used otherwise. Instead use either `buildSupergraphAPIQueryGraph` or
 * `buildFederatedQueryGraph` which are more explicit.
 *
 * @param name - the name to use for the created graph and as "source" name for the schema.
 * @param schema - the schema for which to build the query graph.
 * @param overrideLabelsByCoordinate - A Map of coordinate -> override label to apply to the query graph.
 *   Additional "virtual" edges will be created for progressively overridden fields in order to ensure that
 *   all possibilities are considered during query planning.
 * @returns the query graph corresponding to `schema` "API" (in the sense that no federation
 *   directives are taken into account by this method in the building of the query graph).
 */
export function buildQueryGraph(name: string, schema: Schema, overrideLabelsByCoordinate?: Map<string, string>): QueryGraph {
  return buildGraphInternal(name, schema, false, undefined, overrideLabelsByCoordinate);
}

function buildGraphInternal(
  name: string,
  schema: Schema,
  addAdditionalAbstractTypeEdges: boolean,
  supergraphSchema?: Schema,
  overrideLabelsByCoordinate?: Map<string, string>,
): QueryGraph {
  const builder = new GraphBuilderFromSchema(
    name,
    schema,
    supergraphSchema ? { apiSchema: supergraphSchema.toAPISchema(), isFed1: isFed1Supergraph(supergraphSchema) } : undefined,
    overrideLabelsByCoordinate,
  );
  for (const rootType of schema.schemaDefinition.roots()) {
    builder.addRecursivelyFromRoot(rootType.rootKind, rootType.type);
  }
  if (builder.isFederatedSubgraph) {
    builder.addInterfaceEntityEdges();
  }
  if (addAdditionalAbstractTypeEdges) {
    builder.addAdditionalAbstractTypeEdges();
  }
  return builder.build();
}

/**
 * Builds a "supergraph API" query graph based on the provided supergraph schema.
 *
 * A "supergraph API" query graph is one that is used to reason about queries against said
 * supergraph API, but @see QueryGraph for more details.
 *
 * @param supergraph - the schema of the supergraph for which to build the query graph.
  *  The provided schema should generally be a "supergraph" as generated by composition merging.
  *  Note however that the query graph built by this method is only based on the supergraph
  *  API and doesn't rely on the join spec directives, so it is valid to also directly
  *  pass a schema that directly corresponds to the supergraph API.
 * @returns the built query graph.
 */
export function buildSupergraphAPIQueryGraph(supergraph: Supergraph): QueryGraph {
  const apiSchema = supergraph.apiSchema();

  const overrideLabelsByCoordinate = new Map<string, string>();
  const joinFieldApplications = validateSupergraph(supergraph.schema)[1]
    .fieldDirective(supergraph.schema).applications();
  for (const application of joinFieldApplications) {
    const overrideLabel = application.arguments().overrideLabel;
    if (overrideLabel) {
      overrideLabelsByCoordinate.set(
        (application.parent as FieldDefinition<any>).coordinate,
        overrideLabel
      );
    }
  }
  return buildQueryGraph("supergraph", apiSchema, overrideLabelsByCoordinate);
}

/**
 * Builds a "federated" query graph based on the provided supergraph schema.
 *
 * A "federated" query graph is one that is used to reason about queries made by a
 * gateway/router against a set of federated subgraph services.
 *
 * @see QueryGraph
 *
 * @param supergraph - the supergraph for which to build the query graph.
 * @param forQueryPlanning - whether the build query graph is built for query planning (if
 *   so, it will include some additional edges that don't impact validation but allow
 *   to generate more efficient query plans).
 * @returns the built federated query graph.
 */
export function buildFederatedQueryGraph(supergraph: Supergraph, forQueryPlanning: boolean): QueryGraph {
  const subgraphs = supergraph.subgraphs();
  const graphs = [];
  for (const subgraph of subgraphs) {
    graphs.push(buildGraphInternal(subgraph.name, subgraph.schema, forQueryPlanning, supergraph.schema));
  }
  return federateSubgraphs(supergraph.schema, graphs, forQueryPlanning);
}

function federatedProperties(subgraphs: QueryGraph[]) : [number, Set<SchemaRootKind>, Schema[]] {
  let vertices = 0;
  const rootKinds = new Set<SchemaRootKind>();
  const schemas: Schema[] = [];
  for (const subgraph of subgraphs) {
    vertices += subgraph.verticesCount();
    subgraph.rootKinds().forEach(k => rootKinds.add(k));
    assert(subgraph.sources.size === 1, () => `Subgraphs should only have one sources, got ${subgraph.sources.size} ([${mapKeys(subgraph.sources).join(', ')}])`);
    schemas.push(firstOf(subgraph.sources.values())!);
  }
  return [vertices + rootKinds.size, rootKinds, schemas];
}

function resolvableKeyApplications(
  keyDirective: DirectiveDefinition<{fields: any, resolvable?: boolean}>,
  type: NamedType
): Directive<NamedType, {fields: any, resolvable?: boolean}>[] {
  const applications: Directive<NamedType, {fields: any, resolvable?: boolean}>[] = type.appliedDirectivesOf(keyDirective);
  return applications.filter((application) => application.arguments().resolvable ?? true);
}

function federateSubgraphs(
  supergraph: Schema,
  subgraphs: QueryGraph[],
  forQueryPlanning: boolean,
): QueryGraph {
  const [verticesCount, rootKinds, schemas] = federatedProperties(subgraphs);
  const builder = new GraphBuilder(supergraph, verticesCount);
  rootKinds.forEach(k => builder.createRootVertex(
    k,
    new ObjectType(federatedGraphRootTypeName(k)),
    FEDERATED_GRAPH_ROOT_SOURCE,
    FEDERATED_GRAPH_ROOT_SCHEMA
  ));

  // We first add all the vertices and edges from the subgraphs
  const copyPointers: SubgraphCopyPointer[] = new Array(subgraphs.length);
  for (const [i, subgraph] of subgraphs.entries()) {
    copyPointers[i] = builder.copyGraph(subgraph);
  }

  // We then add the edges from supergraph roots to the subgraph ones.
  // Also, for each root kind, we also add edges from the corresponding root type of each subgraph to the root type of other subgraphs
  // (and for @defer, like for @key, we also add self-link looping on the current subgraph).
  // This essentially encode the fact that if a field return a root type, we can always query any subgraph from that point.
  for (const [i, subgraph] of subgraphs.entries()) {
    const copyPointer = copyPointers[i];
    for (const rootKind of subgraph.rootKinds()) {
      const rootVertex = copyPointer.copiedVertex(subgraph.root(rootKind)!);
      builder.addEdge(builder.root(rootKind)!, rootVertex, subgraphEnteringTransition)

      for (const [j, otherSubgraph] of subgraphs.entries()) {
        const otherRootVertex = otherSubgraph.root(rootKind);
        if (otherRootVertex) {
          const otherCopyPointer = copyPointers[j];
          builder.addEdge(rootVertex, otherCopyPointer.copiedVertex(otherRootVertex), new RootTypeResolution(rootKind));
        }
      }
    }
  }

  // Then we add/update edges for @key and @requires. We do @provides in a second step because its handling requires
  // copying vertex and their edges, and it's easier to reason about this if we know all keys have already been created.
  for (const [i, subgraph] of subgraphs.entries()) {
    const subgraphSchema = schemas[i];
    const subgraphMetadata = federationMetadata(subgraphSchema);
    assert(subgraphMetadata, `Subgraph ${i} is not a valid federation subgraph`);
    const keyDirective = subgraphMetadata.keyDirective();
    const requireDirective = subgraphMetadata.requiresDirective();
    simpleTraversal(
      subgraph,
      v => {
        const type = v.type;
        for (const keyApplication of resolvableKeyApplications(keyDirective, type)) {
          // The @key directive creates an edge from every subgraphs (having that type)
          // to the current subgraph. In other words, the fact this subgraph has a @key means
          // that the service of the current subgraph can be queried for the entity (through
          // _entities) as long as "the other side" can provide the proper field values.
          // Note that we only require that "the other side" can gather the key fields (through
          // the path conditions; note that it's possible those conditions are never satisfiable),
          // but don't care that it defines the same key, because it's not a technical
          // requirement (and while we probably don't want to allow in general a type to be an
          // entity in some subgraphs but not other, this is not the place to impose that
          // restriction, and this may be useful at least temporarily to allow convert a type to
          // an entity).
          assert(isInterfaceType(type) || isObjectType(type), () => `Invalid "@key" application on non Object || Interface type "${type}"`);
          const isInterfaceObject = subgraphMetadata.isInterfaceObjectType(type);
          const conditions = parseFieldSetArgument({ parentType: type, directive: keyApplication, normalize: true });

          // We'll look at adding edges from "other subgraphs" to the current type. So the tail of all the edges
          // we'll build in this branch is always going to be the same.
          const tail = copyPointers[i].copiedVertex(v);
          // Note that each subgraph has a key edge to itself (when i === j below). We usually ignore
          // this edges, but they exists for the special case of @defer, where we technically may have
          // to take such "edge-to-self" as a mean to "re-enter" a subgraph for a deferred section.
          for (const [j, otherSubgraph] of subgraphs.entries()) {
            const otherVertices = otherSubgraph.verticesForType(type.name);
            if (otherVertices.length > 0) {
              // Note that later, when we've handled @provides, this might not be true anymore as @provides may create copy of a
              // certain type. But for now, it's true.
              assert(
                otherVertices.length == 1,
                () => `Subgraph ${j} should have a single vertex for type ${type.name} but got ${otherVertices.length}: ${inspect(otherVertices)}`);

              const otherVertex = otherVertices[0];
              // The edge goes from the otherSubgraph to the current one.
              const head = copyPointers[j].copiedVertex(otherVertex);
              const tail = copyPointers[i].copiedVertex(v);
              builder.addEdge(head, tail, new KeyResolution(), conditions);
            }

            // Additionally, if the key is on an @interfaceObject and this "other" subgraph has some of the implementations
            // of the corresponding interface, then we need an edge from each of those implementations (to the @interfaceObject).
            // This is used when an entity of specific implementation is queried first, but then some of the
            // requested fields are only provided by that @interfaceObject.
            if (isInterfaceObject) {
              const typeInSupergraph = supergraph.type(type.name);
              assert(typeInSupergraph && isInterfaceType(typeInSupergraph), () => `Type ${type} is an interfaceObject in subgraph ${i}; should be an interface in the supergraph`);
              for (const implemTypeInSupergraph of typeInSupergraph.possibleRuntimeTypes()) {
                // That implementation type may or may not exists in "otherSubgraph". If it doesn't, we just have nothing to
                // do for that particular impelmentation. If it does, we'll add the proper edge, but note that we're guaranteed
                // to have at most one vertex for the same reason than mentioned above (only the handling @provides will make it
                // so that there can be more than one vertex per type).
                const implemVertice = otherSubgraph.verticesForType(implemTypeInSupergraph.name)[0];
                if (!implemVertice) {
                  continue;
                }

                const implemHead = copyPointers[j].copiedVertex(implemVertice);
                // The key goes from the implementation type to the @interfaceObject one, so the conditions
                // will be "fetched" on the implementation type, but `conditions` has been parsed on the
                // interface type, so it will use fields from the interface, not the implementation type.
                // So we re-parse the condition using the implementation type: this could fail, but in
                // that case it just mean that key is not usable.
                const implemType = implemVertice.type;
                assert(isCompositeType(implemType), () => `${implemType} should be composite since it implements ${typeInSupergraph} in the supergraph`);
                try {
                  const implConditions = parseFieldSetArgument({ parentType: implemType, directive: keyApplication, validate: false, normalize: true });
                  builder.addEdge(implemHead, tail, new KeyResolution(), implConditions);
                } catch (e) {
                  // Ignored on purpose: it just means the key is not usable on this subgraph.
                }
              }
            }
          }
        }
      },
      e => {
        // Handling @requires
        if (e.transition.kind === 'FieldCollection') {
          const type = e.head.type;
          const field = e.transition.definition;
          assert(isCompositeType(type), () => `Non composite type "${type}" should not have field collection edge ${e}`);
          for (const requiresApplication of field.appliedDirectivesOf(requireDirective)) {
            const conditions = parseFieldSetArgument({ parentType: type, directive: requiresApplication, normalize: true });
            const head = copyPointers[i].copiedVertex(e.head);
            // We rely on the fact that the edge indexes will be the same in the copied builder. But there is no real reason for
            // this to not be the case at this point so...
            const copiedEdge = builder.edge(head, e.index);
            copiedEdge.addToConditions(conditions);
          }
        }
        return true; // Always traverse edges
      }
    );
  }

  /**
   * Handling progressive overrides here. For each progressive @override
   * application (with a label), we want to update the edges to the overridden
   * field within the "to" and "from" subgraphs with their respective override
   * condition (the label and a T/F value). The "from" subgraph will have an
   * override condition of `false`, whereas the "to" subgraph will have an
   * override condition of `true`.
   */
  const subgraphsByName = new Map(subgraphs.map((s) => [s.name, s]));
  for (const [i, toSubgraph] of subgraphs.entries()) {
    const subgraphSchema = schemas[i];
    const subgraphMetadata = federationMetadata(subgraphSchema);
    assert(subgraphMetadata, `Subgraph ${i} is not a valid federation subgraph`);

    for (const application of subgraphMetadata.overrideDirective().applications()) {
      const { from, label } = application.arguments();
      if (!label) continue;
      const fromSubgraph = subgraphsByName.get(from);
      assert(fromSubgraph, () => `Subgraph ${from} not found`);

      function updateEdgeWithOverrideCondition(subgraph: QueryGraph, label: string, condition: boolean) {
        const field = application.parent;
        assert(field instanceof NamedSchemaElement, () => `@override should have been on a field, got ${field}`);
        const typeName = field.parent.name;

        const [vertex, ...unexpectedAdditionalVertices] = subgraph.verticesForType(typeName);
        assert(vertex && unexpectedAdditionalVertices.length === 0, () => `Subgraph ${subgraph.name} should have exactly one vertex for type ${typeName}`);

        const subgraphEdges = subgraph.outEdges(vertex);
        for (const edge of subgraphEdges) {
          if (
            edge.transition.kind === "FieldCollection"
            && edge.transition.definition.name === field.name
          ) {
            const head = copyPointers[subgraphs.indexOf(subgraph)].copiedVertex(vertex);
            const copiedEdge = builder.edge(head, edge.index);

            copiedEdge.overrideCondition = {
              label,
              condition,
            };
          }
        }
      }

      updateEdgeWithOverrideCondition(toSubgraph, label, true);
      updateEdgeWithOverrideCondition(fromSubgraph, label, false);
    }
  }
 
  /** 
   * Now we'll handle instances of @fromContext. For each argument with @fromContext, I want to add its corresponding 
   * context conditions to the edge corresponding to the argument's field
  */
  const subgraphToArgs: Map<string, string[]> = new Map();
  const subgraphToArgIndices: Map<string, Map<string, string>> = new Map();
  
  for (const [i, subgraph] of subgraphs.entries()) {
    const subgraphSchema = schemas[i];
    const subgraphMetadata = federationMetadata(subgraphSchema);
    assert(subgraphMetadata, `Subgraph ${i} is not a valid federation subgraph`);
    
    const contextNameToTypes: Map<string, Set<string>> = new Map();
    
    for (const application of subgraphMetadata.contextDirective().applications()) {
      const { name: context } = application.arguments();
        if (contextNameToTypes.has(context)) {
          contextNameToTypes.get(context)!.add(application.parent.name);
        } else {
          contextNameToTypes.set(context, new Set([application.parent.name]));
        }
    }
    
    const coordinateMap: Map<string, ContextCondition[]> = new Map();
    for (const application of subgraphMetadata.fromContextDirective().applications()) {
      const { field } = application.arguments();
      const { context, selection } = parseContext(field);
      
      assert(context, () => `FieldValue has invalid format. Context not found ${field}`);
      assert(selection, () => `FieldValue has invalid format. Selection not found ${field}`);
      const namedParameter = application.parent.name;
      const argCoordinate = application.parent.coordinate;
      const args = subgraphToArgs.get(subgraph.name) ?? [];
      args.push(argCoordinate);
      subgraphToArgs.set(subgraph.name, args);
      
      const fieldCoordinate = application.parent.parent.coordinate;
      const typesWithContextSet = contextNameToTypes.get(context);
      assert(typesWithContextSet, () => `Context ${context} is never set in subgraph`);
      const z = coordinateMap.get(fieldCoordinate);
      if (z) {
        z.push({ namedParameter, coordinate: argCoordinate, context, selection, typesWithContextSet, subgraphName: subgraph.name, argType: application.parent.type });
      } else {
        coordinateMap.set(fieldCoordinate, [{ namedParameter, coordinate: argCoordinate, context, selection, typesWithContextSet, subgraphName: subgraph.name, argType: application.parent.type }]);
      }
    }
    
    simpleTraversal(
      subgraph,
      _v => { return undefined; },
      e => {
        if (e.head.type.kind === 'ObjectType' && e.transition.kind === 'FieldCollection') {
          const coordinate = `${e.head.type.name}.${e.transition.definition.name}`;
          const requiredContexts = coordinateMap.get(coordinate);
          if (requiredContexts) {
            const headInSupergraph = copyPointers[i].copiedVertex(e.head);
            assert(headInSupergraph, () => `Vertex for type ${e.head.type.name} not found in supergraph`);
            const edgeInSupergraph = builder.edge(headInSupergraph, e.index);
            edgeInSupergraph.addToContextConditions(requiredContexts);
          }
        }
        return true;
      }
    );
   
  }

  // add contextual argument maps to builder
  for (const [i, subgraph] of subgraphs.entries()) {
    const subgraphName = subgraph.name;
    const args = subgraphToArgs.get(subgraph.name);
    if (args) {
      args.sort();
      const argToIndex = new Map();
      for (let idx=0; idx < args.length; idx++) {
        argToIndex.set(args[idx], `contextualArgument_${i+1}_${idx}`);
      }
      subgraphToArgIndices.set(subgraphName, argToIndex);
    }
  }
  
  builder.setContextMaps(subgraphToArgs, subgraphToArgIndices);
      
  // Now we handle @provides
  let provideId = 0;
  for (const [i, subgraph] of subgraphs.entries()) {
    const subgraphSchema = schemas[i];
    const subgraphMetadata = federationMetadata(subgraphSchema);
    assert(subgraphMetadata, `Subgraph ${i} is not a valid federation subgraph`);
    const providesDirective = subgraphMetadata.providesDirective();
    simpleTraversal(
      subgraph,
      _ => undefined,
      e => {
        // Handling @provides
        if (e.transition.kind === 'FieldCollection') {
          const type = e.head.type;
          const field = e.transition.definition;
          assert(isCompositeType(type), () => `Non composite type "${type}" should not have field collection edge ${e}`);
          for (const providesApplication of field.appliedDirectivesOf(providesDirective)) {
            ++provideId;
            const fieldType = baseType(field.type!);
            assert(isCompositeType(fieldType), () => `Invalid @provide on field "${field}" whose type "${fieldType}" is not a composite type`)
            const provided = parseFieldSetArgument({ parentType: fieldType, directive: providesApplication });
            const head = copyPointers[i].copiedVertex(e.head);
            const tail = copyPointers[i].copiedVertex(e.tail);
            // We rely on the fact that the edge indexes will be the same in the copied builder. But there is no real reason for
            // this to not be the case at this point so...
            const copiedEdge = builder.edge(head, e.index);
            // We make a copy of the `fieldType` vertex (with all the same edges), and we change this particular edge to point to the
            // new copy. From that, we can add all the provides edges to the copy.
            const copiedTail = builder.makeCopy(tail, provideId);
            builder.updateEdgeTail(copiedEdge, copiedTail);
            addProvidesEdges(subgraphSchema, builder, copiedTail, provided, provideId);
          }
        }
        return true; // Always traverse edges
      }
    );
  }

  // We now need to finish handling @interfaceObject types. More precisely, there is cases where only a/some implementation(s)
  // of a interface are queried, and that could apply to an interface that is an @interfaceObject in some sugraph. Consider
  // the following example:
  // ```graphql
  // type Query {
  //   getIs: [I]
  // }
  //
  // type I @key(fields: "id") @interfaceObject {
  //   id: ID!
  //   x: Int
  // }
  // ```
  // where we suppose that `I` has some implementations say, `A`, `B` and `C`, in some other subgraph.
  // Now, consider query:
  // ```graphql
  // {
  //   getIs {
  //     ... on B {
  //       x
  //     }
  //   }
  // }
  // ```
  // So here, we query `x` which the subgraph provides, but we only do so for one of the impelementation.
  // So in that case, we essentially need to figure out the `__typename` first (of more precisely, we need
  // to know the real __typename "eventually"; we could theoretically query `x` first, and then get the __typename
  // to know if we should keep the result or discard it, and that could be more efficient in certain case,
  // but as we don't know both 1) if `x` is expansive to resolve and 2) what the ratio of results from `getIs`
  // will be `B` versus some other implementation, it is "safer" to get the __typename first and only resolve `x`
  // when we need to).
  //
  // Long story short, to solve this, we create edges from @interfaceObject types to themselves for every implementation
  // types of the interface: those edges will be taken when we try to take a `... on B` condition, and those edge
  // have __typename as a condition, forcing to find __typename in another subgraph first.
  for (const [i, subgraph] of subgraphs.entries()) {
    const subgraphSchema = schemas[i];
    const subgraphMetadata = federationMetadata(subgraphSchema);
    assert(subgraphMetadata, `Subgraph ${i} is not a valid federation subgraph`);
    const interfaceObjectDirective = subgraphMetadata.interfaceObjectDirective();
    for (const application of interfaceObjectDirective.applications()) {
      const type = application.parent;
      assert(isObjectType(type), '@interfaceObject should have been on an object type');
      const vertex = copyPointers[i].copiedVertex(subgraph.verticesForType(type.name)[0]);
      const supergraphItf = supergraph.type(type.name);
      assert(supergraphItf && isInterfaceType(supergraphItf), () => `${type} has @interfaceObject in subgraph but has kind ${supergraphItf?.kind} in supergraph`)
      const condition = selectionSetOfElement(new Field(type.typenameField()!));
      for (const implementation of supergraphItf.possibleRuntimeTypes()) {
        builder.addEdge(vertex, vertex, new InterfaceObjectFakeDownCast(type, implementation.name), condition);
      }
    }
  }

  return builder.build(FEDERATED_GRAPH_ROOT_SOURCE, forQueryPlanning);
}

function addProvidesEdges(schema: Schema, builder: GraphBuilder, from: Vertex, provided: SelectionSet, provideId: number) {
  const stack: [Vertex, SelectionSet][] = [[from, provided]];
  const source = from.source;
  while (stack.length > 0) {
    const [v, selectionSet] = stack.pop()!;
    // We reverse the selections to cancel the reversing that the stack does.
    for (const selection of selectionSet.selectionsInReverseOrder()) {
      const element = selection.element;
      if (element.kind == 'Field') {
        const fieldDef = element.definition;
        const existingEdge = builder.edg