Skip to content

Commit 686ca71

Browse files
committed
Simplify broadcast's eltype promotion mechanism and make it handle more 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).
1 parent ab984a5 commit 686ca71

File tree

5 files changed

+29
-30
lines changed

5 files changed

+29
-30
lines changed

base/broadcast.jl

Lines changed: 17 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -282,32 +282,27 @@ end
282282
@inline broadcast_elwise_op(f, As...) =
283283
broadcast!(f, similar(Array{promote_eltype_op(f, As...)}, broadcast_indices(As...)), As...)
284284

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

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

292-
# if the first argument is Any, then Nullable should be treated like a
293-
# scalar; if the first argument is Array, then Nullable should be treated
294-
# like a container.
295-
typestuple(::Type, a) = (Base.@_pure_meta; Tuple{eltype(a)})
296-
typestuple(::Type{Any}, a::Nullable) = (Base.@_pure_meta; Tuple{typeof(a)})
297-
typestuple(::Type, T::Type) = (Base.@_pure_meta; Tuple{Type{T}})
298-
typestuple{T}(::Type{T}, a, b...) = (Base.@_pure_meta; Tuple{typestuple(T, a).types..., typestuple(T, b...).types...})
299-
300-
# these functions take the variant of typestuple to be used as first argument
301-
ziptype{T}(::Type{T}, A) = typestuple(T, A)
302-
ziptype{T}(::Type{T}, A, B) = (Base.@_pure_meta; Iterators.Zip2{typestuple(T, A), typestuple(T, B)})
303-
@inline ziptype{T}(::Type{T}, A, B, C, D...) = Iterators.Zip{typestuple(T, A), ziptype(T, B, C, D...)}
304-
305-
_broadcast_type{S}(::Type{S}, f, T::Type, As...) = Base._return_type(f, typestuple(S, T, As...))
306-
_broadcast_type{T}(::Type{T}, f, A, Bs...) = Base._default_eltype(Base.Generator{ziptype(T, A, Bs...), ftype(f, A, Bs...)})
307302

308303
# broadcast methods that dispatch on the type of the final container
309304
@inline function broadcast_c(f, ::Type{Array}, A, Bs...)
310-
T = _broadcast_type(Any, f, A, Bs...)
305+
T = _broadcast_eltype(Any, f, A, Bs...)
311306
shape = broadcast_indices(A, Bs...)
312307
iter = CartesianRange(shape)
313308
if isleaftype(T)
@@ -332,7 +327,7 @@ function broadcast_c(f, ::Type{Tuple}, As...)
332327
end
333328
@inline function broadcast_c(f, ::Type{Nullable}, a...)
334329
nonnull = all(hasvalue, a)
335-
S = _broadcast_type(Array, f, a...)
330+
S = _broadcast_eltype(Array, f, a...)
336331
if isleaftype(S) && null_safe_eltype_op(f, a...)
337332
Nullable{S}(f(map(unsafe_get, a)...), nonnull)
338333
else

base/promotion.jl

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -231,15 +231,13 @@ end
231231
promote_op(::Any...) = (@_pure_meta; Any)
232232
function promote_op{S}(f, ::Type{S})
233233
@_inline_meta
234-
Z = Tuple{_default_type(S)}
235-
T = _default_eltype(Generator{Z, typeof(f)})
234+
T = _return_type(f, Tuple{_default_type(S)})
236235
isleaftype(S) && return isleaftype(T) ? T : Any
237236
return typejoin(S, T)
238237
end
239238
function promote_op{R,S}(f, ::Type{R}, ::Type{S})
240239
@_inline_meta
241-
Z = Iterators.Zip2{Tuple{_default_type(R)}, Tuple{_default_type(S)}}
242-
T = _default_eltype(Generator{Z, typeof(a -> f(a...))})
240+
T = _return_type(f, Tuple{_default_type(R)}, Tuple{_default_type(S)})
243241
isleaftype(R) && isleaftype(S) && return isleaftype(T) ? T : Any
244242
return typejoin(R, S, T)
245243
end

base/sparse/sparse.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ import Base: @get!, acos, acosd, acot, acotd, acsch, asech, asin, asind, asinh,
2626
rotl90, rotr90, round, scale!, setindex!, similar, size, transpose, tril,
2727
triu, vec, permute!, map, map!
2828

29-
import Base.Broadcast: _broadcast_type, broadcast_indices
29+
import Base.Broadcast: _broadcast_eltype, broadcast_indices
3030

3131
export AbstractSparseArray, AbstractSparseMatrix, AbstractSparseVector,
3232
SparseMatrixCSC, SparseVector, blkdiag, dense, droptol!, dropzeros!, dropzeros,

base/sparse/sparsematrix.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1413,7 +1413,7 @@ function map{Tf,N}(f::Tf, A::SparseMatrixCSC, Bs::Vararg{SparseMatrixCSC,N})
14131413
fofzeros = f(_zeros_eltypes(A, Bs...)...)
14141414
fpreszeros = fofzeros == zero(fofzeros)
14151415
maxnnzC = fpreszeros ? min(length(A), _sumnnzs(A, Bs...)) : length(A)
1416-
entrytypeC = _broadcast_type(Any, f, A, Bs...)
1416+
entrytypeC = _broadcast_eltype(Any, f, A, Bs...)
14171417
indextypeC = _promote_indtype(A, Bs...)
14181418
Ccolptr = Vector{indextypeC}(A.n + 1)
14191419
Crowval = Vector{indextypeC}(maxnnzC)
@@ -1443,7 +1443,7 @@ function broadcast{Tf,N}(f::Tf, A::SparseMatrixCSC, Bs::Vararg{SparseMatrixCSC,N
14431443
fofzeros = f(_zeros_eltypes(A, Bs...)...)
14441444
fpreszeros = fofzeros == zero(fofzeros)
14451445
indextypeC = _promote_indtype(A, Bs...)
1446-
entrytypeC = _broadcast_type(Any, f, A, Bs...)
1446+
entrytypeC = _broadcast_eltype(Any, f, A, Bs...)
14471447
Cm, Cn = Base.to_shape(Base.Broadcast.broadcast_indices(A, Bs...))
14481448
maxnnzC = fpreszeros ? _checked_maxnnzbcres(Cm, Cn, A, Bs...) : (Cm * Cn)
14491449
Ccolptr = Vector{indextypeC}(Cn + 1)

test/broadcast.jl

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -363,7 +363,7 @@ StrangeType18623(x,y) = (x,y)
363363
let
364364
f(A, n) = broadcast(x -> +(x, n), A)
365365
@test @inferred(f([1.0], 1)) == [2.0]
366-
g() = (a = 1; Base.Broadcast._broadcast_type(Any, x -> x + a, 1.0))
366+
g() = (a = 1; Base.Broadcast._broadcast_eltype(Any, x -> x + a, 1.0))
367367
@test @inferred(g()) === Float64
368368
end
369369

@@ -409,3 +409,9 @@ Base.Broadcast.broadcast_c(f, ::Type{Array19745}, A, Bs...) =
409409
@test isa(aa .+ 1, Array19745)
410410
@test isa(aa .* aa', Array19745)
411411
end
412+
413+
# Test that broadcast's promotion mechanism handles closures accepting more than one argument.
414+
# (See issue #19641 and referenced issues and pull requests.)
415+
let f() = (a = 1; Base.Broadcast._broadcast_eltype(Any, (x, y) -> x + y + a, 1.0, 1.0))
416+
@test @inferred(f()) == Float64
417+
end

0 commit comments

Comments
 (0)