chunks: Allow weak Chunk references in Thunk args

src/chunks.jl

@@ -266,6 +266,16 @@ function savechunk(data, dir, f)
     Chunk{typeof(data),typeof(fr),typeof(proc),typeof(scope)}(typeof(data), domain(data), fr, proc, scope, true)
 end
 
+struct WeakChunk
+    x::WeakRef
+end
+WeakChunk(c::Chunk) = WeakChunk(WeakRef(c))
+unwrap_weak(c::WeakChunk) = c.x.value
+function unwrap_weak_checked(c::WeakChunk)
+    c = unwrap_weak(c)
+    @assert c !== nothing
+    return c
+end
 
 Base.@deprecate_binding AbstractPart Union{Chunk, Thunk}
 Base.@deprecate_binding Part Chunk

src/sch/Sch.jl

@@ -8,7 +8,7 @@ import Statistics: mean
 import Random: randperm
 
 import ..Dagger
-import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, ThunkFailedException, Chunk, OSProc, AnyScope
+import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, ThunkFailedException, Chunk, WeakChunk, OSProc, AnyScope
 import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak_checked, affinity, tochunk, timespan_start, timespan_finish, procs, move, chunktype, processor, default_enabled, get_processors, get_parent, execute!, rmprocs!, addprocs!, thunk_processor, constrain, cputhreadtime
 
 const OneToMany = Dict{Thunk, Set{Thunk}}

src/sch/eager.jl

@@ -90,7 +90,10 @@ function eager_thunk()
             added_future, future, uid, ref, f, args, opts = take!(EAGER_THUNK_CHAN)
             # preserve inputs until they enter the scheduler
             tid = GC.@preserve args begin
-                _args = map(x->x isa Dagger.EagerThunk ? ThunkID(EAGER_ID_MAP[x.uid], x.thunk_ref) : x, args)
+                _args = map(x->x isa Dagger.EagerThunk ? ThunkID(EAGER_ID_MAP[x.uid], x.thunk_ref) :
+                               x isa Dagger.Chunk ? WeakChunk(x) :
+                               x,
+                            args)
                 add_thunk!(f, h, _args...; future=future, ref=ref, opts...)
             end
             EAGER_ID_MAP[uid] = tid.id

src/sch/util.jl

@@ -113,7 +113,7 @@ function reschedule_inputs!(state, thunk, seen=Set{Thunk}())
             input in seen && continue
 
             # Unseen
-            if istask(input) || (input isa Chunk)
+            if istask(input) || isa(input, Chunk)
                 push!(get!(()->Set{Thunk}(), state.waiting_data, input), thunk)
             end
             istask(input) || continue