From 35638c7309e76e0b92c2f142fd9a7695f9dbf567 Mon Sep 17 00:00:00 2001
From: Twan Koolen <koolen.twan@gmail.com>
Date: Wed, 13 Dec 2017 14:20:40 -0500
Subject: [PATCH] [ci skip]

Address comments.

Reimplement and generalize all-scalar optimization.

Fix allocation tests for sparse broadcast!.
---
 base/broadcast.jl             | 37 ++++++++++++++++++++++++-----------
 base/sparse/higherorderfns.jl | 15 +++++++++-----
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/base/broadcast.jl b/base/broadcast.jl
index 2b9abe1951223..8576824ccaeca 100644
--- a/base/broadcast.jl
+++ b/base/broadcast.jl
@@ -442,21 +442,36 @@ Note that `dest` is only used to store the result, and does not supply
 arguments to `f` unless it is also listed in the `As`,
 as in `broadcast!(f, A, A, B)` to perform `A[:] = broadcast(f, A, B)`.
 """
-broadcast!(f, dest, As...) = broadcast!(f, dest, combine_styles(As...), As...)
-broadcast!(f, dest, ::BroadcastStyle, As...) = broadcast!(f, dest, nothing, As...)
-@inline function broadcast!(f, C, ::Void, A, Bs::Vararg{Any,N}) where N
-    if isa(f, typeof(identity)) && N == 0
-        if isa(A, Number)
-            return fill!(C, A)
-        elseif isa(C, AbstractArray) && isa(A, AbstractArray) && Base.axes(C) == Base.axes(A)
-            return copy!(C, A)
+@inline broadcast!(f, dest, As...) = broadcast!(f, dest, combine_styles(As...), As...)
+@inline broadcast!(f, dest, ::BroadcastStyle, As...) = broadcast!(f, dest, nothing, As...)
+
+# Default behavior (separated out so that it can be called by users who want to extend broadcast!).
+@inline function broadcast!(f, dest, ::Void, As::Vararg{Any, N}) where N
+    if f isa typeof(identity) && N == 1
+        A = As[1]
+        if A isa AbstractArray && Base.axes(dest) == Base.axes(A)
+            return copy!(dest, A)
         end
     end
-    return _broadcast!(f, C, A, Bs...)
+    return _broadcast!(f, dest, As...)
 end
 
-# This indirection allows size-dependent implementations (e.g., see the copying `identity`
-# specialization above)
+# Optimization for the all-Scalar case.
+@inline function broadcast!(f, dest, ::Scalar, As::Vararg{Any, N}) where N
+    if dest isa AbstractArray
+        if f isa typeof(identity) && N == 1
+            return fill!(dest, As[1])
+        else
+            @inbounds for I in eachindex(dest)
+                dest[I] = f(As...)
+            end
+            return dest
+        end
+    end
+    return _broadcast!(f, dest, As...)
+end
+
+# This indirection allows size-dependent implementations.
 @inline function _broadcast!(f, C, A, Bs::Vararg{Any,N}) where N
     shape = broadcast_indices(C)
     @boundscheck check_broadcast_indices(shape, A, Bs...)
diff --git a/base/sparse/higherorderfns.jl b/base/sparse/higherorderfns.jl
index c8f77de2cd416..749903529669a 100644
--- a/base/sparse/higherorderfns.jl
+++ b/base/sparse/higherorderfns.jl
@@ -94,7 +94,7 @@ end
 broadcast(f::Tf, A::SparseVector) where {Tf} = _noshapecheck_map(f, A)
 broadcast(f::Tf, A::SparseMatrixCSC) where {Tf} = _noshapecheck_map(f, A)
 
-function broadcast!(f::Tf, C::SparseVecOrMat, ::Void) where Tf
+@inline function broadcast!(f::Tf, C::SparseVecOrMat, ::BroadcastStyle) where Tf
     isempty(C) && return _finishempty!(C)
     fofnoargs = f()
     if _iszero(fofnoargs) # f() is zero, so empty C
@@ -107,11 +107,16 @@ function broadcast!(f::Tf, C::SparseVecOrMat, ::Void) where Tf
     end
     return C
 end
-function broadcast!(f, dest::SparseVecOrMat, ::Void, A, Bs::Vararg{Any,N}) where N
-    if isa(f, typeof(identity)) && N == 0 && isa(A, Number)
-        return fill!(dest, A)
+@inline function broadcast!(f::Tf, dest::SparseVecOrMat, style::BroadcastStyle, As::Vararg{Any,N}) where {Tf,N}
+    if f isa typeof(identity) && N == 1
+        A = As[1]
+        if A isa Number
+            return fill!(dest, A)
+        elseif A isa AbstractArray && Base.axes(dest) == Base.axes(A)
+            return copy!(dest, A)
+        end
     end
-    return spbroadcast_args!(f, dest, Broadcast.combine_styles(A, Bs...), A, Bs...)
+    return spbroadcast_args!(f, dest, style, As...)
 end
 
 # the following three similar defs are necessary for type stability in the mixed vector/matrix case