Skip to content

Commit

Permalink
Fix the return type in broadcast
Browse files Browse the repository at this point in the history
  • Loading branch information
pabloferz committed Dec 12, 2016
1 parent 579cfcb commit 02596f9
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 28 deletions.
8 changes: 4 additions & 4 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1768,10 +1768,10 @@ function mapslices(f, A::AbstractArray, dims::AbstractVector)
end

# These are needed because map(eltype, As) is not inferrable
promote_eltype_op(::Any) = (@_pure_meta; Any)
promote_eltype_op(op, A) = (@_pure_meta; promote_op(op, eltype(A)))
promote_eltype_op(op, A, B) = (@_pure_meta; promote_op(op, eltype(A), eltype(B)))
promote_eltype_op(op, A, B, C, D...) = (@_pure_meta; promote_eltype_op(op, eltype(A), promote_eltype_op(op, B, C, D...)))
promote_eltype_op(::Any) = Any
promote_eltype_op(op, A) = (@_inline_meta; promote_op(op, eltype(A)))
promote_eltype_op(op, A, B) = (@_inline_meta; promote_op(op, eltype(A), eltype(B)))
promote_eltype_op(op, A, B, C, D...) = (@_inline_meta; promote_eltype_op(op, eltype(A), promote_eltype_op(op, B, C, D...)))

## 1 argument

Expand Down
24 changes: 10 additions & 14 deletions base/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
module Broadcast

using Base.Cartesian
using Base: promote_eltype_op, _default_eltype, linearindices, tail, OneTo, to_shape,
using Base: @pure, promote_eltype_op, _promote_op, linearindices, tail, OneTo, to_shape,
_msk_end, unsafe_bitgetindex, bitcache_chunks, bitcache_size, dumpbitcache
import Base: .+, .-, .*, ./, .\, .//, .==, .<, .!=, .<=, , .%, .<<, .>>, .^
import Base: broadcast
Expand Down Expand Up @@ -257,26 +257,22 @@ 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::DataType, A) = Type{T}
ftype(T::DataType, A...) = Type{T}
ziptype(A) = Tuple{eltype(A)}
ziptype(A, B) = Iterators.Zip2{Tuple{eltype(A)}, Tuple{eltype(B)}}
@inline ziptype(A, B, C, D...) = Iterators.Zip{Tuple{eltype(A)}, ziptype(B, C, D...)}
@pure typestuple(a) = Tuple{eltype(a)}
@pure typestuple(T::Type) = Tuple{Type{T}}
@pure typestuple(a, b...) = Tuple{typestuple(a).types..., typestuple(b...).types...}

# broadcast methods that dispatch on the type of the final container
@inline function broadcast_c(f, ::Type{Array}, As...)
T = _default_eltype(Base.Generator{ziptype(As...), ftype(f, As...)})
shape = broadcast_indices(As...)
@inline function broadcast_c(f, ::Type{Array}, A, Bs...)
T = _promote_op(f, typestuple(A, Bs...))
shape = broadcast_indices(A, Bs...)
iter = CartesianRange(shape)
if isleaftype(T)
return broadcast_t(f, T, shape, iter, As...)
return broadcast_t(f, T, shape, iter, A, Bs...)
end
if isempty(iter)
return similar(Array{T}, shape)
end
return broadcast_t(f, Any, shape, iter, As...)
return broadcast_t(f, Any, shape, iter, A, Bs...)
end
function broadcast_c(f, ::Type{Tuple}, As...)
shape = broadcast_indices(As...)
Expand Down Expand Up @@ -351,7 +347,7 @@ julia> string.(("one","two","three","four"), ": ", 1:4)
"four: 4"
```
"""
@inline broadcast(f, As...) = broadcast_c(f, containertype(As...), As...)
@inline broadcast(f, A, Bs...) = broadcast_c(f, containertype(A, Bs...), A, Bs...)

"""
bitbroadcast(f, As...)
Expand Down
12 changes: 8 additions & 4 deletions base/promotion.jl
Original file line number Diff line number Diff line change
Expand Up @@ -222,18 +222,22 @@ minmax(x::Real, y::Real) = minmax(promote(x, y)...)
# operations, so it is advised against overriding them
_default_type(T::Type) = (@_pure_meta; T)

if isdefined(Core, :Inference)
_promote_op(f::ANY, t::ANY) = Core.Inference.return_type(f, t)
else
_promote_op(f::ANY, t::ANY) = Any
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 = _promote_op(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 = _promote_op(f, Tuple{_default_type(R), _default_type(S)})
isleaftype(R) && isleaftype(S) && return isleaftype(T) ? T : Any
return typejoin(R, S, T)
end
Expand Down
15 changes: 9 additions & 6 deletions test/broadcast.jl
Original file line number Diff line number Diff line change
Expand Up @@ -341,14 +341,17 @@ end
immutable StrangeType18623 end
StrangeType18623(x) = x
StrangeType18623(x,y) = (x,y)
@test @inferred broadcast(StrangeType18623, 1:3) == [1,2,3]
@test @inferred broadcast(StrangeType18623, 1:3, 4:6) == [(1,4),(2,5),(3,6)]
@test @inferred(broadcast(StrangeType18623, 1:3)) == [1,2,3]
@test @inferred(broadcast(StrangeType18623, 1:3, 4:6)) == [(1,4),(2,5),(3,6)]

@test typeof(Int.(Number[1, 2, 3])) === typeof((x->Int(x)).(Number[1, 2, 3]))

@test @inferred broadcast(CartesianIndex, 1:2) == [CartesianIndex(1), CartesianIndex(2)]
@test @inferred broadcast(CartesianIndex, 1:2, 3:4) == [CartesianIndex(1,3), CartesianIndex(2,4)]
@test @inferred(broadcast(CartesianIndex, 1:2)) == [CartesianIndex(1), CartesianIndex(2)]
@test @inferred(broadcast(CartesianIndex, 1:2, 3:4)) == [CartesianIndex(1,3), CartesianIndex(2,4)]

# Issue 18622
@test @inferred muladd.([1.0], [2.0], [3.0])::Vector{Float64} == [5.0]
@test @inferred tuple.(1:3, 4:6, 7:9)::Vector{Tuple{Int,Int,Int}} == [(1,4,7), (2,5,8), (3,6,9)]
@test @inferred(broadcast(muladd, [1.0], [2.0], [3.0])) == [5.0]
@test @inferred(broadcast(tuple, 1:3, 4:6, 7:9)) == [(1,4,7), (2,5,8), (3,6,9)]

# 19419
@test @inferred(broadcast(round, Int, [1])) == [1]

0 comments on commit 02596f9

Please sign in to comment.