parser: Support do-blocks

Author: jpsamaroo

src/thunk.jl

@@ -363,37 +363,47 @@ function replace_broadcast(fn::Symbol)
 end
 
 function _par(ex::Expr; lazy=true, recur=true, opts=())
-    if ex.head == :call && recur
-        f = replace_broadcast(ex.args[1])
-        if length(ex.args) >= 2 && Meta.isexpr(ex.args[2], :parameters)
-            args = ex.args[3:end]
-            kwargs = ex.args[2]
-        else
-            args = ex.args[2:end]
-            kwargs = Expr(:parameters)
+    body = nothing
+    if recur && @capture(ex, f_(allargs__)) || @capture(ex, f_(allargs__) do cargs_ body_ end)
+        f = replace_broadcast(f)
+        args = filter(arg->!Meta.isexpr(arg, :parameters), allargs)
+        kwargs = filter(arg->Meta.isexpr(arg, :parameters), allargs)
+        if !isempty(kwargs)
+            kwargs = only(kwargs).args
+        end
+        if body !== nothing
+            f = quote
+                ($(args...); $(kwargs...))->$f($(args...); $(kwargs...)) do $cargs
+                    $body
+                end
+            end
         end
-        args_ex = _par.(args; lazy=lazy, recur=false)
-        kwargs_ex = _par.(kwargs.args; lazy=lazy, recur=false)
         if lazy
-            return :(Dagger.delayed($f, $Options(;$(opts...)))($(args_ex...); $(kwargs_ex...)))
+            return :(Dagger.delayed($f, $Options(;$(opts...)))($(args...); $(kwargs...)))
         else
             sync_var = Base.sync_varname
             @gensym result
             return quote
-                let args = ($(args_ex...),)
-                    $result = $spawn($f, $Options(;$(opts...)), args...; $(kwargs_ex...))
+                let
+                    $result = $spawn($f, $Options(;$(opts...)), $(args...); $(kwargs...))
                     if $(Expr(:islocal, sync_var))
                         put!($sync_var, schedule(Task(()->wait($result))))
                     end
                     $result
                 end
             end
         end
+    elseif lazy
+        # Recurse into the expression
+        return Expr(ex.head, _par_inner.(ex.args, lazy=lazy, recur=recur, opts=opts)...)
     else
-        return Expr(ex.head, _par.(ex.args, lazy=lazy, recur=recur, opts=opts)...)
+        throw(ArgumentError("Invalid Dagger task expression: $ex"))
     end
 end
-_par(ex; kwargs...) = ex
+_par(ex; kwargs...) = throw(ArgumentError("Invalid Dagger task expression: $ex"))
+
+_par_inner(ex; kwargs...) = ex
+_par_inner(ex::Expr; kwargs...) = _par(ex; kwargs...)
 
 """
     Dagger.spawn(f, args...; kwargs...) -> DTask

test/thunk.jl

@@ -81,7 +81,39 @@ end
     end
     @testset "inner macro" begin
         A = rand(4)
-        @test fetch(@spawn sum(@view A[2:3])) ≈ sum(@view A[2:3])
+        t = @spawn sum(@view A[2:3])
+        @test t isa Dagger.DTask
+        @test fetch(t) ≈ sum(@view A[2:3])
+    end
+    @testset "do block" begin
+        A = rand(4)
+
+        t = @spawn sum(A) do a
+            a + 1
+        end
+        @test t isa Dagger.DTask
+        @test fetch(t) ≈ sum(a->a+1, A)
+
+        t = @spawn sum(A; dims=1) do a
+            a + 1
+        end
+        @test t isa Dagger.DTask
+        @test fetch(t) ≈ sum(a->a+1, A; dims=1)
+
+        do_f = f -> f(42)
+        t = @spawn do_f() do x
+            x + 1
+        end
+        @test t isa Dagger.DTask
+        @test fetch(t) == 43
+    end
+    @testset "invalid expression" begin
+        @test_throws LoadError eval(:(@spawn 1))
+        @test_throws LoadError eval(:(@spawn begin 1 end))
+        @test_throws LoadError eval(:(@spawn begin
+            1+1
+            1+1
+        end))
     end
     @testset "waiting" begin
         a = @spawn sleep(1)