Merge pull request #12 from tkoolen/broadcast

tkoolen · web-flow · commit 84ba361518d5 · 2017-09-20T00:49:31.000-04:00
Add `broadcast!` with `AbstractVector` destination
diff --git a/.travis.yml b/.travis.yml
@@ -15,11 +15,9 @@ notifications:
 git:
   depth: 99999999
 
-## uncomment the following lines to allow failures on nightly julia
-## (tests will run but not make your overall status red)
-#matrix:
-#  allow_failures:
-#  - julia: nightly
+matrix:
+  allow_failures:
+    - julia: nightly
 
 ## uncomment and modify the following lines to manually install system packages
 #addons:
diff --git a/README.md b/README.md
@@ -3,7 +3,7 @@
 [![Build Status](https://travis-ci.org/tkoolen/TypeSortedCollections.jl.svg?branch=master)](https://travis-ci.org/tkoolen/TypeSortedCollections.jl)
 [![codecov.io](http://codecov.io/github/tkoolen/TypeSortedCollections.jl/coverage.svg?branch=master)](http://codecov.io/github/tkoolen/TypeSortedCollections.jl?branch=master)
 
-TypeSortedCollections provides the `TypeSortedCollection` type, which can be used to store type-heterogeneous data in a way that allows operations on the data to be performed in a type-stable manner. It does so by sorting a type-heterogeneous input collection by type upon construction, and storing these elements in a `Tuple` of concretely typed `Vector`s, one for each type. TypeSortedCollections provides type stable methods for `map!`, `foreach`, and `mapreduce` that take at least one `TypeSortedCollection`.
+TypeSortedCollections provides the `TypeSortedCollection` type, which can be used to store type-heterogeneous data in a way that allows operations on the data to be performed in a type-stable manner. It does so by sorting a type-heterogeneous input collection by type upon construction, and storing these elements in a `Tuple` of concretely typed `Vector`s, one for each type. TypeSortedCollections provides type stable methods for `map!`, `foreach`, `broadcast!`, and `mapreduce` that take at least one `TypeSortedCollection`.
 
 An example:
 ```julia
@@ -40,7 +40,7 @@ Note that construction of a `TypeSortedCollection` is of course not type stable,
 See also [FunctionWrappers.jl](https://github.com/yuyichao/FunctionWrappers.jl) for a solution to the related problem of storing and calling multiple callables in a type-stable manner, and [Unrolled.jl](https://github.com/cstjean/Unrolled.jl) for a macro-based solution.
 
 # Iteration order
-By default, `TypeSortedCollection`s do not preserve iteration order, in the sense that the order in which elements are processed in `map!`, `foreach`, and `mapreduce` will not be the same as if these functions were called on the original type-heterogeneous vector:
+By default, `TypeSortedCollection`s do not preserve iteration order, in the sense that the order in which elements are processed in `map!`, `foreach`, `broadcast!`, and `mapreduce` will not be the same as if these functions were called on the original type-heterogeneous vector:
 ```julia
 julia> xs = Number[1.; 2; 3.];
 
@@ -75,7 +75,7 @@ julia> results = similar(xs);
 julia> map!(identity, results, sortedxs) # results of applying `identity` end up in the right location
 3-element Array{Number,1}:
  1.0
- 2  
+ 2
  3.0
 ```
 
@@ -105,3 +105,24 @@ julia> map!(*, results, sortedxs, sortedys)
   9.0
  16.0
 ```
+
+# Broadcasting
+Broadcasting (in place) is implemented for `AbstractVector` return types:
+```julia
+julia> x = 4;
+
+julia> ys = Number[1.; 2; 3];
+
+julia> sortedys = TypeSortedCollection(ys);
+
+julia> results = similar(ys, Float64);
+
+julia> results .= x .* sortedys
+3-element Array{Float64,1}:
+  4.0
+  8.0
+ 12.0
+
+julia> @allocated results .= x .* sortedys
+0
+```
diff --git a/src/TypeSortedCollections.jl b/src/TypeSortedCollections.jl
@@ -91,28 +91,35 @@ end
 Base.@pure num_types(::Type{<:TypeSortedCollection{<:Any, N}}) where {N} = N
 num_types(x::TypeSortedCollection) = num_types(typeof(x))
 
+const TSCOrAbstractVector{N} = Union{<:TypeSortedCollection{<:Any, N}, AbstractVector}
+
 Base.isempty(x::TypeSortedCollection) = all(isempty, x.data)
 Base.empty!(x::TypeSortedCollection) = foreach(empty!, x.data)
 Base.length(x::TypeSortedCollection) = mapreduce(length, +, 0, x.data)
 Base.indices(x::TypeSortedCollection) = x.indices # semantics are a little different from Array, but OK
 
 # Trick from StaticArrays:
-@inline first_tsc(a1::TypeSortedCollection, as::Union{<:TypeSortedCollection, AbstractVector}...) = a1
-@inline first_tsc(a1, as::Union{<:TypeSortedCollection, AbstractVector}...) = first_tsc(as...)
+@inline first_tsc(a1::TypeSortedCollection, as...) = a1
+@inline first_tsc(a1, as...) = first_tsc(as...)
+
+Base.@pure first_tsc_type(a1::Type{<:TypeSortedCollection}, as::Type...) = a1
+Base.@pure first_tsc_type(a1::Type, as::Type...) = first_tsc_type(as...)
 
 # inspired by Base.ith_all
 @inline _getindex_all(::Val, j, vecindex) = ()
 Base.@propagate_inbounds _getindex_all(vali::Val{i}, j, vecindex, a1, as...) where {i} = (_getindex(vali, j, vecindex, a1), _getindex_all(vali, j, vecindex, as...)...)
+@inline _getindex(::Val, j, vecindex, a) = a # for anything that's not an AbstractVector or TypeSortedCollection, don't index (for use in broadcast!)
 @inline _getindex(::Val, j, vecindex, a::AbstractVector) = a[vecindex]
 @inline _getindex(::Val{i}, j, vecindex, a::TypeSortedCollection) where {i} = a.data[i][j]
 @inline _setindex!(::Val, j, vecindex, a::AbstractVector, val) = a[vecindex] = val
 @inline _setindex!(::Val{i}, j, vecindex, a::TypeSortedCollection, val) where {i} = a.data[i][j] = val
 
 @inline lengths_match(a1) = true
-@inline lengths_match(a1, a2, as...) = length(a1) == length(a2) && lengths_match(a2, as...)
+@inline lengths_match(a1::TSCOrAbstractVector, a2::TSCOrAbstractVector, as...) = length(a1) == length(a2) && lengths_match(a2, as...)
+@inline lengths_match(a1::TSCOrAbstractVector, a2, as...) = lengths_match(a1, as...) # case: a2 is not indexable: skip it
 @noinline lengths_match_fail() = throw(DimensionMismatch("Lengths of input collections do not match."))
 
-@inline indices_match(::Val, indices::Vector{Int}, ::AbstractVector) = true
+@inline indices_match(::Val, indices::Vector{Int}, ::Any) = true
 @inline function indices_match(::Val{i}, indices::Vector{Int}, tsc::TypeSortedCollection) where {i}
     tsc_indices = tsc.indices[i]
     length(indices) == length(tsc_indices) || return false
@@ -124,7 +131,7 @@ end
 @inline indices_match(vali::Val, indices::Vector{Int}, a1, as...) = indices_match(vali, indices, a1) && indices_match(vali, indices, as...)
 @noinline indices_match_fail() = throw(ArgumentError("Indices of TypeSortedCollections do not match."))
 
-@generated function Base.map!(f, dest::Union{TypeSortedCollection{<:Any, N}, AbstractArray}, args::Union{TypeSortedCollection{<:Any, N}, AbstractArray}...) where {N}
+@generated function Base.map!(f, dest::TSCOrAbstractVector{N}, args::TSCOrAbstractVector{N}...) where {N}
     expr = Expr(:block)
     push!(expr.args, :(Base.@_inline_meta))
     push!(expr.args, :(leading_tsc = first_tsc(dest, args...)))
@@ -134,7 +141,7 @@ end
         push!(expr.args, quote
             let inds = leading_tsc.indices[$i]
                 @boundscheck indices_match($vali, inds, dest, args...) || indices_match_fail()
-                for j in linearindices(inds)
+                @inbounds for j in linearindices(inds)
                     vecindex = inds[j]
                     _setindex!($vali, j, vecindex, dest, f(_getindex_all($vali, j, vecindex, args...)...))
                 end
@@ -147,7 +154,7 @@ end
     end
 end
 
-@generated function Base.foreach(f, As::Union{<:TypeSortedCollection{<:Any, N}, AbstractVector}...) where {N}
+@generated function Base.foreach(f, As::TSCOrAbstractVector{N}...) where {N}
     expr = Expr(:block)
     push!(expr.args, :(Base.@_inline_meta))
     push!(expr.args, :(leading_tsc = first_tsc(As...)))
@@ -187,4 +194,36 @@ end
     end
 end
 
+## broadcast!
+Base.Broadcast._containertype(::Type{<:TypeSortedCollection}) = TypeSortedCollection
+Base.Broadcast.promote_containertype(::Type{TypeSortedCollection}, _) = TypeSortedCollection
+Base.Broadcast.promote_containertype(_, ::Type{TypeSortedCollection}) = TypeSortedCollection
+Base.Broadcast.promote_containertype(::Type{TypeSortedCollection}, ::Type{Array}) = TypeSortedCollection # handle ambiguities with `Array`
+Base.Broadcast.promote_containertype(::Type{Array}, ::Type{TypeSortedCollection}) = TypeSortedCollection # handle ambiguities with `Array`
+
+@generated function Base.Broadcast.broadcast_c!(f, ::Type, ::Type{TypeSortedCollection}, dest::AbstractVector, A, Bs...)
+    T = first_tsc_type(A, Bs...)
+    N = num_types(T)
+    expr = Expr(:block)
+    push!(expr.args, :(Base.@_inline_meta)) # TODO: good idea?
+    push!(expr.args, :(leading_tsc = first_tsc(A, Bs...)))
+    push!(expr.args, :(@boundscheck lengths_match(dest, A, Bs...) || lengths_match_fail()))
+    for i = 1 : N
+        vali = Val(i)
+        push!(expr.args, quote
+            let inds = leading_tsc.indices[$i]
+                @boundscheck indices_match($vali, inds, A, Bs...) || indices_match_fail()
+                @inbounds for j in linearindices(inds)
+                    vecindex = inds[j]
+                    _setindex!($vali, j, vecindex, dest, f(_getindex_all($vali, j, vecindex, A, Bs...)...))
+                end
+            end
+        end)
+    end
+    quote
+        $expr
+        dest
+    end
+end
+
 end # module
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -9,6 +9,7 @@ g(x::Int64, y1::Float64, y2::Int64) = x * y1 * y2
 g(x::Float64, y1::Float64, y2::Int64) = x + y1 + y2
 g(x::Int64, y1::Float64, y2::Float64) = x * y1 - y2
 g(x::Float64, y1::Float64, y2::Float64) = x + y1 - y2
+g(x::Int64, y1::Int64, y2::Float64) = x - y1 * y2
 end
 
 @testset "ambiguities" begin
@@ -148,3 +149,92 @@ end
     foreach(x -> push!(results, x), sortedx2)
     @test all(x .== results)
 end
+
+@testset "broadcast!" begin
+    x = Number[3.; 4; 5]
+    sortedx = TypeSortedCollection(x)
+    y1 = rand(length(x))
+    y2 = rand(Int)
+    results = similar(x, Float64)
+    broadcast!(M.g, results, sortedx, y1, y2)
+    @test all(results .== M.g.(x, y1, y2))
+
+    results = similar(x, Float64)
+    results .= M.g.(sortedx, y1, y2)
+    @test all(results .== M.g.(x, y1, y2))
+
+    @test (@allocated broadcast!(M.g, results, sortedx, y1, y2)) == 0
+end
+
+@testset "broadcast! with scalars and TSC as second arg" begin
+    x = 3
+    y = Number[3.; 4; 5]
+    z = 5.
+    sortedy = TypeSortedCollection(y)
+    results = similar(y, Float64)
+    results .= M.g.(x, sortedy, z)
+    @test all(results .== M.g.(x, y, z))
+    @test (@allocated results .= M.g.(x, sortedy, z)) == 0
+end
+
+@testset "broadcast! consecutive scalars" begin
+    x = 3
+    y = 4.
+    z = Number[3.; 4; 5.]
+    sortedz = TypeSortedCollection(z)
+    results = similar(z, Float64)
+    results .= M.g.(x, y, sortedz)
+    @test all(results .== M.g.(x, y, z))
+    @test (@allocated results .= M.g.(x, y, sortedz)) == 0
+end
+
+@testset "broadcast! Array first" begin
+    x = rand(Int, 3)
+    y = Number[3.; 4; 5.]
+    z = rand()
+    sortedy = TypeSortedCollection(y)
+    results = similar(y, Float64)
+    results .= M.g.(x, sortedy, z)
+    @test all(results .== M.g.(x, y, z))
+    @test (@allocated results .= M.g.(x, sortedy, z)) == 0
+end
+
+@testset "broadcast! indices mismatch" begin
+    x = Number[3.; 4; 5]
+    sortedx = TypeSortedCollection(x)
+    y1 = rand()
+    y2 = Number[8; 9; Float32(7)]
+    sortedy2 = TypeSortedCollection(y2)
+    results = similar(x, Float64)
+    @test_throws ArgumentError broadcast!(M.g, results, sortedx, y1, sortedy2)
+end
+
+@testset "broadcast! length mismatch" begin
+    x = Number[3.; 4; 5]
+    sortedx = TypeSortedCollection(x)
+    y1 = rand(length(x) + 1)
+    y2 = rand(length(x))
+    results = similar(x, Float64)
+    @test_throws DimensionMismatch results .= M.g.(sortedx, y1, y2)
+
+    y1 = rand()
+    y2 = rand(length(x) + 1)
+    @test_throws DimensionMismatch results .= M.g.(sortedx, y1, y2)
+
+    results = rand(length(x) + 1)
+    y1 = rand()
+    y2 = rand(Int)
+    @test_throws DimensionMismatch results .= M.g.(sortedx, y1, y2)
+end
+
+@testset "broadcast! matching indices" begin
+    x = Number[3.; 4; 5]
+    sortedx = TypeSortedCollection(x)
+    y1 = rand()
+    y2 = [7.; 8.; 9.]
+    sortedy2 = TypeSortedCollection(y2, indices(sortedx))
+    results = similar(x, Float64)
+    broadcast!(M.g, results, sortedx, y1, sortedy2)
+    @test all(results .== M.g.(x, y1, y2))
+    @test (@allocated broadcast!(M.g, results, sortedx, y1, sortedy2)) == 0
+end
