Fix: Improve merge passthrough child optimization (#272)

The optimization now checks if the child's input contains the output
instead of being identical.
The optimizer now runs in a loop to ensure that all possible
optimizations are applied.

Closes #278

Author: kamilkisiela

Cargo.lock

@@ -16,6 +16,7 @@ serde_json = { version = "1.0.140", features = ["preserve_order"] }
 tracing = { version = "0.1.41" }
 # Data Structures
 petgraph = "0.8.0"
+bitflags = "2.9.1"
 # Utils
 lazy-init = "0.5.1"
 thiserror = "2.0.12"

lib/query-planner/Cargo.toml

@@ -4,11 +4,12 @@ use std::hash::{Hash, Hasher};
 use crate::ast::hash::ASTHash;
 use crate::ast::selection_set::SelectionSet;
 
-pub type SelectionId = u64;
+#[derive(Hash, PartialEq, Eq, Copy, Clone)]
+pub struct SelectionId(u64);
 
 pub fn generate_selection_id(type_name: &str, selection_set: &SelectionSet) -> SelectionId {
     let mut hasher = DefaultHasher::new();
     type_name.hash(&mut hasher);
     selection_set.ast_hash::<_, true>(&mut hasher);
-    hasher.finish()
+    SelectionId(hasher.finish())
 }

lib/query-planner/src/ast/minification/selection_id.rs

@@ -1,13 +1,13 @@
 use std::collections::HashMap;
 
 use crate::ast::minification::error::MinificationError;
-use crate::ast::minification::selection_id::generate_selection_id;
+use crate::ast::minification::selection_id::{generate_selection_id, SelectionId};
 use crate::ast::selection_item::SelectionItem;
 use crate::ast::selection_set::SelectionSet;
 use crate::state::supergraph_state::SupergraphState;
 
 pub struct Stats {
-    state: HashMap<u64, usize>,
+    state: HashMap<SelectionId, usize>,
     contains_duplicates: bool,
 }
 
@@ -30,11 +30,11 @@ impl Stats {
         self.contains_duplicates
     }
 
-    pub fn is_duplicated(&self, key: &u64) -> bool {
+    pub fn is_duplicated(&self, key: &SelectionId) -> bool {
         self.state.get(key).unwrap_or(&0) > &1
     }
 
-    pub fn increase(&mut self, key: u64) -> &usize {
+    pub fn increase(&mut self, key: SelectionId) -> &usize {
         let occurrences = self.state.entry(key).or_insert_with(|| 0);
         *occurrences += 1;
 

lib/query-planner/src/ast/minification/stats.rs

@@ -4,14 +4,14 @@ use std::mem::take;
 use crate::ast::document::Document;
 use crate::ast::fragment::FragmentDefinition;
 use crate::ast::minification::error::MinificationError;
-use crate::ast::minification::selection_id::generate_selection_id;
+use crate::ast::minification::selection_id::{generate_selection_id, SelectionId};
 use crate::ast::minification::stats::Stats;
 use crate::ast::operation::OperationDefinition;
 use crate::ast::selection_item::SelectionItem;
 use crate::ast::selection_set::{FieldSelection, InlineFragmentSelection, SelectionSet};
 use crate::state::supergraph_state::SupergraphState;
 
-type Fragments = HashMap<u64, FragmentDefinition>;
+type Fragments = HashMap<SelectionId, FragmentDefinition>;
 
 pub fn transform_operation(
     supergraph: &SupergraphState,
@@ -253,7 +253,7 @@ pub fn get_or_create_fragment(
     stats: &Stats,
     fragments: &mut Fragments,
     next_fragment_name_idx: &mut usize,
-    id: &u64,
+    id: &SelectionId,
     type_name: &str,
     selection_set: &mut SelectionSet,
 ) -> Result<String, MinificationError> {

lib/query-planner/src/ast/minification/transform.rs

@@ -68,15 +68,34 @@ fn build_possible_types_map<'a>(ctx: &NormalizationContext<'a>) -> PossibleTypes
     let mut possible_types = PossibleTypesMap::new();
     let maybe_subgraph_name = ctx.subgraph_name.as_ref();
 
-    let relevant_definitions = ctx.supergraph.definitions.iter().filter(|(_, def)| {
-        if let Some(subgraph_name) = maybe_subgraph_name {
-            def.is_defined_in_subgraph(subgraph_name.as_str())
-        } else {
-            true
+    let mut object_types_list = Vec::new();
+    let mut abstract_types_list = Vec::new();
+
+    for (name, def) in &ctx.supergraph.definitions {
+        match def {
+            SupergraphDefinition::Union(_) | SupergraphDefinition::Interface(_) => {
+                if maybe_subgraph_name.is_none()
+                    || maybe_subgraph_name.is_some_and(|subgraph_name| {
+                        def.is_defined_in_subgraph(subgraph_name.as_str())
+                    })
+                {
+                    abstract_types_list.push((name, def));
+                }
+            }
+            SupergraphDefinition::Object(_) => {
+                if maybe_subgraph_name.is_none()
+                    || maybe_subgraph_name.is_some_and(|subgraph_name| {
+                        def.is_defined_in_subgraph(subgraph_name.as_str())
+                    })
+                {
+                    object_types_list.push((name, def));
+                }
+            }
+            _ => {}
         }
-    });
+    }
 
-    for (type_name, type_def) in relevant_definitions.clone() {
+    for (type_name, type_def) in &abstract_types_list {
         match type_def {
             SupergraphDefinition::Union(union_type) => {
                 let members = union_type
@@ -95,14 +114,14 @@ fn build_possible_types_map<'a>(ctx: &NormalizationContext<'a>) -> PossibleTypes
             }
             SupergraphDefinition::Interface(_) => {
                 let mut object_types: HashSet<&str> = HashSet::new();
-                for (obj_type_name, obj_type_def) in relevant_definitions.clone() {
+                for (obj_type_name, obj_type_def) in &object_types_list {
                     if let SupergraphDefinition::Object(object_type) = obj_type_def {
                         if object_type.join_implements.iter().any(|j| {
                             let belongs = match maybe_subgraph_name {
                                 Some(subgraph_name) => &j.graph_id == *subgraph_name,
                                 None => true,
                             };
-                            belongs && &j.interface == type_name
+                            belongs && &j.interface == *type_name
                         }) {
                             object_types.insert(obj_type_name.as_str());
                         }

lib/query-planner/src/ast/normalization/pipeline/flatten_fragments.rs

@@ -5,7 +5,7 @@ use crate::ast::type_aware_selection::TypeAwareSelection;
 use crate::graph::edge::{Edge, FieldMove, InterfaceObjectTypeMove, PlannerOverrideContext};
 use crate::graph::node::Node;
 use crate::graph::Graph;
-use crate::planner::fetch::fetch_step_data::{FetchStepData, FetchStepKind};
+use crate::planner::fetch::fetch_step_data::{FetchStepData, FetchStepFlags, FetchStepKind};
 use crate::planner::plan_nodes::{FetchNodePathSegment, FetchRewrite, ValueSetter};
 use crate::planner::tree::query_tree::QueryTree;
 use crate::planner::tree::query_tree_node::{MutationFieldPosition, QueryTreeNode};
@@ -236,6 +236,11 @@ impl Display for FetchGraph {
 }
 
 fn create_noop_fetch_step(fetch_graph: &mut FetchGraph, created_from_requires: bool) -> NodeIndex {
+    let flags = if created_from_requires {
+        FetchStepFlags::USED_FOR_REQUIRES
+    } else {
+        FetchStepFlags::empty()
+    };
     fetch_graph.add_step(FetchStepData {
         service_name: SubgraphName::any(),
         response_path: MergePath::default(),
@@ -247,7 +252,7 @@ fn create_noop_fetch_step(fetch_graph: &mut FetchGraph, created_from_requires: b
             selection_set: SelectionSet::default(),
             type_name: "*".to_string(),
         },
-        used_for_requires: created_from_requires,
+        flags,
         condition: None,
         kind: FetchStepKind::Root,
         input_rewrites: None,
@@ -268,6 +273,11 @@ fn create_fetch_step_for_entity_call(
     condition: Option<&Condition>,
     used_for_requires: bool,
 ) -> NodeIndex {
+    let flags = if used_for_requires {
+        FetchStepFlags::USED_FOR_REQUIRES
+    } else {
+        FetchStepFlags::empty()
+    };
     fetch_graph.add_step(FetchStepData {
         service_name: subgraph_name.clone(),
         response_path: response_path.clone(),
@@ -281,7 +291,7 @@ fn create_fetch_step_for_entity_call(
             selection_set: SelectionSet::default(),
             type_name: output_type_name.to_string(),
         },
-        used_for_requires,
+        flags,
         condition: condition.cloned(),
         kind: FetchStepKind::Entity,
         input_rewrites: None,
@@ -311,7 +321,7 @@ fn create_fetch_step_for_root_move(
             selection_set: SelectionSet::default(),
             type_name: type_name.to_string(),
         },
-        used_for_requires: false,
+        flags: FetchStepFlags::empty(),
         condition: None,
         kind: FetchStepKind::Root,
         variable_usages: None,
@@ -373,7 +383,9 @@ fn ensure_fetch_step_for_subgraph(
 
                     // If there are requirements, then we do not re-use
                     // optimizations will try to re-use the existing step later, if possible.
-                    if fetch_step.used_for_requires || requires.is_some() {
+                    if fetch_step.flags.contains(FetchStepFlags::USED_FOR_REQUIRES)
+                        || requires.is_some()
+                    {
                         return None;
                     }
 
@@ -713,13 +725,18 @@ fn process_entity_move_edge(
             output_type_name,
             fetch_step_index.index()
         );
+
         fetch_step.add_input_rewrite(FetchRewrite::ValueSetter(ValueSetter {
             path: vec![
                 FetchNodePathSegment::TypenameEquals(output_type_name.to_string()),
                 FetchNodePathSegment::Key("__typename".to_string()),
             ],
             set_value_to: output_type_name.clone().into(),
         }));
+
+        fetch_step
+            .flags
+            .insert(FetchStepFlags::USED_FOR_TYPE_CONDITION);
     }
 
     let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;

lib/query-planner/src/planner/fetch/fetch_graph.rs

@@ -1,5 +1,6 @@
 use std::{collections::BTreeSet, fmt::Display};
 
+use bitflags::bitflags;
 use petgraph::{graph::NodeIndex, visit::EdgeRef};
 use tracing::trace;
 
@@ -17,14 +18,24 @@ use crate::{
     state::supergraph_state::SubgraphName,
 };
 
+bitflags! {
+    #[derive(Default, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+    pub struct FetchStepFlags: u8 {
+        /// This fetch is for resolving a @requires directive.
+        const USED_FOR_REQUIRES = 1 << 0;
+        /// This fetch is for resolving a type condition on an interface.
+        const USED_FOR_TYPE_CONDITION = 1 << 1;
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct FetchStepData {
     pub service_name: SubgraphName,
     pub response_path: MergePath,
     pub input: TypeAwareSelection,
     pub output: TypeAwareSelection,
     pub kind: FetchStepKind,
-    pub used_for_requires: bool,
+    pub flags: FetchStepFlags,
     pub condition: Option<Condition>,
     pub variable_usages: Option<BTreeSet<String>>,
     pub variable_definitions: Option<Vec<VariableDefinition>>,
@@ -52,10 +63,14 @@ impl Display for FetchStepData {
             self.response_path.join("."),
         )?;
 
-        if self.used_for_requires {
+        if self.flags.contains(FetchStepFlags::USED_FOR_REQUIRES) {
             write!(f, " [@requires]")?;
         }
 
+        if self.flags.contains(FetchStepFlags::USED_FOR_TYPE_CONDITION) {
+            write!(f, " [no_pass_through]")?;
+        }
+
         if let Some(condition) = &self.condition {
             match condition {
                 Condition::Include(var_name) => write!(f, " [@include(if: ${})]", var_name)?,

lib/query-planner/src/planner/fetch/fetch_step_data.rs

@@ -8,11 +8,13 @@ use petgraph::{
 use tracing::{instrument, trace};
 
 use crate::planner::fetch::{
-    error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData,
+    error::FetchGraphError,
+    fetch_graph::FetchGraph,
+    fetch_step_data::{FetchStepData, FetchStepFlags},
 };
 
 impl FetchGraph {
-    /// When a child has the input identical as the output,
+    /// When a child's input contains its output,
     /// it gets squashed into its parent.
     /// Its children becomes children of the parent.
     #[instrument(level = "trace", skip_all)]
@@ -94,6 +96,13 @@ impl FetchStepData {
             return false;
         }
 
+        if other
+            .flags
+            .contains(FetchStepFlags::USED_FOR_TYPE_CONDITION)
+        {
+            return false;
+        }
+
         // if the `other` FetchStep has a single parent and it's `this` FetchStep
         if fetch_graph.parents_of(other_index).count() != 1 {
             return false;
@@ -103,7 +112,7 @@ impl FetchStepData {
             return false;
         }
 
-        other.input.eq(&other.output)
+        other.input.contains(&other.output)
     }
 }
 

lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs

@@ -18,11 +18,25 @@ use crate::{
 impl FetchGraph {
     #[instrument(level = "trace", skip_all)]
     pub fn optimize(&mut self, supergraph_state: &SupergraphState) -> Result<(), FetchGraphError> {
-        self.merge_passthrough_child()?;
-        self.merge_children_with_parents()?;
-        self.merge_siblings()?;
-        self.merge_leafs()?;
-        self.deduplicate_and_prune_fetch_steps()?;
+        // Run optimization passes repeatedly until the graph stabilizes, as one optimization can create
+        // opportunities for others.
+        loop {
+            let node_count_before = self.graph.node_count();
+            let edge_count_before = self.graph.edge_count();
+
+            self.merge_passthrough_child()?;
+            self.merge_children_with_parents()?;
+            self.merge_siblings()?;
+            self.merge_leafs()?;
+            self.deduplicate_and_prune_fetch_steps()?;
+
+            let node_count_after = self.graph.node_count();
+            let edge_count_after = self.graph.edge_count();
+
+            if node_count_before == node_count_after && edge_count_before == edge_count_after {
+                break;
+            }
+        }
         self.turn_mutations_into_sequence()?;
         self.fix_conflicting_type_mismatches(supergraph_state)?;