Merge pull request #19934 from Sacha0/twoargspbcbang

broadcast!(f, C) for sparse vector/matrix C

base/sparse/higherorderfns.jl

@@ -89,6 +89,19 @@ function _noshapecheck_map{Tf,N}(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseVecO
 end
 # (3) broadcast[!] entry points
 broadcast{Tf}(f::Tf, A::SparseVecOrMat) = _noshapecheck_map(f, A)
+function broadcast!{Tf}(f::Tf, C::SparseVecOrMat)
+    isempty(C) && return _finishempty!(C)
+    fofnoargs = f()
+    if _iszero(fofnoargs) # f() is zero, so empty C
+        trimstorage!(C, 0)
+        _finishempty!(C)
+    else # f() is nonzero, so densify C and fill with independent calls to f()
+        _densestructure!(C)
+        storedvals(C)[1] = fofnoargs
+        broadcast!(f, view(storedvals(C), 2:length(storedvals(C))))
+    end
+    return C
+end
 broadcast!{Tf}(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) = map!(f, C, A)
 function broadcast!{Tf,N}(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N})
     _aresameshape(C, A, Bs...) && return _noshapecheck_map!(f, C, A, Bs...)

test/sparse/higherorderfns.jl

@@ -76,6 +76,16 @@ end
     end
 end
 
+@testset "broadcast! implementation specialized for solely an output sparse vector/matrix (no inputs)" begin
+    N, M, p = 10, 12, 0.4
+    V, C = sprand(N, p), sprand(N, M, p)
+    fV, fC = Array(V), Array(C)
+    @test broadcast!(() -> 0, V) == sparse(broadcast!(() -> 0, fV))
+    @test broadcast!(() -> 0, C) == sparse(broadcast!(() -> 0, fC))
+    @test let z = 0, fz = 0; broadcast!(() -> z += 1, V) == broadcast!(() -> fz += 1, fV); end
+    @test let z = 0, fz = 0; broadcast!(() -> z += 1, C) == broadcast!(() -> fz += 1, fC); end
+end
+
 @testset "broadcast[!] implementation specialized for a single (input) sparse vector/matrix" begin
     # broadcast[!] for a single sparse vector/matrix falls back to map[!], tested extensively
     # above. here we simply lightly exercise the relevant broadcast[!] entry points.