Merge pull request #17568 from stevengj/dotview

improve a[...] .= handling of arrays of arrays and dicts of arrays

Author: JeffBezanson

base/broadcast.jl

@@ -5,7 +5,7 @@ module Broadcast
 using Base.Cartesian
 using Base: promote_op, promote_eltype, promote_eltype_op, @get!, _msk_end, unsafe_bitgetindex, linearindices, tail, OneTo, to_shape
 import Base: .+, .-, .*, ./, .\, .//, .==, .<, .!=, .<=, .÷, .%, .<<, .>>, .^
-export broadcast, broadcast!, bitbroadcast
+export broadcast, broadcast!, bitbroadcast, dotview
 export broadcast_getindex, broadcast_setindex!
 
 ## Broadcasting utilities ##
@@ -437,4 +437,24 @@ for (sigA, sigB) in ((BitArray, BitArray),
     end
 end
 
+############################################################
+
+# x[...] .= f.(y...) ---> broadcast!(f, dotview(x, ...), y...).
+# The dotview function defaults to getindex, but we override it in
+# a few cases to get the expected in-place behavior without affecting
+# explicit calls to view.   (All of this can go away if slices
+# are changed to generate views by default.)
+
+dotview(args...) = getindex(args...)
+dotview(A::AbstractArray, args...) = view(A, args...)
+dotview{T<:AbstractArray}(A::AbstractArray{T}, args...) = getindex(A, args...)
+# avoid splatting penalty in common cases:
+for nargs = 0:5
+    args = Symbol[Symbol("x",i) for i = 1:nargs]
+    eval(Expr(:(=), Expr(:call, :dotview, args...),
+                    Expr(:call, :getindex, args...)))
+    eval(Expr(:(=), Expr(:call, :dotview, :(A::AbstractArray), args...),
+                    Expr(:call, :view, :A, args...)))
+end
+
 end # module

doc/manual/functions.rst

@@ -660,7 +660,7 @@ calls do not allocate new arrays over and over again for the results
 except that, as above, the ``broadcast!`` loop is fused with any nested
 "dot" calls.  For example, ``X .= sin.(Y)`` is equivalent to
 ``broadcast!(sin, X, Y)``, overwriting ``X`` with ``sin.(Y)`` in-place.
-If the left-hand side is a ``getindex`` expression, e.g.
+If the left-hand side is an array-indexing expression, e.g.
 ``X[2:end] .= sin.(Y)``, then it translates to ``broadcast!`` on a ``view``,
 e.g. ``broadcast!(sin, view(X, 2:endof(X)), Y)``, so that the left-hand
 side is updated in-place.

src/julia-syntax.scm

@@ -1559,7 +1559,7 @@
       (let* ((ex (partially-expand-ref expr))
              (stmts (butlast (cdr ex)))
              (refex (last    (cdr ex)))
-             (nuref `(call (top view) ,(caddr refex) ,@(cdddr refex))))
+             (nuref `(call (top dotview) ,(caddr refex) ,@(cdddr refex))))
         `(block ,@stmts ,nuref))
       expr))
 

test/broadcast.jl

@@ -274,6 +274,16 @@ let x = [1:4;], y = x
     x[2:end] .= 1:3
     @test y === x == [0,1,2,3]
 end
+let a = [[4, 5], [6, 7]]
+    a[1] .= 3
+    @test a == [[3, 3], [6, 7]]
+end
+let d = Dict(:foo => [1,3,7], (3,4) => [5,9])
+    d[:foo] .+= 2
+    @test d[:foo] == [3,5,9]
+    d[3,4] .-= 1
+    @test d[3,4] == [4,8]
+end
 
 # PR 16988
 @test Base.promote_op(+, Bool) === Int