diff --git a/base/reduce.jl b/base/reduce.jl index 425d8d6f28f2f..d4b7402fa7d45 100644 --- a/base/reduce.jl +++ b/base/reduce.jl @@ -36,31 +36,80 @@ mul_prod(x::Real, y::Real)::Real = x * y ## foldl && mapfoldl -function mapfoldl_impl(f, op, nt::NamedTuple{(:init,)}, itr, i...) - init = nt.init +mapfoldl_impl(f, op, nt, itr) = foldl_impl(op, nt, Generator(f, itr)) + +function foldl_impl(op, nt, itr) + op′, itr′ = _xfadjoint(BottomRF(op), itr) + return _foldl_impl(op′, nt, itr′) +end + +function _foldl_impl(op, nt, itr) + init = get(nt, :init, _InitialValue()) # Unroll the while loop once; if init is known, the call to op may # be evaluated at compile time - y = iterate(itr, i...) + y = iterate(itr) y === nothing && return init - v = op(init, f(y[1])) + v = op(init, y[1]) while true y = iterate(itr, y[2]) y === nothing && break - v = op(v, f(y[1])) + v = op(v, y[1]) end + v isa _InitialValue && reduce_empty_iter(op, itr, IteratorEltype(itr)) return v end -function mapfoldl_impl(f, op, nt::NamedTuple{()}, itr) - y = iterate(itr) - if y === nothing - return Base.mapreduce_empty_iter(f, op, itr, IteratorEltype(itr)) - end - x, i = y - init = mapreduce_first(f, op, x) - return mapfoldl_impl(f, op, (init=init,), itr, i) +struct _InitialValue end + +struct BottomRF{T} + rf::T end +@inline (op::BottomRF)(::_InitialValue, x) = x +@inline (op::BottomRF)(acc, x) = op.rf(acc, x) + +struct MappingRF{F, T} + f::F + rf::T +end + +@inline (op::MappingRF)(acc, x) = op.rf(acc, op.f(x)) + +struct FilteringRF{F, T} + f::F + rf::T +end + +@inline (op::FilteringRF)(acc, x) = op.f(x) ? op.rf(acc, x) : acc + +""" + _xfadjoint(op, itr) -> op′, itr′ + +Given a pair of reducing function `op` and an iterator `itr`, return a pair +`(op′, itr′)` of similar types. If the iterator `itr` is transformed by an +iterator transform `ixf` whose adjoint transducer `xf` is known, `op′ = xf(op)` +and `itr′ = "parent" of itr` is returned. Otherwise, `op` and `itr` are +returned as-is. For example, transducer `rf -> MappingRF(f, rf)` is the +adjoint of iterator transform `itr -> Generator(f, itr)`. + +Nested iterator transforms are converted recursively. That is to say, +given `op` and + + itr = (ixf₁ ∘ ixf₂ ∘ ... ∘ ixfₙ)(itr′) + +what is returned is `itr′` and + + op′ = (xfₙ ∘ ... ∘ xf₂ ∘ xf₁)(op) +""" +_xfadjoint(op, itr) = (op, itr) +_xfadjoint(op, itr::Generator) = + if itr.f === identity + op, itr.iter + else + _xfadjoint(MappingRF(itr.f, op), itr.iter) + end +_xfadjoint(op, itr::Filter) = + _xfadjoint(FilteringRF(itr.flt, op), itr.itr) """ mapfoldl(f, op, itr; [init]) @@ -92,21 +141,14 @@ foldl(op, itr; kw...) = mapfoldl(identity, op, itr; kw...) ## foldr & mapfoldr mapfoldr_impl(f, op, nt::NamedTuple{(:init,)}, itr) = - mapfoldl_impl(f, (x,y) -> op(y,x), nt, Iterators.reverse(itr)) - -# we can't just call mapfoldl_impl with (x,y) -> op(y,x), because -# we need to use the type of op for mapreduce_empty_iter and mapreduce_first. -function mapfoldr_impl(f, op, nt::NamedTuple{()}, itr) - ritr = Iterators.reverse(itr) - y = iterate(ritr) - if y === nothing - return Base.mapreduce_empty_iter(f, op, itr, IteratorEltype(itr)) - end - x, i = y - init = mapreduce_first(f, op, x) - return mapfoldl_impl(f, (x,y) -> op(y,x), (init=init,), ritr, i) + mapfoldl_impl(f, FlipArgs(op), nt, Iterators.reverse(itr)) + +struct FlipArgs{F} + f::F end +@inline (f::FlipArgs)(x, y) = f.f(y, x) + """ mapfoldr(f, op, itr; [init]) @@ -234,6 +276,11 @@ reduce_empty(::typeof(mul_prod), T) = reduce_empty(*, T) reduce_empty(::typeof(mul_prod), ::Type{T}) where {T<:SmallSigned} = one(Int) reduce_empty(::typeof(mul_prod), ::Type{T}) where {T<:SmallUnsigned} = one(UInt) +reduce_empty(op::BottomRF, T) = reduce_empty(op.rf, T) +reduce_empty(op::MappingRF, T) = mapreduce_empty(op.f, op.rf, T) +reduce_empty(op::FilteringRF, T) = reduce_empty(op.rf, T) +reduce_empty(op::FlipArgs, T) = reduce_empty(op.f, T) + """ Base.mapreduce_empty(f, op, T) @@ -251,10 +298,12 @@ mapreduce_empty(::typeof(abs2), op, T) = abs2(reduce_empty(op, T)) mapreduce_empty(f::typeof(abs), ::typeof(max), T) = abs(zero(T)) mapreduce_empty(f::typeof(abs2), ::typeof(max), T) = abs2(zero(T)) -mapreduce_empty_iter(f, op, itr, ::HasEltype) = mapreduce_empty(f, op, eltype(itr)) -mapreduce_empty_iter(f, op::typeof(&), itr, ::EltypeUnknown) = true -mapreduce_empty_iter(f, op::typeof(|), itr, ::EltypeUnknown) = false -mapreduce_empty_iter(f, op, itr, ::EltypeUnknown) = _empty_reduce_error() +# For backward compatibility: +mapreduce_empty_iter(f, op, itr, ItrEltype) = + reduce_empty_iter(MappingRF(f, op), itr, ItrEltype) + +reduce_empty_iter(op, itr, ::HasEltype) = reduce_empty(op, eltype(itr)) +reduce_empty_iter(op, itr, ::EltypeUnknown) = _empty_reduce_error() # handling of single-element iterators """