Skip to content

Commit

Permalink
Simplify broadcast's eltype promotion mechanism and make it handle mo…
Browse files Browse the repository at this point in the history
…re cases.

Re-simplify broadcast's eltype promotion mechanism as in #19421. With benefit of #19667, this simplified mechanism should handle additional cases (e.g. closures accepting more than two arguments). Also rename the mechanism more precisely (_broadcast_type -> _broadcast_eltype).
  • Loading branch information
Sacha0 committed Dec 31, 2016
1 parent 26c8d85 commit c113bea
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 29 deletions.
39 changes: 17 additions & 22 deletions base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -282,32 +282,27 @@ end
@inline broadcast_elwise_op(f, As...) =
broadcast!(f, similar(Array{promote_eltype_op(f, As...)}, broadcast_indices(As...)), As...)

ftype(f, A) = typeof(f)
ftype(f, A...) = typeof(a -> f(a...))
ftype(T::Type, A...) = Type{T}

# nullables need to be treated like scalars sometimes and like containers
# other times, so there are two variants of typestuple.
# _broadcast_eltype is broadcast's primary result-eltype promotion mechanism.
# _broadcast_eltype uses eltypestuple to construct a tuple type of the eltypes
# of the input-array arguments passed to _broadcast_eltype (from an upstream broadcast).
_broadcast_eltype{S}(::Type{S}, f, As...) = Base._return_type(f, eltypestuple(S, As...))
_broadcast_eltype{S}(::Type{S}, f, T::Type, As...) = Base._return_type(f, eltypestuple(S, T, As...)) # 19419 workaround
eltypestuple(::Type, a) = (Base.@_pure_meta; Tuple{eltype(a)})
eltypestuple(::Type, T::Type) = (Base.@_pure_meta; Tuple{Type{T}})
eltypestuple{S}(::Type{S}, a, b...) = (Base.@_pure_meta; Tuple{eltypestuple(S, a).types..., eltypestuple(S, b...).types...})
# nullables need special handling: in some cases they behave like scalars, and in others
# like containers. _broadcast_eltype and eltypestuple presently handles this via its first
# (type) argument, through which callers provide the context in which the nullable
# input-argument appears (and hence how it should be treated). specifically, if the
# first argument is Any, then nullables are treated as scalars, whereas otherwise
# (i.e. if the first argument is not Any), then nullables are treated as containers.
eltypestuple(::Type{Any}, a::Nullable) = (Base.@_pure_meta; Tuple{typeof(a)})

# if the first argument is Any, then Nullable should be treated like a
# scalar; if the first argument is Array, then Nullable should be treated
# like a container.
typestuple(::Type, a) = (Base.@_pure_meta; Tuple{eltype(a)})
typestuple(::Type{Any}, a::Nullable) = (Base.@_pure_meta; Tuple{typeof(a)})
typestuple(::Type, T::Type) = (Base.@_pure_meta; Tuple{Type{T}})
typestuple{T}(::Type{T}, a, b...) = (Base.@_pure_meta; Tuple{typestuple(T, a).types..., typestuple(T, b...).types...})

# these functions take the variant of typestuple to be used as first argument
ziptype{T}(::Type{T}, A) = typestuple(T, A)
ziptype{T}(::Type{T}, A, B) = (Base.@_pure_meta; Iterators.Zip2{typestuple(T, A), typestuple(T, B)})
@inline ziptype{T}(::Type{T}, A, B, C, D...) = Iterators.Zip{typestuple(T, A), ziptype(T, B, C, D...)}

_broadcast_type{S}(::Type{S}, f, T::Type, As...) = Base._return_type(f, typestuple(S, T, As...))
_broadcast_type{T}(::Type{T}, f, A, Bs...) = Base._default_eltype(Base.Generator{ziptype(T, A, Bs...), ftype(f, A, Bs...)})

# broadcast methods that dispatch on the type of the final container
@inline function broadcast_c(f, ::Type{Array}, A, Bs...)
T = _broadcast_type(Any, f, A, Bs...)
T = _broadcast_eltype(Any, f, A, Bs...)
shape = broadcast_indices(A, Bs...)
iter = CartesianRange(shape)
if isleaftype(T)
Expand All @@ -332,7 +327,7 @@ function broadcast_c(f, ::Type{Tuple}, As...)
end
@inline function broadcast_c(f, ::Type{Nullable}, a...)
nonnull = all(hasvalue, a)
S = _broadcast_type(Array, f, a...)
S = _broadcast_eltype(Array, f, a...)
if isleaftype(S) && null_safe_eltype_op(f, a...)
Nullable{S}(f(map(unsafe_get, a)...), nonnull)
else
Expand Down
6 changes: 2 additions & 4 deletions base/promotion.jl
Original file line number Diff line number Diff line change
Expand Up @@ -231,15 +231,13 @@ end
promote_op(::Any...) = (@_pure_meta; Any)
function promote_op{S}(f, ::Type{S})
@_inline_meta
Z = Tuple{_default_type(S)}
T = _default_eltype(Generator{Z, typeof(f)})
T = _return_type(f, Tuple{_default_type(S)})
isleaftype(S) && return isleaftype(T) ? T : Any
return typejoin(S, T)
end
function promote_op{R,S}(f, ::Type{R}, ::Type{S})
@_inline_meta
Z = Iterators.Zip2{Tuple{_default_type(R)}, Tuple{_default_type(S)}}
T = _default_eltype(Generator{Z, typeof(a -> f(a...))})
T = _return_type(f, Tuple{_default_type(R)}, Tuple{_default_type(S)})
isleaftype(R) && isleaftype(S) && return isleaftype(T) ? T : Any
return typejoin(R, S, T)
end
Expand Down
4 changes: 2 additions & 2 deletions base/sparse/higherorderfns.jl
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ function _noshapecheck_map{Tf,N}(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseVecO
fofzeros = f(_zeros_eltypes(A, Bs...)...)
fpreszeros = fofzeros == zero(fofzeros)
maxnnzC = fpreszeros ? min(length(A), _sumnnzs(A, Bs...)) : length(A)
entrytypeC = Base.Broadcast._broadcast_type(Any, f, A, Bs...)
entrytypeC = Base.Broadcast._broadcast_eltype(Any, f, A, Bs...)
indextypeC = _promote_indtype(A, Bs...)
C = _allocres(size(A), indextypeC, entrytypeC, maxnnzC)
return fpreszeros ? _map_zeropres!(f, C, A, Bs...) :
Expand Down Expand Up @@ -101,7 +101,7 @@ function _diffshape_broadcast{Tf,N}(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseV
fofzeros = f(_zeros_eltypes(A, Bs...)...)
fpreszeros = fofzeros == zero(fofzeros)
indextypeC = _promote_indtype(A, Bs...)
entrytypeC = Base.Broadcast._broadcast_type(Any, f, A, Bs...)
entrytypeC = Base.Broadcast._broadcast_eltype(Any, f, A, Bs...)
shapeC = to_shape(Base.Broadcast.broadcast_indices(A, Bs...))
maxnnzC = fpreszeros ? _checked_maxnnzbcres(shapeC, A, Bs...) : _densennz(shapeC)
C = _allocres(shapeC, indextypeC, entrytypeC, maxnnzC)
Expand Down
8 changes: 7 additions & 1 deletion test/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -363,7 +363,7 @@ StrangeType18623(x,y) = (x,y)
let
f(A, n) = broadcast(x -> +(x, n), A)
@test @inferred(f([1.0], 1)) == [2.0]
g() = (a = 1; Base.Broadcast._broadcast_type(Any, x -> x + a, 1.0))
g() = (a = 1; Base.Broadcast._broadcast_eltype(Any, x -> x + a, 1.0))
@test @inferred(g()) === Float64
end

Expand Down Expand Up @@ -409,3 +409,9 @@ Base.Broadcast.broadcast_c(f, ::Type{Array19745}, A, Bs...) =
@test isa(aa .+ 1, Array19745)
@test isa(aa .* aa', Array19745)
end

# Test that broadcast's promotion mechanism handles closures accepting more than one argument.
# (See issue #19641 and referenced issues and pull requests.)
let f() = (a = 1; Base.Broadcast._broadcast_eltype(Any, (x, y) -> x + y + a, 1.0, 1.0))
@test @inferred(f()) == Float64
end

0 comments on commit c113bea

Please sign in to comment.