chore(federation): refactor fragment usage counting

apollo-federation/src/operation/mod.rs

@@ -841,25 +841,6 @@ impl Selection {
         }
     }
 
-    fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, i32>) {
-        match self {
-            Selection::Field(field_selection) => {
-                if let Some(s) = field_selection.selection_set.clone() {
-                    s.collect_used_fragment_names(aggregator)
-                }
-            }
-            Selection::InlineFragment(inline) => {
-                inline.selection_set.collect_used_fragment_names(aggregator);
-            }
-            Selection::FragmentSpread(fragment) => {
-                let current_count = aggregator
-                    .entry(fragment.spread.fragment_name.clone())
-                    .or_default();
-                *current_count += 1;
-            }
-        }
-    }
-
     pub(crate) fn with_updated_selection_set(
         &self,
         selection_set: Option<SelectionSet>,
@@ -1012,18 +993,6 @@ impl Fragment {
         })
     }
 
-    // PORT NOTE: in JS code this is stored on the fragment
-    pub(crate) fn fragment_usages(&self) -> HashMap<Name, i32> {
-        let mut usages = HashMap::new();
-        self.selection_set.collect_used_fragment_names(&mut usages);
-        usages
-    }
-
-    // PORT NOTE: in JS code this is stored on the fragment
-    pub(crate) fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, i32>) {
-        self.selection_set.collect_used_fragment_names(aggregator)
-    }
-
     fn has_defer(&self) -> bool {
         self.selection_set.has_defer()
     }
@@ -2624,12 +2593,6 @@ impl SelectionSet {
         Ok(())
     }
 
-    pub(crate) fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, i32>) {
-        self.selections
-            .iter()
-            .for_each(|(_, s)| s.collect_used_fragment_names(aggregator));
-    }
-
     /// Removes the @defer directive from all selections without removing that selection.
     fn without_defer(&mut self) {
         for (_key, mut selection) in Arc::make_mut(&mut self.selections).iter_mut() {
@@ -3435,17 +3398,6 @@ impl NamedFragments {
         self.fragments.contains_key(name)
     }
 
-    /**
-     * Collect the usages of fragments that are used within the selection of other fragments.
-     */
-    pub(crate) fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, i32>) {
-        for fragment in self.fragments.values() {
-            fragment
-                .selection_set
-                .collect_used_fragment_names(aggregator);
-        }
-    }
-
     /// JS PORT NOTE: In JS implementation this method was named mapInDependencyOrder and accepted a lambda to
     /// apply transformation on the fragments. It was called when rebasing/filtering/expanding selection sets.
     /// JS PORT NOTE: In JS implementation this method was potentially returning `undefined`. In order to simplify the code
@@ -3466,7 +3418,7 @@ impl NamedFragments {
         //       the outcome of `map_to_expanded_selection_sets`.
         let mut fragments_map: IndexMap<Name, FragmentDependencies> = IndexMap::default();
         for fragment in fragments.values() {
-            let mut fragment_usages: HashMap<Name, i32> = HashMap::new();
+            let mut fragment_usages = HashMap::new();
             NamedFragments::collect_fragment_usages(&fragment.selection_set, &mut fragment_usages);
             let usages: Vec<Name> = fragment_usages.keys().cloned().collect::<Vec<Name>>();
             fragments_map.insert(
@@ -3509,10 +3461,10 @@ impl NamedFragments {
         mapped_fragments
     }
 
-    // JS PORT - we need to calculate those for both executable::SelectionSet and SelectionSet
+    /// Just like our `SelectionSet::used_fragments`, but with apollo-compiler types
     fn collect_fragment_usages(
         selection_set: &executable::SelectionSet,
-        aggregator: &mut HashMap<Name, i32>,
+        aggregator: &mut HashMap<Name, u32>,
     ) {
         selection_set.selections.iter().for_each(|s| match s {
             executable::Selection::Field(f) => {
@@ -3605,6 +3557,60 @@ impl RebasedFragments {
     }
 }
 
+// Collect fragment usages from operation types.
+
+impl Selection {
+    fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, u32>) {
+        match self {
+            Selection::Field(field_selection) => {
+                if let Some(s) = &field_selection.selection_set {
+                    s.collect_used_fragment_names(aggregator)
+                }
+            }
+            Selection::InlineFragment(inline) => {
+                inline.selection_set.collect_used_fragment_names(aggregator);
+            }
+            Selection::FragmentSpread(fragment) => {
+                let current_count = aggregator
+                    .entry(fragment.spread.fragment_name.clone())
+                    .or_default();
+                *current_count += 1;
+            }
+        }
+    }
+}
+
+impl SelectionSet {
+    pub(crate) fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, u32>) {
+        for s in self.selections.values() {
+            s.collect_used_fragment_names(aggregator);
+        }
+    }
+
+    pub(crate) fn used_fragments(&self) -> HashMap<Name, u32> {
+        let mut usages = HashMap::new();
+        self.collect_used_fragment_names(&mut usages);
+        usages
+    }
+}
+
+impl Fragment {
+    pub(crate) fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, u32>) {
+        self.selection_set.collect_used_fragment_names(aggregator)
+    }
+}
+
+impl NamedFragments {
+    /// Collect the usages of fragments that are used within the selection of other fragments.
+    pub(crate) fn collect_used_fragment_names(&self, aggregator: &mut HashMap<Name, u32>) {
+        for fragment in self.fragments.values() {
+            fragment
+                .selection_set
+                .collect_used_fragment_names(aggregator);
+        }
+    }
+}
+
 // Collect used variables from operation types.
 
 pub(crate) struct VariableCollector<'s> {

apollo-federation/src/operation/optimize.rs

@@ -1128,8 +1128,6 @@ impl NamedFragments {
         selection_set: &SelectionSet,
         min_usage_to_optimize: u32,
     ) -> Result<SelectionSet, FederationError> {
-        let min_usage_to_optimize: i32 = min_usage_to_optimize.try_into().unwrap_or(i32::MAX);
-
         // Call `reduce_inner` repeatedly until we reach a fix-point, since newly computed
         // selection set may drop some fragment references due to normalization, which could lead
         // to further reduction.
@@ -1168,20 +1166,16 @@ impl NamedFragments {
     }
 
     /// The inner loop body of `reduce` method.
-    /// - Takes i32 `min_usage_to_optimize` since `collect_used_fragment_names` counts usages in
-    ///   i32.
     fn reduce_inner(
         &mut self,
         selection_set: &SelectionSet,
-        min_usage_to_optimize: i32,
+        min_usage_to_optimize: u32,
     ) -> Result<SelectionSet, FederationError> {
-        // Initial computation of fragment usages in `selection_set`.
-        let mut usages = HashMap::new();
-        selection_set.collect_used_fragment_names(&mut usages);
+        let mut usages = selection_set.used_fragments();
 
         // Short-circuiting: Nothing was used => Drop everything (selection_set is unchanged).
         if usages.is_empty() {
-            self.retain(|_, _| false);
+            *self = Default::default();
             return Ok(selection_set.clone());
         }
 
@@ -1252,7 +1246,7 @@ impl NamedFragments {
         )
     }
 
-    fn update_usages(usages: &mut HashMap<Name, i32>, fragment: &Node<Fragment>, usage_count: i32) {
+    fn update_usages(usages: &mut HashMap<Name, u32>, fragment: &Node<Fragment>, usage_count: u32) {
         let mut inner_usages = HashMap::new();
         fragment.collect_used_fragment_names(&mut inner_usages);