FED-47: Finish porting compute_best_plan_from_closed_branches

src/query_graph/graph_path.rs

@@ -9,6 +9,7 @@ use crate::query_graph::condition_resolver::{
 };
 use crate::query_graph::path_tree::OpPathTree;
 use crate::query_graph::{QueryGraph, QueryGraphEdgeTransition, QueryGraphNodeType};
+use crate::query_plan::display::{write_indented_lines, State as IndentedFormatter};
 use crate::query_plan::operation::normalized_field_selection::{
     NormalizedField, NormalizedFieldData, NormalizedFieldSelection,
 };
@@ -430,6 +431,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()
@@ -438,6 +455,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.
@@ -646,11 +669,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>>);

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();

src/query_plan/display.rs

@@ -7,29 +7,40 @@ pub(crate) struct State<'fmt, 'fmt2> {
     output: &'fmt mut fmt::Formatter<'fmt2>,
 }
 
-impl State<'_, '_> {
-    fn write<T: fmt::Display>(&mut self, value: T) -> fmt::Result {
+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 write<T: fmt::Display>(&mut self, value: T) -> fmt::Result {
         write!(self.output, "{}", value)
     }
 
-    fn new_line(&mut self) -> fmt::Result {
+    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(())
     }
 
-    fn indent_no_new_line(&mut self) {
+    pub(crate) fn indent_no_new_line(&mut self) {
         self.indent_level += 1;
     }
 
-    fn indent(&mut self) -> fmt::Result {
+    pub(crate) fn indent(&mut self) -> fmt::Result {
         self.indent_no_new_line();
         self.new_line()
     }
 
-    fn dedent(&mut self) -> fmt::Result {
+    pub(crate) fn dedent(&mut self) -> fmt::Result {
         self.indent_level -= 1;
         self.new_line()
     }
@@ -375,7 +386,7 @@ fn write_selections(
     state.write("}")
 }
 
-fn write_indented_lines<T>(
+pub(crate) fn write_indented_lines<T>(
     state: &mut State<'_, '_>,
     values: &[T],
     mut write_line: impl FnMut(&mut State<'_, '_>, &T) -> fmt::Result,

src/query_plan/fetch_dependency_graph.rs

@@ -106,6 +106,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.
@@ -150,7 +158,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,
@@ -177,7 +185,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>,
@@ -2261,11 +2269,22 @@ fn create_fetch_initial_path(
 
 fn extract_defer_from_operation(
     _dependency_graph: &mut FetchDependencyGraph,
-    _operation: &OpPathElement,
-    _defer_context: &DeferContext,
+    operation: &OpPathElement,
+    defer_context: &DeferContext,
     _node_path: &FetchDependencyGraphNodePath,
 ) -> (Option<OpPathElement>, DeferContext) {
-    todo!() // Port `extractDeferFromOperation`
+    // PORT_NOTE: temporary straw-man implementation (waiting for FED-189)
+    let updated_path_to_defer_parent = defer_context
+        .path_to_defer_parent
+        .with_pushed(operation.clone().into());
+    let updated_context = DeferContext {
+        path_to_defer_parent: updated_path_to_defer_parent.into(),
+        // Following fields are identical to those of `defer_context`.
+        current_defer_ref: defer_context.current_defer_ref.clone(),
+        active_defer_ref: defer_context.active_defer_ref.clone(),
+        is_part_of_query: defer_context.is_part_of_query,
+    };
+    (Some(operation.clone()), updated_context)
 }
 
 fn handle_requires(