Repository: rust-itertools/itertools
Pull request #778 (2023-10-06 18:42:01)

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`.

Philippe-Cholet: I'm not in favor of a breaking change without good reason. Plus, add `fold_with` seems way more reasonable.
I would like another opinion on this before you work on this though.
phimuemue: If we accept this, I'd second @Philippe-Cholet's proposal and implement `fold` in terms of `fold_with`.
tamird: Done.
Philippe-Cholet: I think `fold_with` example should be updated to something that `fold` can not do by itself (not cloning copying some value). Otherwise, there would be no need for such generalization.

I'm wondering if `FI` should be with `FnMut` as well?!
tamird: Done. The example is a bit contrived, let me know if you have a specific suggestion please.
Philippe-Cholet: I did not thought of anything and `Default::default` in particular but I see it could be useful. However, the `Accumulator` struct is artificial and I came up with another idea, short and that could be useful: `|| Vec::with_capacity(64)`.
The resulting vector would be clonable but the capacity of each clone would be 0 so using this closure with `fold_with` does something that `fold(Vec::with_capacity(64), ...)` would not be able.
tamird: That doesn't seem very compelling; such code would still compile albiet with slightly worse performance. I would prefer to demonstrate a use case that can't be written using `fold` at all.
Philippe-Cholet: For such small example, it sure is not very compelling.
Looking at https://doc.rust-lang.org/std/clone/trait.Clone.html#implementors , it's not easy to find a non-clonable std type, leading to create our own type for this. I guess `fold_with` would be used on user-defined structs anyway. Let's go with your `Accumulator`.
We should have one quicktest about this and it would be ready.
tamird: Done.
phimuemue: As for the test, I'm fine with the copy-paste (it's a specific test meant to test one thing only, and the danger that these two tests diverge seems low).

I first thought that I'm fine with it API wise (it's in line with other `_with`-functions). But one question suggested itself: Should we - in the spirit that functions are either simple to use (such as the existing `fold`) or as general as possible (possibly the new `fold_with`), but not something in between - pass `key` and `val` to `init`?
Philippe-Cholet: I did not thought of that, it would definitely be more general. I'm wondering what would be the use case though, especially having `&val`.
- `FI: FnMut() -> R`
- `FI: FnMut(&K) -> R`
- `FI: FnMut(&K, &V) -> R`

But with the simple `fold` remaining as is, I guess we should be as general as possible with `fold_with`.

I would keep the doctest example simple though with `|_key, _val| Default::default()`. It suggests `Default` use instead of `Clone` while keeping "key/val" usages clear enough.
tamird: I did write this with a key parameter at first, but then realized the result is almost the same as `aggregate`, save for one `Option`.

I ended up using `aggregate` in my own code FWIW. Let me know what you prefer.
tamird: @Philippe-Cholet @phimuemue did you want this to take `&K` or is it suitable as is?
phimuemue: > @Philippe-Cholet @phimuemue did you want this to take `&K` or is it suitable as is?

Since you already cooked up a version involving `&K`, let's go with the most general one, i.e. `FnMut(&K, &V)` - it's more general than accepting only `&K` while not being (much) more expensive.

I agree the function is very similar to `aggregate`, but it has its use: It lets you operate outside `Option`-land.
Philippe-Cholet: Sorry for the late response, I took a little break.

> I did write this with a key parameter at first, but then realized the result is almost the same as `aggregate`, save for one `Option`.
> 
> I ended up using `aggregate` in my own code FWIW. Let me know what you prefer.

Maybe it's because I'm not very familiar with `GroupMap` but I don't see what you're saying. But apparently phimuemue does so not a real problem.

Otherwise agree with phimuemue.
tamird: Done.

Patch:
diff --git a/src/grouping_map.rs b/src/grouping_map.rs
index affcb21..8799c6e 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(|_key| 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(&K) -> R,
+        FO: FnMut(R, &K, V) -> R,
+    {
+        self.aggregate(|acc, key, val| {
+            let acc = acc.unwrap_or_else(|| init(key));
+            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(|_: &K| init.clone(), operation)
     }
 
     /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements


Test patch:
diff --git a/tests/quick.rs b/tests/quick.rs
index f0c8f2d..d65e274 100644
--- a/tests/quick.rs
+++ b/tests/quick.rs
@@ -1472,6 +1472,35 @@ quickcheck! {
         }
     }
 
+    fn correct_grouping_map_by_fold_with_modulo_key(a: Vec<u8>, modulo: u8) -> () {
+        #[derive(Debug, Default, PartialEq)]
+        struct Accumulator {
+            acc: u64,
+        }
+
+        let modulo = if modulo == 0 { 1 } else { modulo } as u64; // Avoid `% 0`
+        let lookup = a.iter().map(|&b| b as u64) // Avoid overflows
+            .into_grouping_map_by(|i| i % modulo)
+            .fold_with(|_key| Default::default(), |Accumulator { acc }, &key, val| {
+                assert!(val % modulo == key);
+                let acc = acc + val;
+                Accumulator { acc }
+            });
+
+        let group_map_lookup = a.iter()
+            .map(|&b| b as u64)
+            .map(|i| (i % modulo, i))
+            .into_group_map()
+            .into_iter()
+            .map(|(key, vals)| (key, vals.into_iter().sum())).map(|(key, acc)| (key,Accumulator { acc }))
+            .collect::<HashMap<_,_>>();
+        assert_eq!(lookup, group_map_lookup);
+
+        for (&key, &Accumulator { acc: sum }) in lookup.iter() {
+            assert_eq!(sum, a.iter().map(|&b| b as u64).filter(|&val| val % modulo == key).sum::<u64>());
+        }
+    }
+
     fn correct_grouping_map_by_fold_modulo_key(a: Vec<u8>, modulo: u8) -> () {
         let modulo = if modulo == 0 { 1 } else { modulo } as u64; // Avoid `% 0`
         let lookup = a.iter().map(|&b| b as u64) // Avoid overflows