parser: Support getindex

Author: jpsamaroo

src/thunk.jl

@@ -365,8 +365,14 @@ end
 function _par(ex::Expr; lazy=true, recur=true, opts=())
     f = nothing
     body = nothing
-    if recur && @capture(ex, f_(allargs__)) || @capture(ex, f_(allargs__) do cargs_ body_ end) || @capture(ex, allargs__->body_)
+    arg1 = nothing
+    if recur && @capture(ex, f_(allargs__)) || @capture(ex, f_(allargs__) do cargs_ body_ end) || @capture(ex, allargs__->body_) || @capture(ex, arg1_[allargs__])
         f = replace_broadcast(f)
+        if arg1 !== nothing
+            # Indexing (A[2,3])
+            f = Base.getindex
+            pushfirst!(allargs, arg1)
+        end
         args = filter(arg->!Meta.isexpr(arg, :parameters), allargs)
         kwargs = filter(arg->Meta.isexpr(arg, :parameters), allargs)
         if !isempty(kwargs)

test/thunk.jl

@@ -118,6 +118,23 @@ end
         @test t isa Dagger.DTask
         @test fetch(t) == sum(A; dims=1)
     end
+    @testset "getindex" begin
+        A = rand(4, 4)
+
+        t = @spawn A[1, 2]
+        @test t isa Dagger.DTask
+        @test fetch(t) == A[1, 2]
+
+        B = Dagger.@spawn rand(4, 4)
+        t = @spawn B[1, 2]
+        @test t isa Dagger.DTask
+        @test fetch(t) == fetch(B)[1, 2]
+
+        R = Ref(42)
+        t = @spawn R[]
+        @test t isa Dagger.DTask
+        @test fetch(t) == 42
+    end
     @testset "invalid expression" begin
         @test_throws LoadError eval(:(@spawn 1))
         @test_throws LoadError eval(:(@spawn begin 1 end))