Implement dynamic scope

Author: willtebbutt

src/copyable_task.jl

@@ -1,122 +1,79 @@
+const dynamic_scope = Base.ScopedValues.ScopedValue{Any}(0)
+
+"""
+    get_dynamic_scope()
+
+Returns the dynamic scope associated to `Libtask`. If called from inside a `TapedTask`, this
+will return whatever is contained in its `dynamic_scope` field.
+
+See also [`set_dynamic_scope!`](@ref).
+"""
+get_dynamic_scope() = dynamic_scope[]
+
 __v::Int = 5
 @noinline function produce(x)
     global __v = 4
     return nothing
 end
 
-mutable struct TapedTask{Tmc<:MistyClosure,Targs}
-    const mc::Tmc
+function build_callable(ir::IRCode)
+    seed_id!()
+    bb, refs = derive_copyable_task_ir(BBCode(ir))
+    ir = IRCode(bb)
+    optimised_ir = Mooncake.optimise_ir!(ir)
+    return MistyClosure(optimised_ir, refs...; do_compile=true), refs[end]
+end
+
+mutable struct TapedTask{Tdynamic_scope,Targs,Tmc<:MistyClosure}
+    dynamic_scope::Tdynamic_scope
     args::Targs
-    const position::Base.RefValue{Int32}
-    const deepcopy_types::Type
+    const mc::Tmc
+    const position::Base.RefValue{Int32} # As does this
 end
 
 """
-    Base.copy(t::TapedTask)
-
-Makes a copy of `t` which can be run. For the most part, calls to [`consume`](@ref) on the
-copied task will give the same results as the original. There are, however, substantial
-limitations to this, detailed in the extended help.
+    TapedTask(dynamic_scope::Any, f, args...)
 
-# Extended Help
-
-We call a copy of a `TapedTask` _consistent_ with the original if the call to `==` in the
-loop below always returns `true`:
-```julia
-t = <some_TapedTask>
-tc = copy(t)
-for (v, vc) in zip(t, tc)
-    v == vc
-end
-```
-(provided that `==` is implemented for all `v` that are produced). Convesely, we refer to a
-copy as _inconsistent_ if this property doesn't hold. In order to ensure
-consistency, we need to ensure that independent copies are made of anything which might be
-mutated by the task or its copy during subsequent `consume` calls. Failure to do this can
-cause problems if, for example, a task reads-to and writes-from some memory.
-If we call `consume` on the original task, and then on a copy of it, any changes made by the
-original will be visible to the copy, potentially causing its behaviour to differ. This can
-manifest itself as a race condition if the task and its copies are run concurrently.
-
-To understand a bit more about when a task is / is not consistent, we need to dig into the
-rather specific semantics of `copy`. Calling `copy` on a `TapedTask` does the following:
-1. `copy` the `position` field,
-2. `map`s `_tape_copy` over the `args` field, and
-3. `map`s `_tape_copy` over the all of the data closed over in the `OpaqueClosure` which
-    implements the task (specifically the values _inside_ the `Ref`s) -- call these the
-    `captures`. Except the last elements of this data, because this is `===` to the
-    `position` field -- for this element we use the copy we made in step 1.
-
-`_tape_copy` doesn't actually make a copy of the object at all if it is not either an
-`Array`, a `Ref`, or an instance of one of the types listed in the task's `deepcopy_types`
-field. If it is an instance of one of these types then `_tape_copy` just calls `deepcopy`.
-
-This behaviour is plainly entirely acceptable if the argument to `_tape_copy` is a bits
-type. For any `mutable struct`s which aren't flagged for `deepcopy`ing, we have an immediate
-risk of inconsistency. Similarly, for any `struct` types which aren't bits types (e.g.
-those which contain an `Array`, `Ref`, or some other `mutable struct` either directly as one
-of their fields, or as a field of a field, etc), we have an inconsistency risk.
-
-Furthermore, for anything which _is_ `deepcopy`ed we introduce inconsistency risks. If, for
-example, two elements of the data closed over by the task alias one another, calling
-`deepcopy` on them separately will cause the copies to _not_ alias one another.
-The same thing can happen if one element is `deepcopy`ed and the other not. For example, if
-we have both an `Array` `x` and `view(x, inds)` stored in separate elements of `captures`,
-`x` will be `deepcopy`ed, while `view(x, inds)` will not. In the copy of `captures`, the
-`view` will still be a view into the original `x`, not the `deepcopy`ed version. Again, this
-introduces inconsistency.
-
-Why do we have these semantics? We have them because Libtask has always had them, and at the
-time of writing we're unsure whether AdvancedPS.jl, and by extension Turing.jl rely on this
-behaviour.
-
-What other options do we have? Simply calling `deepcopy` on a `TapedTask` works fine, and
-should reliably result in consistent behaviour between a `TapedTask` and any copies of it.
-This would, therefore, be a preferable implementation. We should try to determine whether
-this is a viable option.
+Construct a `TapedTask` with the specified `dynamic_scope`, for function `f` and positional
+arguments `args`.
 """
-function Base.copy(t::T) where {T<:TapedTask}
-    captures = t.mc.oc.captures
-    new_captures = map(Base.Fix2(_copy_capture, t.deepcopy_types), captures)
-    new_position = new_captures[end] # baked in later on.
-    new_args = map(Base.Fix2(_tape_copy, t.deepcopy_types), t.args)
-    new_mc = Mooncake.replace_captures(t.mc, new_captures)
-    return T(new_mc, new_args, new_position, t.deepcopy_types)
-end
-
-function _copy_capture(r::Ref{T}, deepcopy_types::Type) where {T}
-    new_capture = Ref{T}()
-    if isassigned(r)
-        new_capture[] = _tape_copy(r[], deepcopy_types)
-    end
-    return new_capture
+function TapedTask(dynamic_scope::Any, fargs...)
+    mc, count_ref = build_callable(Base.code_ircode_by_type(typeof(fargs))[1][1])
+    return TapedTask(dynamic_scope, fargs[2:end], mc, count_ref)
 end
 
-_tape_copy(v, deepcopy_types::Type) = v isa deepcopy_types ? deepcopy(v) : v
-
-# Not sure that we need this in the new implementation.
-_tape_copy(box::Core.Box, deepcopy_types::Type) = error("Found a box")
+"""
+    set_dynamic_scope!(t::TapedTask, new_dynamic_scope)::Nothing
 
-@inline consume(t::TapedTask) = t.mc(t.args...)
+Set the `dynamic_scope` of `t` to `new_dynamic_scope`. Any references to 
+`LibTask.dynamic_scope` in future calls to `consume(t)` (either directly, or implicitly via
+iteration) will see this new value.
 
-function initialise!(t::TapedTask, args::Vararg{Any,N})::Nothing where {N}
-    t.position[] = -1
-    t.args = args
+See also: [`get_dynamic_scope`](@ref).
+"""
+function set_dynamic_scope!(t::TapedTask{T}, new_dynamic_scope::T)::Nothing where {T}
+    t.dynamic_scope = new_dynamic_scope
     return nothing
 end
 
-function TapedTask(fargs...; deepcopy_types::Type=Union{})
-    sig = typeof(fargs)
-    mc, count_ref = build_callable(Base.code_ircode_by_type(sig)[1][1])
-    return TapedTask(mc, fargs[2:end], count_ref, Union{deepcopy_types,Array,Ref})
-end
+"""
+    Base.copy(t::TapedTask)
 
-function build_callable(ir::IRCode)
-    seed_id!()
-    bb, refs = derive_copyable_task_ir(BBCode(ir))
-    ir = IRCode(bb)
-    optimised_ir = Mooncake.optimise_ir!(ir)
-    return MistyClosure(optimised_ir, refs...; do_compile=true), refs[end]
+Makes a completely independent copy of `t`. `consume` can be applied to either the copy of
+`t` or the original without advancing the state of the other.
+"""
+Base.copy(t::T) where {T<:TapedTask} = deepcopy(t)
+
+"""
+    consume(t::TapedTask)
+
+Run `t` until it makes a call to `produce`. If this is the first time that `t` has been
+called, it start execution from the entry point. If `consume` has previously been called on
+`t`, it will resume from the last `produce` call. If there are no more `produce` calls,
+`nothing` will be returned.
+"""
+@inline function consume(t::TapedTask)
+    return Base.ScopedValues.with(() -> t.mc(t.args...), dynamic_scope => t.dynamic_scope)
 end
 
 """
@@ -288,7 +245,7 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple}
             n += 1
             ssa_id_to_ref_index_map[id] = n
             ref_index_to_ssa_id_map[n] = id
-            ref_index_to_type_map[n] = stmt.type
+            ref_index_to_type_map[n] = CC.widenconst(stmt.type)
         end
     end
 
@@ -382,8 +339,25 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple}
                 push!(inst_pairs, (id, inst))
             elseif stmt isa Nothing
                 push!(inst_pairs, (id, inst))
+            elseif stmt isa GlobalRef
+                ref_ind = ssa_id_to_ref_index_map[id]
+                expr = Expr(:call, set_ref_at!, refs_id, ref_ind, stmt)
+                push!(inst_pairs, (id, new_inst(expr)))
+            elseif stmt isa Core.PiNode
+                if stmt.val isa ID
+                    ref_ind = ssa_id_to_ref_index_map[stmt.val]
+                    val_id = ID()
+                    expr = Expr(:call, get_ref_at, refs_id, ref_ind)
+                    push!(inst_pairs, (val_id, new_inst(expr)))
+                    push!(inst_pairs, (id, new_inst(Core.PiNode(val_id, stmt.typ))))
+                else
+                    push!(inst_pairs, (id, inst))
+                end
+                set_ind = ssa_id_to_ref_index_map[id]
+                set_expr = Expr(:call, set_ref_at!, refs_id, set_ind, id)
+                push!(inst_pairs, (ID(), new_inst(set_expr)))
             else
-                throw(error("Unhandled stmt $stmt"))
+                throw(error("Unhandled stmt $stmt of type $(typeof(stmt))"))
             end
         end
 
@@ -451,7 +425,9 @@ function derive_copyable_task_ir(ir::BBCode)::Tuple{BBCode,Tuple}
 end
 
 # Helper used in `derive_copyable_task_ir`.
-@inline get_ref_at(refs::R, n::Int) where {R<:Tuple} = refs[n][]
+@inline function get_ref_at(refs::R, n::Int) where {R<:Tuple}
+    return refs[n][]
+end
 
 # Helper used in `derive_copyable_task_ir`.
 @inline function set_ref_at!(refs::R, n::Int, val) where {R<:Tuple}

src/test_utils.jl

@@ -6,6 +6,7 @@ using ..Libtask: TapedTask
 
 struct Testcase
     name::String
+    dynamic_scope::Any
     fargs::Tuple
     expected_iteration_results::Vector
 end
@@ -14,7 +15,7 @@ function (case::Testcase)()
     testset = @testset "$(case.name)" begin
 
         # Construct the task.
-        t = TapedTask(case.fargs...)
+        t = TapedTask(case.dynamic_scope, case.fargs...)
 
         # Iterate through t. Record the results, and take a copy after each iteration.
         iteration_results = []
@@ -39,21 +40,22 @@ function test_cases()
     return Testcase[
         Testcase(
             "single block",
+            nothing,
             (single_block, 5.0),
             [sin(5.0), sin(sin(5.0)), sin(sin(sin(5.0))), sin(sin(sin(sin(5.0))))],
         ),
-        Testcase("produce old", (produce_old_value, 5.0), [sin(5.0), sin(5.0)]),
-        Testcase("branch on old value l", (branch_on_old_value, 2.0), [true, 2.0]),
-        Testcase("branch on old value r", (branch_on_old_value, -1.0), [false, -2.0]),
-        Testcase("no produce", (no_produce_test, 5.0, 4.0), []),
-        Testcase("new object", (new_object_test, 5, 4), [C(5, 4), C(5, 4)]),
-        Testcase("branching test l", (branching_test, 5.0, 4.0), [string(sin(5.0))]),
-        Testcase("branching test r", (branching_test, 4.0, 5.0), [sin(4.0) * cos(5.0)]),
-        Testcase("unused argument test", (unused_argument_test, 3), [1]),
-        Testcase("test with const", (test_with_const,), [1]),
-        Testcase("while loop", (while_loop,), collect(1:9)),
+        Testcase("produce old", nothing, (produce_old_value, 5.0), [sin(5.0), sin(5.0)]),
+        Testcase("branch on old value l", nothing, (branch_on_old_value, 2.0), [true, 2.0]),
+        Testcase("branch on old value r", nothing, (branch_on_old_value, -1.0), [false, -2.0]),
+        Testcase("no produce", nothing, (no_produce_test, 5.0, 4.0), []),
+        Testcase("new object", nothing, (new_object_test, 5, 4), [C(5, 4), C(5, 4)]),
+        Testcase("branching test l", nothing, (branching_test, 5.0, 4.0), [string(sin(5.0))]),
+        Testcase("branching test r", nothing, (branching_test, 4.0, 5.0), [sin(4.0) * cos(5.0)]),
+        Testcase("unused argument test", nothing, (unused_argument_test, 3), [1]),
+        Testcase("test with const", nothing, (test_with_const,), [1]),
+        Testcase("while loop", nothing, (while_loop,), collect(1:9)),
         Testcase(
-            "foreigncall tester", (foreigncall_tester, "hi"), [Ptr{UInt8}, Ptr{UInt8}]
+            "foreigncall tester", nothing, (foreigncall_tester, "hi"), [Ptr{UInt8}, Ptr{UInt8}]
         ),
 
         # Failing tests