Ported expand_fragments

apollo-federation/src/query_plan/operation.rs

@@ -231,6 +231,22 @@ impl Operation {
 
         self
     }
+
+    fn with_updated_selection_set_and_fragments(
+        &self,
+        new_selection_set: SelectionSet,
+        new_fragments: Option<NamedFragments>,
+    ) -> Self {
+        Self {
+            schema: self.schema.clone(),
+            root_kind: self.root_kind,
+            name: self.name.clone(),
+            variables: self.variables.clone(),
+            selection_set: new_selection_set,
+            named_fragments: new_fragments.unwrap_or_default(),
+            directives: self.directives.clone(),
+        }
+    }
 }
 
 /// An analogue of the apollo-compiler type `SelectionSet` with these changes:
@@ -915,9 +931,16 @@ impl Selection {
             Selection::Field(field) => Ok(Selection::from(
                 field.with_updated_selection_set(selection_set),
             )),
-            Selection::InlineFragment(inline_fragment) => Ok(Selection::from(
-                inline_fragment.with_updated_selection_set(selection_set),
-            )),
+            Selection::InlineFragment(inline_fragment) => {
+                let Some(selection_set) = selection_set else {
+                    return Err(FederationError::internal(
+                        "updating inline fragment without a sub-selection set",
+                    ));
+                };
+                Ok(inline_fragment
+                    .with_updated_selection_set(selection_set)
+                    .into())
+            }
             Selection::FragmentSpread(_) => {
                 Err(FederationError::internal("unexpected fragment spread"))
             }
@@ -965,6 +988,44 @@ impl Selection {
             Ok(self.clone())
         }
     }
+
+    /// Expand fragments that are not in the `fragments_to_keep`.
+    fn expand_fragments(
+        &self,
+        fragments_to_keep: &NamedFragments,
+    ) -> Result<SelectionOrSet, FederationError> {
+        match self {
+            Selection::FragmentSpread(fragment) => {
+                if fragments_to_keep.contains(&fragment.spread.data().fragment_name) {
+                    // Keep this spread
+                    Ok(self.clone().into())
+                } else {
+                    // Expand the fragment
+                    let expanded_sub_selections =
+                        fragment.selection_set.expand_fragments(fragments_to_keep)?;
+                    if self.element()?.parent_type_position()
+                        == fragment.spread.data().type_condition_position
+                    {
+                        // The fragment is of the same type as the parent, so we can just use
+                        // the expanded sub-selections directly.
+                        Ok(expanded_sub_selections.into())
+                    } else {
+                        // Create an inline fragment since type condition is necessary.
+                        let inline =
+                            InlineFragmentSelection::from_fragment_spread_selection(fragment)?;
+                        Ok(Selection::from(inline).into())
+                    }
+                }
+            }
+
+            // Otherwise, expand the sub-selections.
+            _ => Ok(self
+                .map_selection_set(|selection_set| {
+                    Ok(Some(selection_set.expand_fragments(fragments_to_keep)?))
+                })?
+                .into()),
+        }
+    }
 }
 
 impl From<FieldSelection> for Selection {
@@ -1625,14 +1686,10 @@ mod normalized_inline_fragment_selection {
     }
 
     impl InlineFragmentSelection {
-        pub(crate) fn with_updated_selection_set(
-            &self,
-            selection_set: Option<SelectionSet>,
-        ) -> Self {
+        pub(crate) fn with_updated_selection_set(&self, selection_set: SelectionSet) -> Self {
             Self {
                 inline_fragment: self.inline_fragment.clone(),
-                //FIXME
-                selection_set: selection_set.unwrap(),
+                selection_set,
             }
         }
 
@@ -1869,6 +1926,20 @@ impl FromIterator<Selection> for SelectionMapperReturn {
     }
 }
 
+// Note: `expand_fragments` methods may return a selection or a selection set. The items in a
+//       selection set needs to be cloned here, since it's sub-selections are contained in an
+//       `Arc`.
+impl From<SelectionOrSet> for SelectionMapperReturn {
+    fn from(value: SelectionOrSet) -> Self {
+        match value {
+            SelectionOrSet::Selection(selection) => selection.into(),
+            SelectionOrSet::SelectionSet(selections) => {
+                Vec::from_iter(selections.selections.values().cloned()).into()
+            }
+        }
+    }
+}
+
 impl SelectionSet {
     pub(crate) fn empty(
         schema: ValidFederationSchema,
@@ -2258,6 +2329,14 @@ impl SelectionSet {
         Ok(())
     }
 
+    /// Expand fragments that are not in the `fragments_to_keep`.
+    fn expand_fragments(
+        &self,
+        fragments_to_keep: &NamedFragments,
+    ) -> Result<Self, FederationError> {
+        self.lazy_map(|selection| Ok(selection.expand_fragments(fragments_to_keep)?.into()))
+    }
+
     /// Modifies the provided selection set to optimize the handling of __typename selections for query planning.
     ///
     /// __typename information can always be provided by any subgraph declaring that type. While this data can be