improve #17546 for a[...] .= handling of arrays of arrays and dictionaries of arrays

Author: stevengj

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 ##
@@ -440,4 +440,28 @@ for (sigA, sigB) in ((BitArray, BitArray),
     end
 end
 
+############################################################
+
+# x[...] .= f.(y...) ---> broadcast!(f, dotview(x, ...), y...).
+# The dotview function defaults to view, 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...) = view(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, :view, args...)))
+end
+
+# for a[i...] .= ... where a is an array-of-arrays, just pass a[i...] directly
+# to broadcast!
+dotview{T<:AbstractArray,N,I<:Integer}(a::AbstractArray{T,N}, i::Vararg{I,N}) =
+    a[i...]
+
+# dict[k] .= ... should work if dict[k] is an array
+dotview(a::Associative, k) = a[k]
+dotview(a::Associative, k1, k2, ks...) = a[tuple(k1,k2,ks...)]
+
 end # module

src/julia-syntax.scm

@@ -1549,7 +1549,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