Add search state and field routing skeleton for incremental planner

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: tninesling

apollo-federation/src/query_plan/incremental_planner/field_routing/conditions.rs

@@ -0,0 +1,117 @@
+//! Condition satisfiability: can a set of @requires / @key fields be resolved
+//! at a given query graph node?
+
+use std::sync::Arc;
+
+use petgraph::graph::NodeIndex;
+
+use super::FieldRoutingSearchSpace;
+use crate::error::FederationError;
+use crate::operation::SelectionSet;
+use crate::schema::ValidFederationSchema;
+use crate::schema::position::CompositeTypeDefinitionPosition;
+
+impl FieldRoutingSearchSpace {
+    /// Schema-based check: can every field in `conditions` be found in the
+    /// subgraph's schema at `type_pos`?
+    pub(super) fn can_satisfy(
+        &self,
+        conditions: &Arc<SelectionSet>,
+        type_pos: &CompositeTypeDefinitionPosition,
+        _subgraph: &Arc<str>,
+        schema: &ValidFederationSchema,
+    ) -> bool {
+        can_satisfy_conditions(conditions, type_pos, schema)
+    }
+
+    /// Graph-based check: can every field in `conditions` be resolved at
+    /// `node` via outgoing edges?
+    pub(super) fn conditions_resolvable_at_node(
+        &self,
+        node: NodeIndex,
+        conditions: &Arc<SelectionSet>,
+    ) -> Result<bool, FederationError> {
+        for selection in conditions.selections.values() {
+            match selection {
+                crate::operation::Selection::Field(field_sel) => {
+                    if self.edge_for_field(node, &field_sel.field).is_none() {
+                        return Ok(false);
+                    }
+                }
+                crate::operation::Selection::InlineFragment(frag_sel) => {
+                    if self
+                        .edge_for_inline_fragment(node, &frag_sel.inline_fragment)
+                        .is_none()
+                    {
+                        return Ok(false);
+                    }
+                }
+            }
+        }
+        Ok(true)
+    }
+
+    /// Do any fields in `conditions` carry @requires at `node`? If so,
+    /// the conditions cannot be resolved in-place and need their own
+    /// entity fetch.
+    pub(super) fn conditions_have_requires(
+        &self,
+        node: NodeIndex,
+        conditions: &Arc<SelectionSet>,
+    ) -> Result<bool, FederationError> {
+        for selection in conditions.selections.values() {
+            if let crate::operation::Selection::Field(field_sel) = selection
+                && let Some(edge_idx) = self.edge_for_field(node, &field_sel.field)
+            {
+                let edge = self.query_graph.edge_weight(edge_idx)?;
+                if edge.conditions.is_some() {
+                    return Ok(true);
+                }
+            }
+        }
+        Ok(false)
+    }
+}
+
+/// Check whether every field in `conditions` exists in the subgraph schema
+/// at the given type position.
+fn can_satisfy_conditions(
+    conditions: &SelectionSet,
+    type_pos: &CompositeTypeDefinitionPosition,
+    schema: &ValidFederationSchema,
+) -> bool {
+    for selection in conditions.selections.values() {
+        match selection {
+            crate::operation::Selection::Field(field_sel) => {
+                let field_name = field_sel.field.name();
+                if field_name.as_str() == "__typename" {
+                    continue;
+                }
+                let field_exists = type_pos
+                    .field(field_name.clone())
+                    .is_ok_and(|field_pos| field_pos.get(schema.schema()).is_ok());
+                if !field_exists {
+                    return false;
+                }
+            }
+            crate::operation::Selection::InlineFragment(frag_sel) => {
+                if let Some(type_cond) = &frag_sel.inline_fragment.type_condition_position {
+                    let inner_type_name = type_cond.type_name();
+                    let inner_pos = schema
+                        .get_type(inner_type_name)
+                        .ok()
+                        .and_then(|ty| CompositeTypeDefinitionPosition::try_from(ty).ok());
+                    match inner_pos {
+                        Some(pos) => {
+                            if !can_satisfy_conditions(&frag_sel.selection_set, &pos, schema) {
+                                return false;
+                            }
+                        }
+                        None => return false,
+                    }
+                }
+            }
+        }
+    }
+    true
+}

apollo-federation/src/query_plan/incremental_planner/field_routing/mod.rs

@@ -0,0 +1,168 @@
+//! Field-level routing as a BULB search space.
+//!
+//! The planner walks the operation selection-by-selection ("pendings"),
+//! consulting the query graph for where each field can be resolved:
+//! - [`state`]: mutable search state (pending stack, checkpoints).
+//! - [`routing`]: enumerating and ranking options for a selection.
+//! - [`commit`]: applying a chosen option to the fetch graph.
+//! - [`conditions`]: condition satisfiability for @requires / @key.
+//! - [`requires`]: hop-edge inputs and condition paths.
+//!
+//! This file holds the search-space type and the
+//! [`BulbSearchSpace`] implementation.
+
+mod conditions;
+mod requires;
+pub(super) mod state;
+
+use std::sync::Arc;
+
+use apollo_compiler::Name;
+use petgraph::graph::EdgeIndex;
+use petgraph::graph::NodeIndex;
+use petgraph::visit::EdgeRef;
+#[allow(unused_imports)]
+pub(crate) use state::PendingSelection;
+#[allow(unused_imports)]
+pub(crate) use state::PlanState;
+
+use super::shared_path::SharedPath;
+use crate::error::FederationError;
+use crate::operation::Field;
+use crate::operation::InlineFragment;
+use crate::operation::Selection;
+use crate::operation::SelectionId;
+use crate::operation::SelectionSet;
+use crate::query_graph::OverrideConditions;
+use crate::query_graph::QueryGraph;
+use crate::query_graph::graph_path::operation::OpPathElement;
+use crate::schema::ValidFederationSchema;
+use crate::schema::position::CompositeTypeDefinitionPosition;
+
+/// Cache key for routing options. Captures the selection identity at a QG node.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub(super) enum RoutingCacheKey {
+    Field(Name),
+    InlineFragment(Option<Name>),
+}
+
+/// Subgraph, type position, and schema at a query graph node.
+pub(super) struct NodeSource {
+    pub(super) subgraph: Arc<str>,
+    pub(super) type_pos: CompositeTypeDefinitionPosition,
+    pub(super) schema: ValidFederationSchema,
+}
+
+/// Search space presenting field-level routing decisions as a BULB problem.
+pub(crate) struct FieldRoutingSearchSpace {
+    pub(crate) query_graph: Arc<QueryGraph>,
+    pub(crate) supergraph_schema: ValidFederationSchema,
+    pub(crate) override_conditions: OverrideConditions,
+    pub(crate) inconsistent_abstract_types: Arc<apollo_compiler::collections::IndexSet<Name>>,
+}
+
+impl FieldRoutingSearchSpace {
+    pub(super) fn node_source(&self, node: NodeIndex) -> Result<NodeSource, FederationError> {
+        let data = self.query_graph.node_weight(node)?;
+        Ok(NodeSource {
+            subgraph: data.source.clone(),
+            type_pos: data.type_.clone().try_into()?,
+            schema: self.query_graph.schema_by_source(&data.source)?.clone(),
+        })
+    }
+
+    /// Select `__typename` in `fetch_node` at `base_path` so the executor
+    /// can identify the concrete type for entity representations.
+    pub(super) fn append_typename(
+        &self,
+        state: &mut PlanState,
+        fetch_node: NodeIndex,
+        base_path: &SharedPath<Arc<OpPathElement>>,
+        source: &NodeSource,
+    ) {
+        let typename = Arc::new(OpPathElement::Field(Field::new_introspection_typename(
+            &source.schema,
+            &source.type_pos,
+            None,
+        )));
+        state
+            .graph
+            .append_selection(fetch_node, &base_path.pushed(typename), None);
+    }
+
+    /// Op path at which selections enter an entity fetch group: entity
+    /// fetches start from the `_Entity` union, so everything nests under a
+    /// `... on <ConcreteType>` condition rebased onto the supergraph schema
+    /// (which OpPaths reference).
+    pub(super) fn entity_root_path(
+        &self,
+        type_name: &Name,
+    ) -> Result<SharedPath<Arc<OpPathElement>>, FederationError> {
+        let rebased: CompositeTypeDefinitionPosition =
+            self.supergraph_schema.get_type(type_name)?.try_into()?;
+        let condition = InlineFragment {
+            schema: self.supergraph_schema.clone(),
+            parent_type_position: rebased.clone(),
+            type_condition_position: Some(rebased),
+            directives: Default::default(),
+            selection_id: SelectionId::new(),
+        };
+        Ok(SharedPath::new().pushed(Arc::new(OpPathElement::InlineFragment(condition))))
+    }
+
+    /// Can condition fields simply be selected in the fetch at `node`?
+    /// True when the subgraph resolves every field itself and none carries
+    /// @requires (which draws on an entity representation and needs its own
+    /// fetch). The graph-based check complements the schema-based one:
+    /// @external fields may still resolve at `node` when it is a
+    /// provides-copy created by an ancestor's @provides.
+    pub(super) fn can_resolve_in_place(
+        &self,
+        node: NodeIndex,
+        conditions: &Arc<SelectionSet>,
+        source: &NodeSource,
+    ) -> Result<bool, FederationError> {
+        let satisfiable = self.can_satisfy(
+            conditions,
+            &source.type_pos,
+            &source.subgraph,
+            &source.schema,
+        ) || self.conditions_resolvable_at_node(node, conditions)?;
+        Ok(satisfiable && !self.conditions_have_requires(node, conditions)?)
+    }
+
+    /// Outgoing edge indices from a query graph node, sorted and filtered.
+    pub(super) fn out_edge_indices(&self, node: NodeIndex) -> Vec<EdgeIndex> {
+        self.query_graph
+            .out_edges(node)
+            .into_iter()
+            .map(|edge_ref| edge_ref.id())
+            .collect()
+    }
+
+    /// Find the outgoing edge for a field at a query graph node.
+    pub(super) fn edge_for_field(&self, node: NodeIndex, field: &Field) -> Option<EdgeIndex> {
+        self.query_graph
+            .edge_for_field(node, field, &self.override_conditions)
+    }
+
+    /// Find the outgoing downcast edge for an inline fragment at a query
+    /// graph node.
+    pub(super) fn edge_for_inline_fragment(
+        &self,
+        node: NodeIndex,
+        fragment: &InlineFragment,
+    ) -> Option<EdgeIndex> {
+        self.query_graph.edge_for_inline_fragment(node, fragment)
+    }
+}
+
+/// Short human-readable label for a selection, for logging.
+pub(super) fn selection_label(selection: &Selection) -> String {
+    match selection {
+        Selection::Field(f) => f.field.field_position.to_string(),
+        Selection::InlineFragment(f) => {
+            format!("... on {:?}", f.inline_fragment.type_condition_position)
+        }
+    }
+}

apollo-federation/src/query_plan/incremental_planner/field_routing/requires.rs

@@ -0,0 +1,62 @@
+//! Key-hop input path helpers: building the unconditioned input path and
+//! extracting trailing condition fragments for entity fetch op paths.
+
+use std::sync::Arc;
+
+use super::super::shared_path::SharedPath;
+use crate::operation::DirectiveList;
+use crate::query_graph::graph_path::operation::OpPathElement;
+
+/// The trailing inline-fragment elements of `op_path` (after the last field)
+/// that carry directives: @skip/@include conditions at the current position,
+/// which a key hop must carry into the entity fetch's op path or the hopped
+/// selections lose their gating.
+pub(super) fn trailing_condition_fragments(
+    op_path: &SharedPath<Arc<OpPathElement>>,
+) -> Vec<Arc<OpPathElement>> {
+    let mut trailing = Vec::new();
+    for element in op_path.iter() {
+        match element.as_ref() {
+            OpPathElement::Field(_) => trailing.clear(),
+            OpPathElement::InlineFragment(frag) => {
+                if !frag.directives.is_empty() {
+                    trailing.push(element.clone());
+                }
+            }
+        }
+    }
+    trailing
+}
+
+/// `op_path` with @skip/@include stripped from its inline-fragment elements,
+/// for appending key input selections. Inputs must be selected
+/// unconditionally: an input gated by one branch's Boolean condition leaves
+/// the representation incomplete whenever a different branch executes.
+/// Condition-only fragments are dropped; type-conditioned fragments keep the
+/// downcast without the conditions.
+pub(super) fn unconditioned_input_path(
+    op_path: &SharedPath<Arc<OpPathElement>>,
+) -> SharedPath<Arc<OpPathElement>> {
+    let mut elements = Vec::with_capacity(op_path.len());
+    for element in op_path.iter() {
+        match element.as_ref() {
+            OpPathElement::Field(_) => elements.push(element.clone()),
+            OpPathElement::InlineFragment(frag) => {
+                let stripped: DirectiveList = frag
+                    .directives
+                    .iter()
+                    .filter(|d| d.name != "skip" && d.name != "include")
+                    .cloned()
+                    .collect();
+                if stripped.len() == frag.directives.len() {
+                    elements.push(element.clone());
+                } else if !stripped.is_empty() || frag.type_condition_position.is_some() {
+                    elements.push(Arc::new(OpPathElement::InlineFragment(
+                        frag.with_updated_directives(stripped),
+                    )));
+                }
+            }
+        }
+    }
+    SharedPath::from_vec(elements)
+}