Use OpTree concatenation in serial query planning (#5636)

apollo-federation/src/query_graph/path_tree.rs

@@ -23,7 +23,6 @@ use crate::query_graph::QueryGraphNode;
 // Typescript doesn't have a native way of associating equality/hash functions with types, so they
 // were passed around manually. This isn't the case with Rust, where we instead implement trigger
 // equality via `PartialEq` and `Hash`.
-#[derive(Clone)]
 pub(crate) struct PathTree<TTrigger, TEdge>
 where
     TTrigger: Eq + Hash,
@@ -45,6 +44,34 @@ where
     pub(crate) childs: Vec<Arc<PathTreeChild<TTrigger, TEdge>>>,
 }
 
+impl<TTrigger, TEdge> Clone for PathTree<TTrigger, TEdge>
+where
+    TTrigger: Eq + Hash,
+    TEdge: Copy + Into<Option<EdgeIndex>>,
+{
+    fn clone(&self) -> Self {
+        Self {
+            graph: self.graph.clone(),
+            node: self.node,
+            local_selection_sets: self.local_selection_sets.clone(),
+            childs: self.childs.clone(),
+        }
+    }
+}
+
+impl<TTrigger, TEdge> PartialEq for PathTree<TTrigger, TEdge>
+where
+    TTrigger: Eq + Hash,
+    TEdge: Copy + PartialEq + Into<Option<EdgeIndex>>,
+{
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.graph, &other.graph)
+            && self.node == other.node
+            && self.local_selection_sets == other.local_selection_sets
+            && self.childs == other.childs
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct PathTreeChild<TTrigger, TEdge>
 where
@@ -61,6 +88,19 @@ where
     pub(crate) tree: Arc<PathTree<TTrigger, TEdge>>,
 }
 
+impl<TTrigger, TEdge> PartialEq for PathTreeChild<TTrigger, TEdge>
+where
+    TTrigger: Eq + Hash,
+    TEdge: Copy + PartialEq + Into<Option<EdgeIndex>>,
+{
+    fn eq(&self, other: &Self) -> bool {
+        self.edge == other.edge
+            && self.trigger == other.trigger
+            && self.conditions == other.conditions
+            && self.tree == other.tree
+    }
+}
+
 /// A `PathTree` whose triggers are operation elements (essentially meaning that the constituent
 /// `GraphPath`s were guided by a GraphQL operation).
 pub(crate) type OpPathTree = PathTree<OpGraphPathTrigger, Option<EdgeIndex>>;
@@ -283,6 +323,39 @@ where
             })
     }
 
+    /// Appends the children of the other `OpTree` onto the children of this tree.
+    ///
+    /// ## Panics
+    /// Like `Self::merge`, this method will panic if the graphs of the two `OpTree`s below to
+    /// different allocations (i.e. they don't below to the same graph) or if they below to
+    /// different root nodes.
+    pub(crate) fn extend(&mut self, other: &Self) {
+        assert!(
+            Arc::ptr_eq(&self.graph, &other.graph),
+            "Cannot merge path tree build on another graph"
+        );
+        assert_eq!(
+            self.node, other.node,
+            "Cannot merge path trees rooted different nodes"
+        );
+        if self == other {
+            return;
+        }
+        if other.childs.is_empty() {
+            return;
+        }
+        if self.childs.is_empty() {
+            self.clone_from(other);
+            return;
+        }
+        self.childs.extend_from_slice(&other.childs);
+        self.local_selection_sets
+            .extend_from_slice(&other.local_selection_sets);
+    }
+
+    /// ## Panics
+    /// This method will panic if the graphs of the two `OpTree`s below to different allocations
+    /// (i.e. they don't below to the same graph) or if they below to different root nodes.
     pub(crate) fn merge(self: &Arc<Self>, other: &Arc<Self>) -> Arc<Self> {
         if Arc::ptr_eq(self, other) {
             return self.clone();

apollo-federation/src/query_plan/query_planner.rs

@@ -8,8 +8,6 @@ use apollo_compiler::Name;
 use indexmap::IndexMap;
 use indexmap::IndexSet;
 use itertools::Itertools;
-use petgraph::csr::NodeIndex;
-use petgraph::stable_graph::IndexType;
 
 use crate::error::FederationError;
 use crate::error::SingleFederationError;
@@ -560,7 +558,7 @@ fn compute_root_serial_dependency_graph(
             // }
             // then we should _not_ merge the 2 `mut1` fields (contrarily to what happens on queried fields).
 
-            prev_path = OpPathTree::merge(&prev_path, &new_path);
+            Arc::make_mut(&mut prev_path).extend(&new_path);
             fetch_dependency_graph = FetchDependencyGraph::new(
                 supergraph_schema.clone(),
                 federated_query_graph.clone(),
@@ -590,14 +588,14 @@ fn compute_root_serial_dependency_graph(
     Ok(digest)
 }
 
-fn only_root_subgraph(graph: &FetchDependencyGraph) -> Result<NodeIndex, FederationError> {
+fn only_root_subgraph(graph: &FetchDependencyGraph) -> Result<Arc<str>, FederationError> {
     let mut iter = graph.root_node_by_subgraph_iter();
-    let (Some((_, index)), None) = (iter.next(), iter.next()) else {
+    let (Some((name, _)), None) = (iter.next(), iter.next()) else {
         return Err(FederationError::internal(format!(
             "{graph} should have only one root."
         )));
     };
-    Ok(index.index() as u32)
+    Ok(name.clone())
 }
 
 pub(crate) fn compute_root_fetch_groups(

apollo-federation/tests/query_plan/build_query_plan_tests.rs

@@ -41,6 +41,7 @@ mod interface_object;
 mod interface_type_explosion;
 mod introspection_typename_handling;
 mod merged_abstract_types_handling;
+mod mutations;
 mod named_fragments;
 mod named_fragments_preservation;
 mod provides;

apollo-federation/tests/query_plan/build_query_plan_tests/mutations.rs

@@ -0,0 +1,135 @@
+const SUBGRAPH_A: &str = r#"
+      type Foo @key(fields: "id") {
+        id: ID!
+        bar: String
+      }
+
+      type Query {
+        foo: Foo
+      }
+
+      type Mutation {
+        updateFooInA: Foo
+      }
+"#;
+
+const SUBGRAPH_B: &str = r#"
+      type Mutation {
+        updateFooInB: Foo
+      }
+
+      type Foo @key(fields: "id") {
+        id: ID!
+        baz: Int
+      }
+"#;
+
+#[test]
+fn adjacent_mutations_get_merged() {
+    let planner = planner!(
+        SubgraphA: SUBGRAPH_A,
+        SubgraphB: SUBGRAPH_B,
+    );
+    assert_plan!(
+        &planner,
+        r#"
+        mutation TestMutation {
+          updateInAOne: updateFooInA {
+            id
+            bar
+          }
+          updateInATwo: updateFooInA {
+            id
+            bar
+          }
+          updateInBOne: updateFooInB {
+            id
+            baz
+          }
+        }
+        "#,
+        @r###"
+        QueryPlan {
+          Sequence {
+            Fetch(service: "SubgraphA") {
+              {
+                updateInAOne: updateFooInA {
+                  id
+                  bar
+                }
+                updateInATwo: updateFooInA {
+                  id
+                  bar
+                }
+              }
+            },
+            Fetch(service: "SubgraphB") {
+              {
+                updateInBOne: updateFooInB {
+                  id
+                  baz
+                }
+              }
+            },
+          },
+        }
+        "###
+    );
+}
+
+#[test]
+fn non_adjacent_mutations_do_not_get_merged() {
+    let planner = planner!(
+        SubgraphA: SUBGRAPH_A,
+        SubgraphB: SUBGRAPH_B,
+    );
+    assert_plan!(
+        &planner,
+        r#"
+        mutation TestMutation {
+          updateInAOne: updateFooInA {
+            id
+            bar
+          }
+          updateInBOne: updateFooInB {
+            id
+            baz
+          }
+          updateInATwo: updateFooInA {
+            id
+            bar
+          }
+        }
+        "#,
+        @r###"
+        QueryPlan {
+          Sequence {
+            Fetch(service: "SubgraphA") {
+              {
+                updateInAOne: updateFooInA {
+                  id
+                  bar
+                }
+              }
+            },
+            Fetch(service: "SubgraphB") {
+              {
+                updateInBOne: updateFooInB {
+                  id
+                  baz
+                }
+              }
+            },
+            Fetch(service: "SubgraphA") {
+              {
+                updateInATwo: updateFooInA {
+                  id
+                  bar
+                }
+              }
+            },
+          },
+        }
+        "###
+    );
+}

apollo-federation/tests/query_plan/supergraphs/adjacent_mutations_get_merged.graphql

@@ -0,0 +1,67 @@
+# Composed from subgraphs with hash: 2655e7da6754e73955fece01d7cbb5f21085bdbb
+schema
+  @link(url: "https://specs.apollo.dev/link/v1.0")
+  @link(url: "https://specs.apollo.dev/join/v0.3", for: EXECUTION)
+{
+  query: Query
+  mutation: Mutation
+}
+
+directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
+
+directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
+
+directive @join__graph(name: String!, url: String!) on ENUM_VALUE
+
+directive @join__implements(graph: join__Graph!, interface: String!) repeatable on OBJECT | INTERFACE
+
+directive @join__type(graph: join__Graph!, key: join__FieldSet, extension: Boolean! = false, resolvable: Boolean! = true, isInterfaceObject: Boolean! = false) repeatable on OBJECT | INTERFACE | UNION | ENUM | INPUT_OBJECT | SCALAR
+
+directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+
+directive @link(url: String, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
+
+type Foo
+  @join__type(graph: SUBGRAPHA, key: "id")
+  @join__type(graph: SUBGRAPHB, key: "id")
+{
+  id: ID!
+  bar: String @join__field(graph: SUBGRAPHA)
+  baz: Int @join__field(graph: SUBGRAPHB)
+}
+
+scalar join__FieldSet
+
+enum join__Graph {
+  SUBGRAPHA @join__graph(name: "SubgraphA", url: "none")
+  SUBGRAPHB @join__graph(name: "SubgraphB", url: "none")
+}
+
+scalar link__Import
+
+enum link__Purpose {
+  """
+  `SECURITY` features provide metadata necessary to securely resolve fields.
+  """
+  SECURITY
+
+  """
+  `EXECUTION` features provide metadata necessary for operation execution.
+  """
+  EXECUTION
+}
+
+type Mutation
+  @join__type(graph: SUBGRAPHA)
+  @join__type(graph: SUBGRAPHB)
+{
+  updateFooInA: Foo @join__field(graph: SUBGRAPHA)
+  updateFooInB: Foo @join__field(graph: SUBGRAPHB)
+}
+
+type Query
+  @join__type(graph: SUBGRAPHA)
+  @join__type(graph: SUBGRAPHB)
+{
+  foo: Foo @join__field(graph: SUBGRAPHA)
+}