From 1cee14962556e31a3d1c902ac01011e5e9fef6b7 Mon Sep 17 00:00:00 2001
From: Tamir Duberstein <tamird@gmail.com>
Date: Fri, 6 Oct 2023 14:36:58 -0400
Subject: [PATCH] GroupMap: add fold_with

This is a generalization of `fold` which takes a function rather than a
value, which removes the need for a `Clone` bound. `fold` is implemented
in terms of `fold_with`.
---
 src/grouping_map.rs | 51 ++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 46 insertions(+), 5 deletions(-)

diff --git a/src/grouping_map.rs b/src/grouping_map.rs
index affcb2121..476583cdc 100644
--- a/src/grouping_map.rs
+++ b/src/grouping_map.rs
@@ -115,6 +115,50 @@ where
         destination_map
     }
 
+    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
+    /// of each group sequentially, passing the previously accumulated value, a reference to the key
+    /// and the current element as arguments, and stores the results in a new map.
+    ///
+    /// `init` is called to obtain the initial value of each accumulator.
+    ///
+    /// `operation` is a function that is invoked on each element with the following parameters:
+    ///  - the current value of the accumulator of the group;
+    ///  - a reference to the key of the group this element belongs to;
+    ///  - the element from the source being accumulated.
+    ///
+    /// Return a `HashMap` associating the key of each group with the result of folding that group's elements.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// #[derive(Debug, Default)]
+    /// struct Accumulator {
+    ///   acc: usize,
+    /// }
+    ///
+    /// let lookup = (1..=7)
+    ///     .into_grouping_map_by(|&n| n % 3)
+    ///     .fold_with(Default::default, |Accumulator { acc }, _key, val| {
+    ///         let acc = acc + val;
+    ///         Accumulator { acc }
+    ///      });
+    ///
+    /// assert_eq!(lookup[&0].acc, 3 + 6);
+    /// assert_eq!(lookup[&1].acc, 1 + 4 + 7);
+    /// assert_eq!(lookup[&2].acc, 2 + 5);
+    /// assert_eq!(lookup.len(), 3);
+    /// ```
+    pub fn fold_with<FI, FO, R>(self, mut init: FI, mut operation: FO) -> HashMap<K, R>
+    where
+        FI: FnMut() -> R,
+        FO: FnMut(R, &K, V) -> R,
+    {
+        self.aggregate(|acc, key, val| {
+            let acc = acc.unwrap_or_else(&mut init);
+            Some(operation(acc, key, val))
+        })
+    }
+
     /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
     /// of each group sequentially, passing the previously accumulated value, a reference to the key
     /// and the current element as arguments, and stores the results in a new map.
@@ -140,15 +184,12 @@ where
     /// assert_eq!(lookup[&2], 2 + 5);
     /// assert_eq!(lookup.len(), 3);
     /// ```
-    pub fn fold<FO, R>(self, init: R, mut operation: FO) -> HashMap<K, R>
+    pub fn fold<FO, R>(self, init: R, operation: FO) -> HashMap<K, R>
     where
         R: Clone,
         FO: FnMut(R, &K, V) -> R,
     {
-        self.aggregate(|acc, key, val| {
-            let acc = acc.unwrap_or_else(|| init.clone());
-            Some(operation(acc, key, val))
-        })
+        self.fold_with(|| init.clone(), operation)
     }
 
     /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements