Refactor SelectionSet::merge_selections_into and remove unnecessary d…

…irectives in Selection::from_element (#5919)

apollo-federation/src/operation/mod.rs

@@ -34,6 +34,7 @@ use crate::compat::coerce_executable_values;
 use crate::error::FederationError;
 use crate::error::SingleFederationError;
 use crate::error::SingleFederationError::Internal;
+use crate::query_graph::graph_path::op_slice_condition_directives;
 use crate::query_graph::graph_path::OpPathElement;
 use crate::query_plan::conditions::Conditions;
 use crate::query_plan::FetchDataKeyRenamer;
@@ -788,6 +789,7 @@ impl Selection {
     pub(crate) fn from_element(
         element: OpPathElement,
         sub_selections: Option<SelectionSet>,
+        unnecessary_directives: Option<&DirectiveList>,
     ) -> Result<Selection, FederationError> {
         // PORT_NOTE: This is TODO item is copied from the JS `selectionOfElement` function.
         // TODO: validate that the subSelection is ok for the element
@@ -799,7 +801,21 @@ impl Selection {
                         "unexpected inline fragment without sub-selections",
                     ));
                 };
-                Ok(InlineFragmentSelection::new(inline_fragment, sub_selections).into())
+                if let Some(unnecessary_directives) = unnecessary_directives {
+                    let directives = inline_fragment
+                        .directives
+                        .iter()
+                        .filter(|dir| !unnecessary_directives.contains(dir))
+                        .cloned()
+                        .collect::<DirectiveList>();
+                    Ok(InlineFragmentSelection::new(
+                        inline_fragment.with_updated_directives(directives),
+                        sub_selections,
+                    )
+                    .into())
+                } else {
+                    Ok(InlineFragmentSelection::new(inline_fragment, sub_selections).into())
+                }
             }
         }
     }
@@ -1670,6 +1686,16 @@ mod inline_fragment_selection {
                 selection_set: self.selection_set.clone(),
             }
         }
+        pub(crate) fn with_updated_directives_and_selection_set(
+            &self,
+            directives: impl Into<DirectiveList>,
+            selection_set: SelectionSet,
+        ) -> Self {
+            Self {
+                inline_fragment: self.inline_fragment.with_updated_directives(directives),
+                selection_set,
+            }
+        }
     }
 
     impl HasSelectionKey for InlineFragmentSelection {
@@ -1905,7 +1931,7 @@ impl SelectionSet {
                         .flat_map(SelectionSet::split_top_level_fields)
                         .filter_map(move |set| {
                             let parent_type = ele.parent_type_position();
-                            Selection::from_element(ele.clone(), Some(set))
+                            Selection::from_element(ele.clone(), Some(set), None)
                                 .ok()
                                 .map(|sel| SelectionSet::from_selection(parent_type, sel))
                         }),
@@ -2025,7 +2051,7 @@ impl SelectionSet {
             type_position,
             selections: Arc::new(SelectionMap::new()),
         };
-        merged.merge_selections_into(normalized_selections.iter())?;
+        merged.merge_selections_into(normalized_selections.iter(), false)?;
         Ok(merged)
     }
 
@@ -2122,7 +2148,7 @@ impl SelectionSet {
             type_position: self.type_position.clone(),
             selections: Arc::new(SelectionMap::new()),
         };
-        expanded.merge_selections_into(expanded_selections.iter())?;
+        expanded.merge_selections_into(expanded_selections.iter(), false)?;
         Ok(expanded)
     }
 
@@ -2496,7 +2522,7 @@ impl SelectionSet {
                 sibling_typename.alias().cloned(),
             );
             let typename_selection =
-                Selection::from_element(field_element.into(), /*subselection*/ None)?;
+                Selection::from_element(field_element.into(), /*subselection*/ None, None)?;
             Ok([typename_selection, updated].into_iter().collect())
         })
     }
@@ -2613,11 +2639,13 @@ impl SelectionSet {
                 let mut selection = Arc::make_mut(&mut self.selections)
                     .entry(ele.key())
                     .or_insert(|| {
+                        let unnecessary_directives = op_slice_condition_directives(path);
                         Selection::from_element(
                             element,
                             // We immediately add a selection afterward to make this selection set
                             // valid.
                             Some(SelectionSet::empty(self.schema.clone(), sub_selection_type)),
+                            Some(&unnecessary_directives),
                         )
                     })?;
                 match &mut selection {
@@ -2649,9 +2677,11 @@ impl SelectionSet {
                     if !ele.is_terminal()? {
                         return Ok(());
                     } else {
+                        let unneeded_directives = op_slice_condition_directives(path);
                         // add leaf
-                        let selection = Selection::from_element(element, None)?;
-                        self.add_local_selection(&selection)?
+                        let selection =
+                            Selection::from_element(element, None, Some(&unneeded_directives))?;
+                        self.add_local_selection(&selection, true)?
                     }
                 } else {
                     let selection_set = selection_set
@@ -2666,8 +2696,13 @@ impl SelectionSet {
                         })
                         .transpose()?
                         .map(|selection_set| selection_set.without_unnecessary_fragments());
-                    let selection = Selection::from_element(element, selection_set)?;
-                    self.add_local_selection(&selection)?
+                    let unneeded_directives = op_slice_condition_directives(path);
+                    let selection = Selection::from_element(
+                        element,
+                        selection_set,
+                        Some(&unneeded_directives),
+                    )?;
+                    self.add_local_selection(&selection, true)?;
                 }
             }
             // If we don't have any path, we rebase and merge in the given sub selections at the root.