Nested special broadcast, throw error when special broadcast fails (#114)

* Use convert for ones broadcasted with matrices

* Flag dodgy broadcast

* Zeros broadcasting for more complicated broadcast

* Update fillbroadcast.jl

* increase coverage

* Update runtests.jl

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "FillArrays"
 uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
-version = "0.9.2"
+version = "0.9.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/fillbroadcast.jl

@@ -27,12 +27,11 @@ function broadcasted(::DefaultArrayStyle, op, a::AbstractFill, b::AbstractFill)
     return Fill(val, broadcast_shape(axes(a), axes(b)))
 end
 
-function _broadcasted_zeros(a, b)
-    return Zeros{promote_type(eltype(a), eltype(b))}(broadcast_shape(axes(a), axes(b)))
-end
-function _broadcasted_ones(a, b)
-    return Ones{promote_type(eltype(a), eltype(b))}(broadcast_shape(axes(a), axes(b)))
-end
+_broadcasted_eltype(a) = eltype(a)
+_broadcasted_eltype(a::Base.Broadcast.Broadcasted) = Base.Broadcast.combine_eltypes(a.f, a.args)
+
+_broadcasted_zeros(a, b) = Zeros{promote_type(_broadcasted_eltype(a), _broadcasted_eltype(b))}(broadcast_shape(axes(a), axes(b)))
+_broadcasted_ones(a, b) = Ones{promote_type(_broadcasted_eltype(a), _broadcasted_eltype(b))}(broadcast_shape(axes(a), axes(b)))
 
 broadcasted(::DefaultArrayStyle, ::typeof(+), a::Zeros, b::Zeros) = _broadcasted_zeros(a, b)
 broadcasted(::DefaultArrayStyle, ::typeof(+), a::Ones, b::Zeros) = _broadcasted_ones(a, b)
@@ -51,6 +50,7 @@ for op in (:*, :/)
         broadcasted(::DefaultArrayStyle, ::typeof($op), a::Zeros, b::Number) = _broadcasted_zeros(a, b)
         broadcasted(::DefaultArrayStyle, ::typeof($op), a::Zeros, b::AbstractRange) = _broadcasted_zeros(a, b)
         broadcasted(::DefaultArrayStyle, ::typeof($op), a::Zeros, b::AbstractArray) = _broadcasted_zeros(a, b)
+        broadcasted(::DefaultArrayStyle, ::typeof($op), a::Zeros, b::Base.Broadcast.Broadcasted) = _broadcasted_zeros(a, b)
         broadcasted(::DefaultArrayStyle{1}, ::typeof($op), a::Zeros, b::AbstractRange) = _broadcasted_zeros(a, b)
     end
 end
@@ -62,6 +62,7 @@ for op in (:*, :\)
         broadcasted(::DefaultArrayStyle, ::typeof($op), a::Number, b::Zeros) = _broadcasted_zeros(a, b)
         broadcasted(::DefaultArrayStyle, ::typeof($op), a::AbstractRange, b::Zeros) = _broadcasted_zeros(a, b)
         broadcasted(::DefaultArrayStyle, ::typeof($op), a::AbstractArray, b::Zeros) = _broadcasted_zeros(a, b)
+        broadcasted(::DefaultArrayStyle, ::typeof($op), a::Base.Broadcast.Broadcasted, b::Zeros) = _broadcasted_zeros(a, b)
         broadcasted(::DefaultArrayStyle{1}, ::typeof($op), a::AbstractRange, b::Zeros) = _broadcasted_zeros(a, b)
     end
 end
@@ -76,24 +77,48 @@ _range_convert(::Type{AbstractVector{T}}, a::AbstractRange{T}) where T = a
 _range_convert(::Type{AbstractVector{T}}, a::AbstractUnitRange) where T = convert(T,first(a)):convert(T,last(a))
 _range_convert(::Type{AbstractVector{T}}, a::AbstractRange) where T = convert(T,first(a)):step(a):convert(T,last(a))
 
+
+# TODO: replacing with the following will support more general broadcasting.
+# function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractFill, b::AbstractRange)
+#     broadcast_shape(axes(a), axes(b)) # check axes
+#     r1 = b[1] * getindex_value(a)
+#     T = typeof(r1)
+#     if length(b) == 1 # Need a fill, but for type stability use StepRangeLen
+#         StepRangeLen{T}(r1, zero(T), length(a))
+#     else
+#         StepRangeLen{T}(r1, convert(T, getindex_value(a) * step(b)), length(b))
+#     end
+# end
+
+# function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractRange, b::AbstractFill)
+#     broadcast_shape(axes(a), axes(b)) # check axes
+#     r1 = a[1] * getindex_value(b)
+#     T = typeof(r1)
+#     if length(a) == 1 # Need a fill, but for type stability use StepRangeLen
+#         StepRangeLen{T}(r1, zero(T), length(b))
+#     else
+#         StepRangeLen{T}(r1, convert(T, step(a) * getindex_value(b)), length(a))
+#     end
+# end
+
 function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::Ones{T}, b::AbstractRange{V}) where {T,V}
-    broadcast_shape(axes(a), axes(b)) # Check sizes are compatible.
+    broadcast_shape(axes(a), axes(b)) == axes(b) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
     return _range_convert(AbstractVector{promote_type(T,V)}, b)
 end
 
 function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractRange{V}, b::Ones{T}) where {T,V}
-    broadcast_shape(axes(a), axes(b)) # Check sizes are compatible.
+    broadcast_shape(axes(a), axes(b)) == axes(a) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
     return _range_convert(AbstractVector{promote_type(T,V)}, a)
 end
 
 
 function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractFill, b::AbstractRange)
-    broadcast_shape(axes(a), axes(b)) # Check sizes are compatible.
+    broadcast_shape(axes(a), axes(b)) == axes(b) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
     return broadcasted(*, getindex_value(a), b)
 end
 
 function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractRange, b::AbstractFill)
-    broadcast_shape(axes(a), axes(b)) # Check sizes are compatible.
+    broadcast_shape(axes(a), axes(b)) == axes(a) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
     return broadcasted(*, a, getindex_value(b))
 end
 

test/runtests.jl

@@ -537,13 +537,22 @@ end
     @test imag(Ones{ComplexF64}(10)) isa Zeros{Float64}
     @test imag(Ones{ComplexF64}(10,10)) isa Zeros{Float64}
 
-    rnge = range(-5.0, step=1.0, length=10)
-    @test broadcast(*, Fill(5.0, 10), rnge) == broadcast(*, 5.0, rnge)
-    @test broadcast(*, Zeros(10, 10), rnge) == zeros(10, 10)
-    @test broadcast(*, rnge, Zeros(10, 10)) == zeros(10, 10)
-    @test_throws DimensionMismatch broadcast(*, Fill(5.0, 11), rnge)
-    @test broadcast(*, rnge, Fill(5.0, 10)) == broadcast(*, rnge, 5.0)
-    @test_throws DimensionMismatch broadcast(*, rnge, Fill(5.0, 11))
+    @testset "range broadcast" begin
+        rnge = range(-5.0, step=1.0, length=10)
+        @test broadcast(*, Fill(5.0, 10), rnge) == broadcast(*, 5.0, rnge)
+        @test broadcast(*, Zeros(10, 10), rnge) == zeros(10, 10)
+        @test broadcast(*, rnge, Zeros(10, 10)) == zeros(10, 10)
+        @test broadcast(*, Ones{Int}(10), rnge) ≡ rnge
+        @test broadcast(*, rnge, Ones{Int}(10)) ≡ rnge
+        @test_throws DimensionMismatch broadcast(*, Fill(5.0, 11), rnge)
+        @test broadcast(*, rnge, Fill(5.0, 10)) == broadcast(*, rnge, 5.0)
+        @test_throws DimensionMismatch broadcast(*, rnge, Fill(5.0, 11))
+
+        # following should pass using alternative implementation in code
+        deg = 5:5
+        @test_throws ArgumentError @inferred(broadcast(*, Fill(5.0, 10), deg)) == broadcast(*, fill(5.0,10), deg)
+        @test_throws ArgumentError @inferred(broadcast(*, deg, Fill(5.0, 10))) == broadcast(*, deg, fill(5.0,10))
+    end
 
     @testset "Special Zeros/Ones" begin
         @test broadcast(+,Zeros(5)) ≡ broadcast(-,Zeros(5)) ≡ Zeros(5)
@@ -566,6 +575,11 @@ end
         @test_throws DimensionMismatch broadcast(*, 1:6, Ones(3))
         @test_throws DimensionMismatch broadcast(*, Fill(1,3), 1:6)
         @test_throws DimensionMismatch broadcast(*, 1:6, Fill(1,3))
+
+        @testset "Nested" begin
+            @test randn(5) .\ rand(5) .* Zeros(5) ≡ Zeros(5)
+            @test broadcast(*, Zeros(5), Base.Broadcast.broadcasted(\, randn(5), rand(5))) ≡ Zeros(5)
+        end
     end
 
     @testset "support Ref" begin