Fix @~ macro for case a dot call is wrapping a normal call (#55)

Author: tkf

src/lazymacro.jl

@@ -17,7 +17,8 @@ is_call(ex::Expr) =
     ex.head == :call && !startswith(String(ex.args[1]), ".")
 
 is_dotcall(ex::Expr) =
-    ex.head == :. || (ex.head == :call && startswith(String(ex.args[1]), "."))
+    (ex.head == :. && ex.args[2].head === :tuple) ||
+    (ex.head == :call && startswith(String(ex.args[1]), "."))
 # e.g., `f.(x, y, z)` or `x .+ y .+ z`
 
 lazy_expr(x) = x
@@ -40,11 +41,17 @@ end
 bc_expr_impl(x) = x
 function bc_expr_impl(ex::Expr)
     # walk down chain of dot calls
-    if is_dotcall(ex)
-        return Expr(ex.head,
-                    lazy_expr(ex.args[1]), # function name (`f`, `.+`, etc.)
-                    bc_expr_impl.(ex.args[2:end])...) # arguments
+    if ex.head == :. && ex.args[2].head === :tuple
+        @assert length(ex.args) == 2  # argument is always expressed as a tuple
+        f = ex.args[1]  # function name
+        args = ex.args[2].args
+        return Expr(ex.head, lazy_expr(f), Expr(:tuple, bc_expr_impl.(args)...))
+    elseif ex.head == :call && startswith(String(ex.args[1]), ".")
+        f = ex.args[1]  # function name (e.g., `.+`)
+        args = ex.args[2:end]
+        return Expr(ex.head, lazy_expr(f), bc_expr_impl.(args)...)
     else
+        @assert !is_dotcall(ex)
         return lazy_expr(ex)
     end
 end

test/macrotests.jl

@@ -1,6 +1,7 @@
 module MacroTests
 
 using Test, LazyArrays, MacroTools
+using Base.Broadcast: Broadcasted
 
 A = randn(6, 6)
 B = BroadcastArray(+, A, 2)
@@ -28,7 +29,7 @@ testparams = [
     @testset "$label" for (label, ex) in testparams
         desired = @eval $ex
         lazy = @eval @~ $ex
-        @test lazy isa Union{Broadcast.Broadcasted, LazyArrays.Applied}
+        @test lazy isa Union{Broadcasted, Applied}
 
         @testset ".= @~ $label" begin
             actual = zero(desired)
@@ -57,4 +58,16 @@ testparams = [
     end
 end
 
+@testset "@~ laziness" begin
+    A = ones(1, 1)
+    x = [1]
+
+    bc = @~ exp.(A * x)
+    @test bc.args isa Tuple{Applied}
+
+    bc = @~ exp.(A * x .+ 1)
+    @test bc.args isa Tuple{Broadcasted}
+    @test bc.args[1].args isa Tuple{Applied, Int}
+end
+
 end  # module