From 4b23f42d99e411cfd032de2f9233180f9dc388c0 Mon Sep 17 00:00:00 2001
From: Sacha Verweij <sacha@stanford.edu>
Date: Sun, 8 Jan 2017 12:55:17 -0800
Subject: [PATCH] Implement and test broadcast!(f, C) for sparse vector/matrix
 C.

---
 base/sparse/higherorderfns.jl | 13 +++++++++++++
 test/sparse/higherorderfns.jl | 10 ++++++++++
 2 files changed, 23 insertions(+)

diff --git a/base/sparse/higherorderfns.jl b/base/sparse/higherorderfns.jl
index 8bae5a82c4043..7d3c369512cb4 100644
--- a/base/sparse/higherorderfns.jl
+++ b/base/sparse/higherorderfns.jl
@@ -81,6 +81,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...)
diff --git a/test/sparse/higherorderfns.jl b/test/sparse/higherorderfns.jl
index 57f4eff8e0230..e6dbcc842f93c 100644
--- a/test/sparse/higherorderfns.jl
+++ b/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.