ok

Author: dotansimha

lib/query-planner/src/ast/mismatch_finder.rs

@@ -8,7 +8,7 @@ use crate::{
         selection_item::SelectionItem,
         selection_set::{FieldSelection, SelectionSet},
     },
-    planner::fetch::selections::FetchStepSelections,
+    planner::fetch::{selections::FetchStepSelections, state::MultiTypeFetchStep},
     state::{
         subgraph_state::{SubgraphDefinition, SubgraphState},
         supergraph_state::{SubgraphName, SupergraphState, TypeNode},
@@ -31,19 +31,19 @@ impl<'a> SelectionMismatchFinder<'a> {
     pub fn find_mismatches_in_node(
         &self,
         subgraph_name: &SubgraphName,
-        selections: &FetchStepSelections,
+        selections: &FetchStepSelections<MultiTypeFetchStep>,
     ) -> MismatchesFound {
         let mut mismtaches_found = MismatchesFound::new();
 
-        for (definition_name, selection_set) in selections.selections() {
+        for (definition_name, selection_set) in selections.iter_selections() {
             let subgraph_state = self
                 .supergraph_state
                 .subgraphs_state
                 .get(subgraph_name)
                 .unwrap();
 
             let entrypoint_type = subgraph_state.definitions.get(definition_name).unwrap();
-            let start_path = MergePath::new_empty_root(definition_name);
+            let start_path = MergePath::default();
 
             handle_selection_set(
                 definition_name,

lib/query-planner/src/ast/mod.rs

@@ -12,5 +12,5 @@ pub(crate) mod merge_path;
 pub(crate) mod type_aware_selection;
 pub mod value;
 
-// pub(crate) mod mismatch_finder;
-// pub(crate) mod safe_merge;
+pub(crate) mod mismatch_finder;
+pub(crate) mod safe_merge;

lib/query-planner/src/ast/safe_merge.rs

@@ -59,7 +59,7 @@ impl SafeSelectionSetMerger {
             source,
             (self_used_for_requires, other_used_for_requires),
             as_first,
-            MergePath::empty(),
+            MergePath::default(),
             &mut aliases_performed,
         );
 

lib/query-planner/src/ast/type_aware_selection.rs

@@ -154,7 +154,7 @@ fn selection_item_is_subset_of(source: &SelectionItem, target: &SelectionItem) -
     }
 }
 
-fn selection_items_are_subset_of(source: &[SelectionItem], target: &[SelectionItem]) -> bool {
+pub fn selection_items_are_subset_of(source: &[SelectionItem], target: &[SelectionItem]) -> bool {
     target.iter().all(|target_node| {
         source
             .iter()

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

@@ -7,7 +7,7 @@ use crate::graph::node::Node;
 use crate::graph::Graph;
 use crate::planner::fetch::fetch_step_data::{FetchStepData, FetchStepKind};
 use crate::planner::fetch::selections::FetchStepSelections;
-use crate::planner::fetch::state::SingleTypeFetchStep;
+use crate::planner::fetch::state::{MultiTypeFetchStep, SingleTypeFetchStep};
 use crate::planner::plan_nodes::{FetchNodePathSegment, FetchRewrite, ValueSetter};
 use crate::planner::tree::query_tree::QueryTree;
 use crate::planner::tree::query_tree_node::{MutationFieldPosition, QueryTreeNode};
@@ -40,19 +40,48 @@ impl Default for FetchGraph<SingleTypeFetchStep> {
 }
 
 impl FetchGraph<SingleTypeFetchStep> {
+    pub fn to_multi_type(self) -> FetchGraph<MultiTypeFetchStep> {
+        let new_graph = self
+            .graph
+            .map(|_, w| w.clone().into_multi_type(), |_, _| ());
+
+        FetchGraph {
+            graph: new_graph,
+            root_index: self.root_index,
+        }
+    }
+
     pub fn new() -> Self {
         Self {
             graph: StableDiGraph::new(),
             root_index: None,
         }
     }
+}
 
-    pub fn all_nodes(&self) -> NodeReferences<'_, FetchStepData<SingleTypeFetchStep>> {
+impl<State> FetchGraph<State> {
+    pub fn all_nodes(&self) -> NodeReferences<'_, FetchStepData<State>> {
         self.graph.node_references()
     }
+
+    pub fn get_step_data_mut(
+        &mut self,
+        index: NodeIndex,
+    ) -> Result<&mut FetchStepData<State>, FetchGraphError> {
+        self.graph
+            .node_weight_mut(index)
+            .ok_or(FetchGraphError::MissingStep(
+                index.index(),
+                String::from("when getting mutable step data"),
+            ))
+    }
 }
 
-impl FetchGraph<SingleTypeFetchStep> {
+impl<State> FetchGraph<State> {
+    pub fn step_indices(&self) -> NodeIndices<FetchStepData<State>> {
+        self.graph.node_indices()
+    }
+
     pub fn parents_of(&self, index: NodeIndex) -> petgraph::stable_graph::Edges<'_, (), Directed> {
         self.graph.edges_directed(index, Direction::Incoming)
     }
@@ -65,14 +94,10 @@ impl FetchGraph<SingleTypeFetchStep> {
         has_path_connecting(&self.graph, ancestor, descendant, None)
     }
 
-    pub fn step_indices(&self) -> NodeIndices<FetchStepData<SingleTypeFetchStep>> {
-        self.graph.node_indices()
-    }
-
     pub fn get_step_data(
         &self,
         index: NodeIndex,
-    ) -> Result<&FetchStepData<SingleTypeFetchStep>, FetchGraphError> {
+    ) -> Result<&FetchStepData<State>, FetchGraphError> {
         self.graph
             .node_weight(index)
             .ok_or(FetchGraphError::MissingStep(
@@ -81,51 +106,6 @@ impl FetchGraph<SingleTypeFetchStep> {
             ))
     }
 
-    pub fn get_step_data_mut(
-        &mut self,
-        index: NodeIndex,
-    ) -> Result<&mut FetchStepData<SingleTypeFetchStep>, FetchGraphError> {
-        self.graph
-            .node_weight_mut(index)
-            .ok_or(FetchGraphError::MissingStep(
-                index.index(),
-                String::from("when getting mutable step data"),
-            ))
-    }
-
-    pub fn get_pair_of_steps_mut(
-        &mut self,
-        index1: NodeIndex,
-        index2: NodeIndex,
-    ) -> Result<
-        (
-            &mut FetchStepData<SingleTypeFetchStep>,
-            &mut FetchStepData<SingleTypeFetchStep>,
-        ),
-        FetchGraphError,
-    > {
-        // `index_twice_mut` panics when indexes are equal
-        if index1 == index2 {
-            return Err(FetchGraphError::SameNodeIndex(index1.index()));
-        }
-
-        // `index_twice_mut` panics when nodes do not exist
-        if self.graph.node_weight(index1).is_none() {
-            return Err(FetchGraphError::MissingStep(
-                index1.index(),
-                String::from("when checking existence"),
-            ));
-        }
-        if self.graph.node_weight(index2).is_none() {
-            return Err(FetchGraphError::MissingStep(
-                index2.index(),
-                String::from("when checking existence"),
-            ));
-        }
-
-        Ok(self.graph.index_twice_mut(index1, index2))
-    }
-
     pub fn connect(&mut self, parent_index: NodeIndex, child_index: NodeIndex) -> EdgeIndex {
         self.graph.update_edge(parent_index, child_index, ())
     }
@@ -141,13 +121,13 @@ impl FetchGraph<SingleTypeFetchStep> {
         self.graph.remove_node(index).is_some_and(|_| true)
     }
 
-    pub fn add_step(&mut self, data: FetchStepData<SingleTypeFetchStep>) -> NodeIndex {
+    pub fn add_step(&mut self, data: FetchStepData<State>) -> NodeIndex {
         self.graph.add_node(data)
     }
 
     pub fn bfs<F>(&self, root_index: NodeIndex, mut visitor: F) -> Option<NodeIndex>
     where
-        F: FnMut(&NodeIndex, &FetchStepData<SingleTypeFetchStep>) -> bool,
+        F: FnMut(&NodeIndex, &FetchStepData<State>) -> bool,
     {
         self.graph.node_weight(root_index)?;
 
@@ -167,7 +147,9 @@ impl FetchGraph<SingleTypeFetchStep> {
 
         None
     }
+}
 
+impl FetchGraph<MultiTypeFetchStep> {
     #[instrument(level = "trace", skip_all)]
     pub fn collect_variable_usages(&mut self) -> Result<(), FetchGraphError> {
         let nodes_idx = self.graph.node_indices().collect::<Vec<_>>();
@@ -183,9 +165,42 @@ impl FetchGraph<SingleTypeFetchStep> {
 
         Ok(())
     }
+
+    pub fn get_pair_of_steps_mut(
+        &mut self,
+        index1: NodeIndex,
+        index2: NodeIndex,
+    ) -> Result<
+        (
+            &mut FetchStepData<MultiTypeFetchStep>,
+            &mut FetchStepData<MultiTypeFetchStep>,
+        ),
+        FetchGraphError,
+    > {
+        // `index_twice_mut` panics when indexes are equal
+        if index1 == index2 {
+            return Err(FetchGraphError::SameNodeIndex(index1.index()));
+        }
+
+        // `index_twice_mut` panics when nodes do not exist
+        if self.graph.node_weight(index1).is_none() {
+            return Err(FetchGraphError::MissingStep(
+                index1.index(),
+                String::from("when checking existence"),
+            ));
+        }
+        if self.graph.node_weight(index2).is_none() {
+            return Err(FetchGraphError::MissingStep(
+                index2.index(),
+                String::from("when checking existence"),
+            ));
+        }
+
+        Ok(self.graph.index_twice_mut(index1, index2))
+    }
 }
 
-impl Display for FetchGraph<SingleTypeFetchStep> {
+impl<State> Display for FetchGraph<State> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         writeln!(f, "Nodes:")?;
         for node_index in self.graph.node_indices() {
@@ -262,7 +277,7 @@ fn create_noop_fetch_step(
         variable_usages: None,
         variable_definitions: None,
         mutation_field_position: None,
-        // internal_aliases_locations: Vec::new(),
+        internal_aliases_locations: Vec::new(),
     })
 }
 
@@ -293,7 +308,7 @@ fn create_fetch_step_for_entity_call(
         variable_usages: None,
         variable_definitions: None,
         mutation_field_position: None,
-        // internal_aliases_locations: Vec::new(),
+        internal_aliases_locations: Vec::new(),
     })
 }
 
@@ -320,7 +335,7 @@ fn create_fetch_step_for_root_move(
         input_rewrites: None,
         output_rewrites: None,
         mutation_field_position,
-        // internal_aliases_locations: Vec::new(),
+        internal_aliases_locations: Vec::new(),
     });
 
     fetch_graph.connect(root_step_index, idx);
@@ -1574,7 +1589,7 @@ fn process_query_node(
     }
 }
 
-pub fn find_graph_roots(graph: &FetchGraph<SingleTypeFetchStep>) -> Vec<NodeIndex> {
+pub fn find_graph_roots<State>(graph: &FetchGraph<State>) -> Vec<NodeIndex> {
     let mut roots = Vec::new();
 
     // Iterate over all nodes in the graph
@@ -1598,7 +1613,7 @@ pub fn build_fetch_graph_from_query_tree(
     supergraph: &SupergraphState,
     override_context: &PlannerOverrideContext,
     query_tree: QueryTree,
-) -> Result<FetchGraph<SingleTypeFetchStep>, FetchGraphError> {
+) -> Result<FetchGraph<MultiTypeFetchStep>, FetchGraphError> {
     let mut fetch_graph = FetchGraph::new();
 
     process_query_node(
@@ -1633,7 +1648,8 @@ pub fn build_fetch_graph_from_query_tree(
 
     // fine to unwrap as we have already checked the length
     fetch_graph.root_index = Some(*root_indexes.first().unwrap());
-    // fetch_graph.optimize(supergraph)?;
+    let mut fetch_graph = fetch_graph.to_multi_type();
+    fetch_graph.optimize(supergraph)?;
     fetch_graph.collect_variable_usages()?;
 
     trace!("fetch graph after optimizations:");