Merge pull request #547 from JuliaArrays/teh/inferbcast

Fix inference on axes for broadcasting

Author: timholy

src/SOneTo.jl

@@ -58,3 +58,6 @@ end
 Base.@pure function Base.checkindex(::Type{Bool}, ::SOneTo{n1}, ::SOneTo{n2}) where {n1, n2}
     return n1::Int >= n2::Int
 end
+
+Base.promote_rule(a::Type{Base.OneTo{T}}, ::Type{SOneTo{n}}) where {T,n} =
+    Base.OneTo{promote_type(T, Int)}

src/abstractarray.jl

@@ -13,6 +13,8 @@ Base.axes(s::StaticArray) = _axes(Size(s))
 @pure function _axes(::Size{sizes}) where {sizes}
     map(SOneTo, sizes)
 end
+Base.axes(rv::Adjoint{<:Any,<:StaticVector})   = (SOneTo(1), axes(rv.parent)...)
+Base.axes(rv::Transpose{<:Any,<:StaticVector}) = (SOneTo(1), axes(rv.parent)...)
 
 function Base.summary(io::IO, a, inds::Tuple{SOneTo, Vararg{SOneTo}})
     print(io, Base.dims2string(length.(inds)), " ")

src/broadcast.jl

@@ -4,6 +4,8 @@
 
 import Base.Broadcast:
 BroadcastStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle, materialize!
+import Base.Broadcast: _bcs1  # for SOneTo axis information
+using Base.Broadcast: _bcsm
 # Add a new BroadcastStyle for StaticArrays, derived from AbstractArrayStyle
 # A constructor that changes the style parameter N (array dimension) is also required
 struct StaticArrayStyle{N} <: AbstractArrayStyle{N} end
@@ -37,6 +39,10 @@ end
     _broadcast!(f, destsize, dest, argsizes, as...)
 end
 
+# Resolving priority between dynamic and static axes
+_bcs1(a::SOneTo, b::SOneTo) = _bcsm(b, a) ? b : (_bcsm(a, b) ? a : throw(DimensionMismatch("arrays could not be broadcast to a common size")))
+_bcs1(a::SOneTo, b::Base.OneTo) = _bcs1(Base.OneTo(a), b)
+_bcs1(a::Base.OneTo, b::SOneTo) = _bcs1(a, Base.OneTo(b))
 
 ###################################################
 ## Internal broadcast machinery for StaticArrays ##

test/broadcast.jl

@@ -49,6 +49,10 @@ end
         @test @inferred(v .- m) === @SMatrix [0 -1; 1 0]
         @test @inferred(m .^ v) === @SMatrix [1 2; 81 256]
         @test @inferred(v .^ m) === @SMatrix [1 1; 64 256]
+        # Issue #546
+        @test @inferred(m ./ (v .* v')) === @SMatrix [1.0 0.5; 0.75 0.25]
+        testinf(m, v) = m ./ (v .* v')
+        @test @inferred(testinf(m, v)) === @SMatrix [1.0 0.5; 0.75 0.25]
     end
 
     @testset "2x2 StaticMatrix with 1x2 StaticMatrix" begin

test/core.jl

@@ -183,4 +183,11 @@
         @test !StaticArrays.sizematch(Size(2, 2), a)
         @test !StaticArrays.sizematch(Size(3, StaticArrays.Dynamic()), a)
     end
+
+    @testset "SOneTo" begin
+        a = Base.OneTo(3)
+        b = StaticArrays.SOneTo{2}()
+        @test @inferred(promote(a, b)) === (a, Base.OneTo(2))
+        @test @inferred(promote(b, a)) === (Base.OneTo(2), a)
+    end
 end