at-dot macro for adding dots to function calls

NEWS.md

@@ -41,6 +41,9 @@ Language changes
   * `isa` is now parsed as an infix operator with the same precedence as `in`
     ([#19677]).
 
+  * `@.` is now parsed as `@__dot__`, and can be used to add dots to
+    every function call, operator, and assignment in an expression ([#20321]).
+
 Breaking changes
 ----------------
 

base/broadcast.jl

@@ -5,9 +5,9 @@ module Broadcast
 using Base.Cartesian
 using Base: linearindices, tail, OneTo, to_shape,
             _msk_end, unsafe_bitgetindex, bitcache_chunks, bitcache_size, dumpbitcache,
-            nullable_returntype, null_safe_eltype_op, hasvalue
+            nullable_returntype, null_safe_eltype_op, hasvalue, isoperator
 import Base: broadcast, broadcast!
-export broadcast_getindex, broadcast_setindex!, dotview
+export broadcast_getindex, broadcast_setindex!, dotview, @__dot__
 
 typealias ScalarType Union{Type{Any}, Type{Nullable}}
 
@@ -509,4 +509,64 @@ Base.@propagate_inbounds dotview(args...) = getindex(args...)
 Base.@propagate_inbounds dotview(A::AbstractArray, args...) = view(A, args...)
 Base.@propagate_inbounds dotview{T<:AbstractArray}(A::AbstractArray{T}, args...) = getindex(A, args...)
 
+
+############################################################
+# The parser turns @. into a call to the __dot__ macro,
+# which converts all function calls and assignments into
+# broadcasting "dot" calls/assignments:
+
+dottable(x) = false # avoid dotting spliced objects (e.g. view calls inserted by @view)
+dottable(x::Symbol) = !isoperator(x) || first(string(x)) != '.' || x == :.. # don't add dots to dot operators
+dottable(x::Expr) = x.head != :$
+undot(x) = x
+function undot(x::Expr)
+    if x.head == :.=
+        Expr(:(=), x.args...)
+    elseif x.head == :block # occurs in for x=..., y=...
+        Expr(:block, map(undot, x.args)...)
+    else
+        x
+    end
+end
+__dot__(x) = x
+function __dot__(x::Expr)
+    dotargs = map(__dot__, x.args)
+    if x.head == :call && dottable(x.args[1])
+        Expr(:., dotargs[1], Expr(:tuple, dotargs[2:end]...))
+    elseif x.head == :$
+        x.args[1]
+    elseif x.head == :let # don't add dots to "let x=... assignments
+        Expr(:let, dotargs[1], map(undot, dotargs[2:end])...)
+    elseif x.head == :for # don't add dots to for x=... assignments
+        Expr(:for, undot(dotargs[1]), dotargs[2])
+    elseif (x.head == :(=) || x.head == :function || x.head == :macro) &&
+           Meta.isexpr(x.args[1], :call) # function or macro definition
+        Expr(x.head, x.args[1], dotargs[2])
+    else
+        head = string(x.head)
+        if last(head) == '=' && first(head) != '.'
+            Expr(Symbol('.',head), dotargs...)
+        else
+            Expr(x.head, dotargs...)
+        end
+    end
+end
+"""
+    @. expr
+
+Convert every function call or operator in `expr` into a "dot call"
+(e.g. convert `f(x)` to `f.(x)`), and convert every assignment in `expr`
+to a "dot assignment" (e.g. convert `+=` to `.+=`).
+
+If you want to *avoid* adding dots for selected function calls in
+`expr`, splice those function calls in with `\$`.  For example,
+`@. sqrt(abs(\$sort(x)))` is equivalent to `sqrt.(abs.(sort(x)))`
+(no dot for `sort`).
+
+(`@.` is equivalent to a call to `@__dot__`.)
+"""
+macro __dot__(x)
+    esc(__dot__(x))
+end
+
 end # module

base/exports.jl

@@ -1382,6 +1382,7 @@ export
     @polly,
 
     @assert,
+    @__dot__,
     @enum,
     @label,
     @goto,

doc/src/manual/functions.md

@@ -642,12 +642,17 @@ overwriting `X` with `sin.(Y)` in-place. If the left-hand side is an array-index
 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.
 
+Since adding dots to many operations and function calls in an expression
+can be tedious and lead to code that is difficult to read, the macro
+[`@.`](@ref @__dot__) is provided to convert *every* function call,
+operation, and assignment in an expression into the "dotted" version.
+
 ```jldoctest
 julia> Y = [1.0, 2.0, 3.0, 4.0];
 
 julia> X = similar(Y); # pre-allocate output array
 
-julia> X .= sin.(cos.(Y))
+julia> @. X = sin(cos(Y)) # equivalent to X .= sin.(cos.(Y))
 4-element Array{Float64,1}:
   0.514395
  -0.404239

doc/src/manual/mathematical-operations.md

@@ -151,13 +151,14 @@ it can combine arrays and scalars, arrays of the same size (performing
 the operation elementwise), and even arrays of different shapes (e.g.
 combining row and column vectors to produce a matrix). Moreover, like
 all vectorized "dot calls," these "dot operators" are
-*fusing*. For example, if you compute `2 .* A.^2 .+ sin.(A)` for an
-array `A`, it performs a *single* loop over `A`, computing `2a^2 + sin(a)`
+*fusing*. For example, if you compute `2 .* A.^2 .+ sin.(A)` (or
+equivalently `@. 2A^2 + sin(A)`, using the [`@.`](@ref @__dot__) macro) for
+an array `A`, it performs a *single* loop over `A`, computing `2a^2 + sin(a)`
 for each element of `A`. In particular, nested dot calls like `f.(g.(x))`
 are fused, and "adjacent" binary operators like `x .+ 3 .* x.^2` are
 equivalent to nested dot calls `(+).(x, (*).(3, (^).(x, 2)))`.
 
-Furthermore, "dotted" updating operators like `a .+= b` are parsed
+Furthermore, "dotted" updating operators like `a .+= b` (or `@. a += b`) are parsed
 as `a .= a .+ b`, where `.=` is a fused *in-place* assignment operation
 (see the [dot syntax documentation](@ref man-vectorized)).
 

doc/src/manual/performance-tips.md

@@ -883,11 +883,12 @@ resulting loops can be fused with surrounding computations. For example,
 consider the two functions:
 
 ```julia
-f(x) = 3 * x.^2 + 4 * x + 7 * x.^3
-fdot(x) = 3 .* x.^2 .+ 4 .* x .+ 7 .* x.^3
+f(x) = 3x.^2 + 4x + 7x.^3
+fdot(x) = @. 3x^2 + 4x + 7x^3 # equivalent to 3 .* x.^2 .+ 4 .* x .+ 7 .* x.^3
 ```
 
-Both `f` and `fdot` compute the same thing.  However, `fdot` is
+Both `f` and `fdot` compute the same thing.  However, `fdot`
+(defined with the help of the [`@.`](@ref @__dot__) macro) is
 significantly faster when applied to an array:
 
 ```julia

doc/src/stdlib/arrays.md

@@ -41,13 +41,19 @@ Base.linspace
 Base.logspace
 ```
 
-## Mathematical operators and functions
+## Broadcast and vectorization
 
-All mathematical operations and functions are supported for arrays
+See also the [dot syntax for vectorizing functions](@ref man-vectorized);
+for example, `f.(args...)` implicitly calls `broadcast(f, args...)`.
+Rather than relying on "vectorized" methods of functions like `sin`
+to operate on arrays, you should use `sin.(a)` to vectorize via `broadcast`.
 
 ```@docs
 Base.broadcast
 Base.Broadcast.broadcast!
+Base.@__dot__
+Base.Broadcast.broadcast_getindex
+Base.Broadcast.broadcast_setindex!
 ```
 
 ## Indexing, Assignment, and Concatenation
@@ -63,8 +69,6 @@ Base.parent
 Base.parentindexes
 Base.slicedim
 Base.setindex!(::AbstractArray, ::Any, ::Any...)
-Base.Broadcast.broadcast_getindex
-Base.Broadcast.broadcast_setindex!
 Base.isassigned
 Base.cat
 Base.vcat

src/julia-parser.scm

@@ -2079,7 +2079,9 @@
           ((eqv? t #\@)
            (take-token s)
            (with-space-sensitive
-            (let ((head (parse-unary-prefix s)))
+            (let ((head (if (eq? (peek-token s) '|.|)
+                            (begin (take-token s) '__dot__)
+                            (parse-unary-prefix s))))
               (if (eq? head '__LINE__)
                   (input-port-line (ts:port s))
                   (begin

test/broadcast.jl

@@ -217,7 +217,7 @@ let A = [sqrt(i)+j for i = 1:3, j=1:4]
 end
 let x = sin.(1:10)
     @test atan2.((x->x+1).(x), (x->x+2).(x)) == broadcast(atan2, x+1, x+2) == broadcast(atan2, x.+1, x.+2)
-    @test sin.(atan2.([x+1,x+2]...)) == sin.(atan2.(x+1,x+2))
+    @test sin.(atan2.([x+1,x+2]...)) == sin.(atan2.(x+1,x+2)) == @. sin(atan2(x+1,x+2))
     @test sin.(atan2.(x, 3.7)) == broadcast(x -> sin(atan2(x,3.7)), x)
     @test atan2.(x, 3.7) == broadcast(x -> atan2(x,3.7), x) == broadcast(atan2, x, 3.7)
 end
@@ -226,6 +226,9 @@ let g = Int[]
     f17300(x) = begin; push!(g, x); x+2; end
     f17300.(f17300.(f17300.(1:3)))
     @test g == [1,3,5, 2,4,6, 3,5,7]
+    empty!(g)
+    @. f17300(f17300(f17300(1:3)))
+    @test g == [1,3,5, 2,4,6, 3,5,7]
 end
 # fusion with splatted args:
 let x = sin.(1:10), a = [x]
@@ -244,6 +247,28 @@ let x = [1:4;]
     @test sin.(f17300kw.(x, y=1)) == sin.(f17300kw.(x; y=1)) == sin.(x .+ 1)
 end
 
+# splice escaping of @.
+let x = [4, -9, 1, -16]
+    @test [2, 3, 4, 5] == @.(1 + sqrt($sort(abs(x))))
+end
+
+# interaction of @. with let
+@test [1,4,9] == @. let x = [1,2,3]; x^2; end
+
+# interaction of @. with for loops
+let x = [1,2,3], y = x
+    @. for i = 1:3
+        y = y^2 # should convert to y .= y.^2
+    end
+    @test x == [1,256,6561]
+end
+
+# interaction of @. with function definitions
+let x = [1,2,3]
+    @. f(x) = x^2
+    @test f(x) == [1,4,9]
+end
+
 # PR #17510: Fused in-place assignment
 let x = [1:4;], y = x
     y .= 2:5
@@ -259,15 +284,15 @@ let x = [1:4;], y = x
     @test y === x == [9,9,9,9]
     y .-= 1
     @test y === x == [8,8,8,8]
-    y .-= 1:4
+    @. y -= 1:4          # @. should convert to .-=
     @test y === x == [7,6,5,4]
     x[1:2] .= 1
     @test y === x == [1,1,5,4]
-    x[1:2] .+= [2,3]
+    @. x[1:2] .+= [2,3]  # use .+= to make sure @. works with dotted assignment
     @test y === x == [3,4,5,4]
-    x[:] .= 0
+    @. x[:] .= 0         # use .= to make sure @. works with dotted assignment
     @test y === x == [0,0,0,0]
-    x[2:end] .= 1:3
+    @. x[2:end] = 1:3    # @. should convert to .=
     @test y === x == [0,1,2,3]
 end
 let a = [[4, 5], [6, 7]]

test/subarray.jl

@@ -479,7 +479,7 @@ Y = 4:-1:1
 
 @test isa(@view(X[1:3]), SubArray)
 
-@test X[1:end] == @view X[1:end]
+@test X[1:end] == @.(@view X[1:end]) # test compatibility of @. and @view
 @test X[1:end-3] == @view X[1:end-3]
 @test X[1:end,2,2] == @view X[1:end,2,2]
 # @test X[1,1:end-2] == @view X[1,1:end-2] # TODO: Re-enable after partial linear indexing deprecation
@@ -518,14 +518,16 @@ end
     @test x == [5,6,35,4]
     x[Y[2:3]] .= 7:8
     @test x == [5,8,7,4]
-    x[(3,)..., ()...] .+= 3
+    @. x[(3,)..., ()...] += 3 # @. should convert to .+=, test compatibility with @views
     @test x == [5,8,10,4]
     i = Int[]
     # test that lhs expressions in update operations are evaluated only once:
     x[push!(i,4)[1]] += 5
     @test x == [5,8,10,9] && i == [4]
     x[push!(i,3)[end]] += 2
     @test x == [5,8,12,9] && i == [4,3]
+    @. x[3:end] = 0       # make sure @. works with end expressions in @views
+    @test x == [5,8,0,0]
 end
 @views @test isa(X[1:3], SubArray)
 @test X[1:end] == @views X[1:end]