make x[...] .= ... assign in-place (#17546)

* make x[...] .= ... assign in-place (fixes bug in #17510)

* doc fix

* You're the top! You're the Coliseum.

* better var name

Author: stevengj

doc/manual/functions.rst

@@ -660,6 +660,10 @@ 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.
+``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.
 
 (In future versions of Julia, operators like ``.*`` will also be handled with
 the same mechanism: they will be equivalent to ``broadcast`` calls and

src/julia-syntax.scm

@@ -1542,6 +1542,15 @@
         (cadr expr)  ;; eta reduce `x->f(x)` => `f`
         `(-> ,argname (block ,@splat ,expr)))))
 
+(define (ref-to-view expr)
+  (if (and (pair? expr) (eq? (car expr) 'ref))
+      (let* ((ex (partially-expand-ref expr))
+             (stmts (butlast (cdr ex)))
+             (refex (last    (cdr ex)))
+             (nuref `(call (top view) ,(caddr refex) ,@(cdddr refex))))
+        `(block ,@stmts ,nuref))
+      expr))
+
 ; fuse nested calls to expr == f.(args...) into a single broadcast call,
 ; or a broadcast! call if lhs is non-null.
 (define (expand-fuse-broadcast lhs rhs)
@@ -1657,14 +1666,15 @@
                         (cons farg new-fargs) (cons arg new-args) renames varfarg vararg))))))
           (cf (cdadr f) args '() '() '() '() '()))
         e)) ; (not (fuse? e))
-  (let ((e (compress-fuse (dot-to-fuse rhs)))) ; an expression '(fuse func args) if expr is a dot call
+  (let ((e (compress-fuse (dot-to-fuse rhs))) ; an expression '(fuse func args) if expr is a dot call
+        (lhs-view (ref-to-view lhs))) ; x[...] expressions on lhs turn in to view(x, ...) to update x in-place
     (if (fuse? e)
         (if (null? lhs)
-            (expand-forms `(call broadcast ,(from-lambda (cadr e)) ,@(caddr e)))
-            (expand-forms `(call broadcast! ,(from-lambda (cadr e)) ,lhs ,@(caddr e))))
+            (expand-forms `(call (top broadcast) ,(from-lambda (cadr e)) ,@(caddr e)))
+            (expand-forms `(call (top broadcast!) ,(from-lambda (cadr e)) ,lhs-view ,@(caddr e))))
         (if (null? lhs)
             (expand-forms e)
-            (expand-forms `(call broadcast! identity ,lhs ,e))))))
+            (expand-forms `(call (top broadcast!) identity ,lhs-view ,e))))))
 
 ;; table mapping expression head to a function expanding that form
 (define expand-table

test/broadcast.jl

@@ -265,6 +265,14 @@ let x = [1:4;], y = x
     @test y === x == [8,8,8,8]
     y .-= 1:4
     @test y === x == [7,6,5,4]
+    x[1:2] .= 1
+    @test y === x == [1,1,5,4]
+    x[1:2] .+= [2,3]
+    @test y === x == [3,4,5,4]
+    x[:] .= 0
+    @test y === x == [0,0,0,0]
+    x[2:end] .= 1:3
+    @test y === x == [0,1,2,3]
 end
 
 # PR 16988