From 21fef11ef016534ee9a1a6f9c283e951caf67a25 Mon Sep 17 00:00:00 2001
From: Takafumi Arakaki <aka.tkf@gmail.com>
Date: Sun, 26 Jul 2020 09:29:04 -0700
Subject: [PATCH] Support foldxl etc.

---
 src/LazyGroupBy.jl  | 29 ++++++++++++++++-------------
 test/test_foldl.jl  | 17 +++++++++++++----
 test/test_reduce.jl | 12 ++++++++++++
 3 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/src/LazyGroupBy.jl b/src/LazyGroupBy.jl
index 779cd0a..35b2f4e 100644
--- a/src/LazyGroupBy.jl
+++ b/src/LazyGroupBy.jl
@@ -22,6 +22,17 @@ if !@isdefined only
     using Compat: only
 end
 
+const FoldFunction0 = Union{typeof(foldl),typeof(reduce),typeof(dreduce)}
+if @isdefined foldxl
+    const FoldFunction1 = Union{typeof(foldxl),typeof(foldxt),typeof(foldxd)}
+else
+    const FoldFunction1 = Union{}
+    const foldxl = foldl
+    const foldxt = reduce
+    const foldxd = dreduce
+end
+const FoldFunction = Union{FoldFunction0,FoldFunction1}
+
 # Used only for `KWBroadcastedInner` below:
 function air end
 struct Aired{T}
@@ -98,21 +109,13 @@ function _impl_fold(fold, kwargs, rf, xf, group)
     return d
 end
 
-impl(::typeof(foldl), kwargs, rf, xf::Transducer, group::GroupedBy) =
-    _impl_fold(foldl, kwargs, rf, xf, group)
-impl(::typeof(reduce), kwargs, rf, xf::Transducer, group::GroupedBy) =
-    _impl_fold(reduce, kwargs, rf, xf, group)
-impl(::typeof(dreduce), kwargs, rf, xf::Transducer, group::GroupedBy) =
-    _impl_fold(dreduce, kwargs, rf, xf, group)
+impl(fold::FoldFunction, kwargs, rf, xf::Transducer, group::GroupedBy) =
+    _impl_fold(fold, kwargs, rf, xf, group)
+impl(fold::FoldFunction, kwargs, rf, xs) = impl(fold, kwargs, rf, Map(identity), xs)
 
-impl(::typeof(foldl), kwargs, rf, xs) = impl(foldl, kwargs, rf, Map(identity), xs)
 impl(::typeof(mapfoldl), kwargs, f, rf, xs) = impl(foldl, kwargs, rf, Map(f), xs)
-
-impl(::typeof(reduce), kwargs, rf, xs) = impl(reduce, kwargs, rf, Map(identity), xs)
 impl(::typeof(mapreduce), kwargs, f, rf, xs) = impl(reduce, kwargs, rf, Map(f), xs)
 
-impl(::typeof(dreduce), kwargs, rf, xs) = impl(dreduce, kwargs, rf, Map(identity), xs)
-
 impl(::typeof(collect), kwargs, xf::Transducer, group::GroupedBy) = foldl(
     right,
     GroupBy(group.key, opcompose(Map(last), xf, Map(SingletonVector)), append!!),
@@ -121,7 +124,7 @@ impl(::typeof(collect), kwargs, xf::Transducer, group::GroupedBy) = foldl(
 impl(::typeof(collect), kwargs, group::GroupedBy) =
     impl(collect, kwargs, IdentityTransducer(), group)
 
-impl(::typeof(tcollect), kwargs, xf::Transducer, group::GroupedBy) = reduce(
+impl(::typeof(tcollect), kwargs, xf::Transducer, group::GroupedBy) = foldxt(
     right,
     GroupBy(group.key, opcompose(Map(last), xf, Map(SingletonVector)), append!!),
     group.collection,
@@ -129,7 +132,7 @@ impl(::typeof(tcollect), kwargs, xf::Transducer, group::GroupedBy) = reduce(
 impl(::typeof(tcollect), kwargs, group::GroupedBy) =
     impl(tcollect, kwargs, IdentityTransducer(), group)
 
-impl(::typeof(dcollect), kwargs, xf::Transducer, group::GroupedBy) = dreduce(
+impl(::typeof(dcollect), kwargs, xf::Transducer, group::GroupedBy) = foldxd(
     right,
     GroupBy(group.key, opcompose(Map(last), xf, Map(SingletonVector)), append!!),
     group.collection;
diff --git a/test/test_foldl.jl b/test/test_foldl.jl
index e6c3bf0..7805b81 100644
--- a/test/test_foldl.jl
+++ b/test/test_foldl.jl
@@ -2,12 +2,21 @@ module TestFoldl
 
 using Test
 using LazyGroupBy
+using LazyGroupBy: foldxl
 
 @testset "tuple" begin
-    @test foldl.(tuple, grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5])) ==
-          Dict(false => ((0, 4), 0), true => ((((((7, 3), 1), 5), 9), 3), 5))
-    @test foldl.(tuple, grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5]); init = -1) ==
-          Dict(false => (((-1, 0), 4), 0), true => (((((((-1, 7), 3), 1), 5), 9), 3), 5))
+    @testset for fold in [foldl, foldxl]
+        foldl = nothing
+
+        @test fold.(tuple, grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5])) ==
+              Dict(false => ((0, 4), 0), true => ((((((7, 3), 1), 5), 9), 3), 5))
+
+        @test fold.(tuple, grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5]); init = -1) ==
+              Dict(
+            false => (((-1, 0), 4), 0),
+            true => (((((((-1, 7), 3), 1), 5), 9), 3), 5),
+        )
+    end
 end
 
 extrema_rf((min1, max1), (min2, max2)) = (min(min1, min2), max(max1, max2))
diff --git a/test/test_reduce.jl b/test/test_reduce.jl
index ee7d6b6..46785ac 100644
--- a/test/test_reduce.jl
+++ b/test/test_reduce.jl
@@ -2,6 +2,8 @@ module TestReduce
 
 using Test
 using LazyGroupBy
+using LazyGroupBy: foldxl, foldxt
+using Transducers: Map
 
 extrema_rf((min1, max1), (min2, max2)) = (min(min1, min2), max(max1, max2))
 
@@ -11,6 +13,16 @@ extrema_rf((min1, max1), (min2, max2)) = (min(min1, min2), max(max1, max2))
         extrema_rf,
         grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5]),
     ) == Dict(false => (0, 4), true => (1, 9))
+
+    @testset for fold in [foldl, foldxl, foldxt, reduce]
+        foldl = nothing
+
+        @test fold.(
+            extrema_rf,
+            Map(x -> (x, x)),
+            grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5]),
+        ) == Dict(false => (0, 4), true => (1, 9))
+    end
 end
 
 end  # module