FED-47: Finish porting compute_best_plan_from_closed_branches (apol…

…lographql/federation-next#277)

- refactor: moved `query_plan::display::State` to `crate::indented_display` module since it's quite reusable across the board.

apollo-federation/src/indented_display.rs

@@ -0,0 +1,65 @@
+use std::fmt;
+
+pub(crate) struct State<'fmt, 'fmt2> {
+    indent_level: usize,
+    output: &'fmt mut fmt::Formatter<'fmt2>,
+}
+
+impl<'a, 'b> State<'a, 'b> {
+    pub(crate) fn new(output: &'a mut fmt::Formatter<'b>) -> State<'a, 'b> {
+        Self {
+            indent_level: 0,
+            output,
+        }
+    }
+
+    pub(crate) fn indent_level(&self) -> usize {
+        self.indent_level
+    }
+
+    pub(crate) fn write<T: fmt::Display>(&mut self, value: T) -> fmt::Result {
+        write!(self.output, "{}", value)
+    }
+
+    pub(crate) fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
+        self.output.write_fmt(args)
+    }
+
+    pub(crate) fn new_line(&mut self) -> fmt::Result {
+        self.write("\n")?;
+        for _ in 0..self.indent_level {
+            self.write("  ")?
+        }
+        Ok(())
+    }
+
+    pub(crate) fn indent_no_new_line(&mut self) {
+        self.indent_level += 1;
+    }
+
+    pub(crate) fn indent(&mut self) -> fmt::Result {
+        self.indent_no_new_line();
+        self.new_line()
+    }
+
+    pub(crate) fn dedent(&mut self) -> fmt::Result {
+        self.indent_level -= 1;
+        self.new_line()
+    }
+}
+
+pub(crate) fn write_indented_lines<T>(
+    state: &mut State<'_, '_>,
+    values: &[T],
+    mut write_line: impl FnMut(&mut State<'_, '_>, &T) -> fmt::Result,
+) -> fmt::Result {
+    if !values.is_empty() {
+        state.indent_no_new_line();
+        for value in values {
+            state.new_line()?;
+            write_line(state, value)?;
+        }
+        state.dedent()?;
+    }
+    Ok(())
+}

apollo-federation/src/lib.rs

@@ -3,6 +3,7 @@
 mod api_schema;
 mod compat;
 pub mod error;
+mod indented_display;
 pub mod link;
 pub mod merge;
 pub mod query_graph;

apollo-federation/src/query_graph/graph_path.rs

@@ -1,4 +1,5 @@
 use crate::error::FederationError;
+use crate::indented_display::{write_indented_lines, State as IndentedFormatter};
 use crate::is_leaf_type;
 use crate::link::federation_spec_definition::get_federation_spec_definition_from_subgraph;
 use crate::link::graphql_definition::{
@@ -477,6 +478,22 @@ impl Display for OpGraphPathContext {
 #[derive(Clone)]
 pub(crate) struct SimultaneousPaths(pub(crate) Vec<Arc<OpGraphPath>>);
 
+impl SimultaneousPaths {
+    pub(crate) fn fmt_indented(&self, f: &mut IndentedFormatter) -> std::fmt::Result {
+        match self.0.as_slice() {
+            [] => f.write("<no path>"),
+
+            [first] => f.write_fmt(format_args!("{{ {first} }}")),
+
+            _ => {
+                f.write("{")?;
+                write_indented_lines(f, &self.0, |f, elem| f.write(elem))?;
+                f.write("}")
+            }
+        }
+    }
+}
+
 impl std::fmt::Debug for SimultaneousPaths {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_list()
@@ -485,6 +502,12 @@ impl std::fmt::Debug for SimultaneousPaths {
     }
 }
 
+impl std::fmt::Display for SimultaneousPaths {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.fmt_indented(&mut IndentedFormatter::new(f))
+    }
+}
+
 /// One of the options for an `OpenBranch` (see the documentation of that struct for details). This
 /// includes the simultaneous paths we are traversing for the option, along with metadata about the
 /// traversal.
@@ -693,11 +716,33 @@ enum UnadvanceableReason {
 /// set, and the `SimultaneousPaths` ends at the node at which that query is made instead of a node
 /// for the leaf field. The selection set gets copied "as-is" into the `FetchNode`, and also avoids
 /// extra `GraphPath` creation and work during `PathTree` merging.
+#[derive(Debug)]
 pub(crate) struct ClosedPath {
     pub(crate) paths: SimultaneousPaths,
     pub(crate) selection_set: Option<Arc<NormalizedSelectionSet>>,
 }
 
+impl ClosedPath {
+    pub(crate) fn flatten(
+        &self,
+    ) -> impl Iterator<Item = (&OpGraphPath, Option<&Arc<NormalizedSelectionSet>>)> {
+        self.paths
+            .0
+            .iter()
+            .map(|path| (path.as_ref(), self.selection_set.as_ref()))
+    }
+}
+
+impl std::fmt::Display for ClosedPath {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        if let Some(ref selection_set) = self.selection_set {
+            write!(f, "{} -> {}", self.paths, selection_set)
+        } else {
+            write!(f, "{}", self.paths)
+        }
+    }
+}
+
 /// A list of the options generated during query planning for a specific "closed branch", which is a
 /// full/closed path in a GraphQL operation (i.e. one that ends in a leaf field).
 pub(crate) struct ClosedBranch(pub(crate) Vec<Arc<ClosedPath>>);

apollo-federation/src/query_graph/path_tree.rs

@@ -19,7 +19,7 @@ use std::sync::Arc;
 // 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(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct PathTree<TTrigger, TEdge>
 where
     TTrigger: Eq + Hash,
@@ -74,7 +74,7 @@ impl OpPathTree {
     pub(crate) fn from_op_paths(
         graph: Arc<QueryGraph>,
         node: NodeIndex,
-        paths: &[(&OpGraphPath, &Arc<NormalizedSelectionSet>)],
+        paths: &[(&OpGraphPath, Option<&Arc<NormalizedSelectionSet>>)],
     ) -> Result<Self, FederationError> {
         assert!(
             !paths.is_empty(),
@@ -166,7 +166,7 @@ where
         node: NodeIndex,
         graph_paths_and_selections: Vec<(
             impl Iterator<Item = GraphPathItem<'inputs, TTrigger, TEdge>>,
-            &'inputs Arc<NormalizedSelectionSet>,
+            Option<&'inputs Arc<NormalizedSelectionSet>>,
         )>,
     ) -> Result<Self, FederationError>
     where
@@ -184,15 +184,18 @@ where
 
         struct PathTreeChildInputs<'inputs, GraphPathIter> {
             conditions: Option<Arc<OpPathTree>>,
-            sub_paths_and_selections: Vec<(GraphPathIter, &'inputs Arc<NormalizedSelectionSet>)>,
+            sub_paths_and_selections:
+                Vec<(GraphPathIter, Option<&'inputs Arc<NormalizedSelectionSet>>)>,
         }
 
         let mut local_selection_sets = Vec::new();
 
         for (mut graph_path_iter, selection) in graph_paths_and_selections {
             let Some((generic_edge, trigger, conditions)) = graph_path_iter.next() else {
                 // End of an input `GraphPath`
-                local_selection_sets.push(selection.clone());
+                if let Some(selection) = selection {
+                    local_selection_sets.push(selection.clone());
+                }
                 continue;
             };
             let for_edge = match merged.entry(generic_edge) {
@@ -277,7 +280,7 @@ where
             })
     }
 
-    fn merge(self: &Arc<Self>, other: &Arc<Self>) -> Arc<Self> {
+    pub(crate) fn merge(self: &Arc<Self>, other: &Arc<Self>) -> Arc<Self> {
         if Arc::ptr_eq(self, other) {
             return self.clone();
         }
@@ -503,7 +506,10 @@ mod tests {
                 .unwrap();
         let selection_set = Arc::new(normalized_operation.selection_set);
 
-        let paths = vec![(&path1, &selection_set), (&path2, &selection_set)];
+        let paths = vec![
+            (&path1, Some(&selection_set)),
+            (&path2, Some(&selection_set)),
+        ];
         let path_tree =
             OpPathTree::from_op_paths(query_graph.to_owned(), NodeIndex::new(0), &paths).unwrap();
         let computed = path_tree.to_string();

apollo-federation/src/query_plan/display.rs

@@ -1,40 +1,8 @@
 use super::*;
+use crate::indented_display::{write_indented_lines, State};
 use apollo_compiler::executable;
 use std::fmt;
 
-pub(crate) struct State<'fmt, 'fmt2> {
-    indent_level: usize,
-    output: &'fmt mut fmt::Formatter<'fmt2>,
-}
-
-impl State<'_, '_> {
-    fn write<T: fmt::Display>(&mut self, value: T) -> fmt::Result {
-        write!(self.output, "{}", value)
-    }
-
-    fn new_line(&mut self) -> fmt::Result {
-        self.write("\n")?;
-        for _ in 0..self.indent_level {
-            self.write("  ")?
-        }
-        Ok(())
-    }
-
-    fn indent_no_new_line(&mut self) {
-        self.indent_level += 1;
-    }
-
-    fn indent(&mut self) -> fmt::Result {
-        self.indent_no_new_line();
-        self.new_line()
-    }
-
-    fn dedent(&mut self) -> fmt::Result {
-        self.indent_level -= 1;
-        self.new_line()
-    }
-}
-
 impl QueryPlan {
     fn write_indented(&self, state: &mut State<'_, '_>) -> fmt::Result {
         let Self {
@@ -264,7 +232,7 @@ impl PrimaryDeferBlock {
                 state.write(
                     sub_selection
                         .serialize()
-                        .initial_indent_level(state.indent_level),
+                        .initial_indent_level(state.indent_level()),
                 )?;
                 if node.is_some() {
                     state.write(":")?;
@@ -345,15 +313,15 @@ fn write_operation(
         state.write(
             operation
                 .serialize()
-                .initial_indent_level(state.indent_level),
+                .initial_indent_level(state.indent_level()),
         )?
     }
     for fragment in operation_document.fragments.values() {
         state.new_line()?;
         state.write(
             fragment
                 .serialize()
-                .initial_indent_level(state.indent_level),
+                .initial_indent_level(state.indent_level()),
         )?
     }
     Ok(())
@@ -370,27 +338,11 @@ fn write_selections(
     }
     state.write("{")?;
     write_indented_lines(state, selections, |state, sel| {
-        state.write(sel.serialize().initial_indent_level(state.indent_level))
+        state.write(sel.serialize().initial_indent_level(state.indent_level()))
     })?;
     state.write("}")
 }
 
-fn write_indented_lines<T>(
-    state: &mut State<'_, '_>,
-    values: &[T],
-    mut write_line: impl FnMut(&mut State<'_, '_>, &T) -> fmt::Result,
-) -> fmt::Result {
-    if !values.is_empty() {
-        state.indent_no_new_line();
-        for value in values {
-            state.new_line()?;
-            write_line(state, value)?;
-        }
-        state.dedent()?;
-    }
-    Ok(())
-}
-
 impl fmt::Display for FetchDataPathElement {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
@@ -421,7 +373,7 @@ macro_rules! impl_display {
         $(
             impl fmt::Display for $Ty {
                 fn fmt(&self, output: &mut fmt::Formatter<'_>) -> fmt::Result {
-                    self.write_indented(&mut State { indent_level: 0, output })
+                    self.write_indented(&mut State::new(output))
                 }
             }
         )+

apollo-federation/src/query_plan/fetch_dependency_graph.rs

@@ -109,6 +109,14 @@ impl FetchIdGenerator {
     }
 }
 
+impl Clone for FetchIdGenerator {
+    fn clone(&self) -> Self {
+        Self {
+            next: AtomicU64::new(self.next.load(std::sync::atomic::Ordering::Relaxed)),
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct FetchSelectionSet {
     /// The selection set to be fetched from the subgraph.
@@ -153,7 +161,7 @@ type FetchDependencyGraphPetgraph =
 ///
 /// In the graph, two fetches are connected if one of them (the parent/head) must be performed
 /// strictly before the other one (the child/tail).
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct FetchDependencyGraph {
     /// The supergraph schema that generated the federated query graph.
     supergraph_schema: ValidFederationSchema,
@@ -180,7 +188,7 @@ pub(crate) struct FetchDependencyGraph {
 }
 
 // TODO: Write docstrings
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct DeferTracking {
     pub(crate) top_level_deferred: IndexSet<DeferRef>,
     pub(crate) deferred: IndexMap<DeferRef, DeferredInfo>,