Support safe exception regions in TapedTask

AGENTS.md

@@ -0,0 +1,97 @@
+# AGENTS.md
+
+## Purpose
+
+Libtask.jl provides resumable, **copyable** functions (coroutines) for Julia, with optional
+function-specific globals. [Turing.jl](https://github.com/TuringLang/Turing.jl) (via
+[AdvancedPS.jl](https://github.com/TuringLang/AdvancedPS.jl)) uses it for particle-based
+inference, where copying a partially-run task — not just resuming it — is the load-bearing
+feature.
+
+Mechanism: `Base.code_ircode_by_type` recovers typed IR; `derive_copyable_task_ir` splits it
+at every `produce` site, lifts variables that must survive a suspension into `Ref`s, and emits
+a re-enterable `MistyClosure`. `copy(::TapedTask)` deep-copies those `Ref`s to fork state.
+
+Priorities, in order: (1) `produce`/`consume`/`copy` correctness — a copy resumes
+independently and reproduces a fresh run exactly; (2) faithful control flow (branches, loops,
+phi nodes, exception regions), failing loudly and locally on unsupported constructs
+(`assert_can_handle_control_flow`) rather than resuming wrongly; (3) type stability and low
+allocation on the `consume` hot path; (4) minimal dependencies and a small IR-transform core.
+
+## Repository Layout
+
+- `src/Libtask.jl` — module entry, exports (`TapedTask`, `consume`, `produce`,
+  `get_taped_globals`, `set_taped_globals!`, `NotInTapedTaskError`).
+- `src/copyable_task.jl` — public API + `build_callable` (IR-derivation/compile/cache entry),
+  the `MistyClosure` cache, and `copy`.
+- `src/transformation.jl` — core IR transform: `produce`-site detection (`is_produce_stmt`,
+  `stmt_might_produce`), block splitting, `Ref` creation, reassembly.
+- `src/refelim.jl` — `eliminate_refs`: live-variable analysis dropping `Ref` loads/stores not
+  needed across `produce` boundaries.
+- `src/bbcode.jl` — `BasicBlockCode`, the basic-block IR repr **copied from Mooncake.jl** to
+  avoid a dependency; keep in sync with Mooncake when fixing IR bugs.
+- `src/utils.jl` — IR helpers (`replace_captures`, `optimise_ir!`, `misty_closure`,
+  `Core.Compiler` version shims).
+- `src/test_utils.jl` — `TestUtils`: the `Testcase` driver.
+- `test/` — `runtests.jl` runs Aqua + `copyable_task.jl`; `test/integration/turing/` is a
+  separate-environment Turing.jl suite.
+- `docs/src/` — `index.md` (user) and `internals.md` (documented-internals docstring index).
+
+## Working Conventions
+
+- **Lives on `Core.Compiler` internals**, unstable across Julia versions — the most fragile
+  part of the codebase. Guard version-specific behaviour with `@static if VERSION ...` (e.g.
+  `set_valid_world!`, the `get_mi` shim); test on every supported version. The `[compat]`
+  range (`~1.10.8, 1.11.6`) and the CI matrix (`.github/workflows/CI.yml`) move together.
+- Fix IR problems by making the transform correct, not by special-casing the public API.
+  `transformation.jl`/`refelim.jl` are the single source of truth for the resumable IR.
+- `produce` and `get_taped_globals(::Type)` are **reserved statements**, not real functions:
+  `produce` is `@noinline` with a dummy side-effect so the transform can find/rewrite it (it
+  never runs under `consume`). Don't "optimise away" what resists inlining/constant-folding.
+- Concrete `Core.Compiler` types (`IRCode`, `SSAValue`, …) in the transform layer are
+  load-bearing, not over-constraint; annotate as specifically as dispatch requires.
+- **Workflow for any bug fix or new feature**: investigate → root-cause → understand the big
+  picture → *verify the fix before committing to it*. Reproduce with an MWE, inspect derived IR
+  (`Libtask.generate_ir`), and confirm the hypothesis by temporarily editing or monkey-patching
+  (e.g. `@eval`-ing a replacement method, hacking the local checkout) until the MWE behaves —
+  only then write the real change. Don't edit the transform on a guess. Prefer targeted fixes
+  over new helpers or refactors; run `minimise` before committing.
+- Write clear, local errors for unsupported inputs (naked `produce`, unhandled control flow) —
+  a clear `ArgumentError` at construction beats a confusing failure in compiled IR.
+- Internals may change freely; exports, the `produce`/`consume`/`copy` contract, and names in
+  `docs/src/internals.md` need tests, docs, and stable errors — update docs when they change.
+- Format with JuliaFormatter, **blue** style: `julia -e 'using JuliaFormatter; format(".")'`.
+
+## Concurrency
+
+- `build_callable` is serialized under `build_callable_lock` (`ReentrantLock`) — deliberate:
+  IR derivation mutates a shared `ID` counter (`seed_id!`/`ID()`) and the `MistyClosure` cache,
+  and concurrent derivation segfaulted (#227). Keep all ID generation and cache mutation inside
+  this lock; don't add a second IR-deriving code path outside it.
+- `bbcode.jl` indexes its `ID` counter by `Threads.threadid()` — suspect (task migration makes
+  IDs unstable); the lock, not `threadid()`, is what makes it safe today.
+- `copy(::TapedTask)` is `deepcopy`: the copy must share **no** mutable state (`Ref`s,
+  position). Any change to state storage must keep this — test that mutating a copy doesn't
+  affect the source.
+
+## Testing
+
+- MWE first, then the smallest focused test, then broader groups. Prefer extending existing
+  `Testcase`s over adding new ones; prune with `minimise`.
+- Register behaviour as `TestUtils.Testcase`s where possible — the driver iterates the task,
+  `copy`s after every iteration, and checks every copy resumes to the same result sequence
+  (plus optional allocation flags), exercising the copy/resume contract for free. Reserve
+  bespoke `test/` code for what the driver can't express.
+- Cover control-flow variety (branches, loops, nested produce, exception regions) and varied
+  value types. Bug fixes land with a focused regression test; version/world-age fixes get an
+  isolated direct test. Check `consume` allocations (`@allocated` post-warmup) and stability
+  (`@code_warntype`) for perf-sensitive paths.
+- Never disable tests or weaken perf/allocation assertions to get CI green — ask first.
+- `test/integration/turing/` runs in its own environment and is part of the contract: transform
+  or semantics changes may need updates there even when core tests pass.
+
+## Documentation
+
+`docs/make.jl` drives Documenter; `index.md` is user-facing, `internals.md` indexes documented
+internals. Update docstrings (and `internals.md` when adding/removing a documented internal)
+when changing the public API, transform internals, or developer helpers.

CLAUDE.md

@@ -0,0 +1 @@
+@AGENTS.md

Project.toml

@@ -3,7 +3,7 @@ uuid = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 license = "MIT"
 desc = "Tape based task copying in Turing"
 repo = "https://github.com/TuringLang/Libtask.jl.git"
-version = "0.9.18"
+version = "0.9.19"
 
 [deps]
 MistyClosures = "dbe65cb8-6be2-42dd-bbc5-4196aaced4f4"

docs/src/internals.md

@@ -5,11 +5,13 @@ Libtask.generate_ir
 Libtask.produce_value
 Libtask.is_produce_stmt
 Libtask.stmt_might_produce
+Libtask.assert_can_handle_control_flow
 Libtask.inc_args
 Libtask.get_type
 Libtask._typeof
 Libtask.replace_captures
 Libtask.BasicBlockCode
+Libtask.BasicBlockCode.IDEnterNode
 Libtask.opaque_closure
 Libtask.misty_closure
 Libtask.optimise_ir!

src/bbcode.jl

@@ -23,10 +23,13 @@ export ID,
     IDPhiNode,
     IDGotoNode,
     IDGotoIfNot,
+    IDEnterNode,
+    remake_enter,
     Switch,
     BBlock,
     phi_nodes,
     terminator,
+    is_implicit_terminator,
     insert_before_terminator!,
     collect_stmts,
     compute_all_predecessors,
@@ -87,6 +90,29 @@ end
 
 Base.copy(node::IDGotoIfNot) = IDGotoIfNot(copy(node.cond), copy(node.dest))
 
+"""
+    IDEnterNode
+
+Represent an exception or dynamic-scope `enter` using stable block IDs. `catch_id` is
+`nothing` for a scope-only `enter`; the optional `scope` field mirrors `Core.EnterNode`.
+"""
+struct IDEnterNode
+    catch_id::Union{ID,Nothing}
+    scope::Any
+    IDEnterNode(catch_id::Union{ID,Nothing}) = new(catch_id)
+    IDEnterNode(catch_id::Union{ID,Nothing}, @nospecialize(scope)) = new(catch_id, scope)
+end
+
+function remake_enter(node::IDEnterNode, catch_id::Union{ID,Nothing}, scope_map=identity)
+    return if isdefined(node, :scope)
+        IDEnterNode(catch_id, scope_map(node.scope))
+    else
+        IDEnterNode(catch_id)
+    end
+end
+
+Base.copy(node::IDEnterNode) = remake_enter(node, node.catch_id)
+
 struct Switch
     conds::Vector{Any}
     dests::Vector{ID}
@@ -99,12 +125,19 @@ end
 
 const Terminator = Union{Switch,IDGotoIfNot,IDGotoNode,ReturnNode}
 
+is_implicit_terminator(x) = x isa IDEnterNode || Meta.isexpr(x, :leave)
+
 mutable struct BBlock
     id::ID
     inst_ids::Vector{ID}
     insts::InstVector
     function BBlock(id::ID, inst_ids::Vector{ID}, insts::InstVector)
         @assert length(inst_ids) == length(insts)
+        if length(insts) > 1
+            @assert all(
+                inst -> !is_implicit_terminator(inst.stmt), @view(insts[begin:(end - 1)])
+            ) "`enter` and `:leave` must terminate their basic block"
+        end
         return new(id, inst_ids, insts)
     end
 end
@@ -192,22 +225,32 @@ compute_all_successors(ir::BBCode)::Dict{ID,Vector{ID}} = _compute_all_successor
         is_final_block = n == length(blks)
         t = terminator(blk)
         if t === nothing
-            return is_final_block ? ID[] : ID[blks[n + 1].id]
+            normal_succs = is_final_block ? ID[] : ID[blks[n + 1].id]
         elseif t isa IDGotoNode
-            return [t.label]
+            normal_succs = [t.label]
         elseif t isa IDGotoIfNot
-            return is_final_block ? ID[t.dest] : ID[t.dest, blks[n + 1].id]
+            normal_succs = is_final_block ? ID[t.dest] : ID[t.dest, blks[n + 1].id]
         elseif t isa ReturnNode
-            return ID[]
+            normal_succs = ID[]
         elseif t isa Switch
-            return vcat(t.dests, t.fallthrough_dest)
+            normal_succs = vcat(t.dests, t.fallthrough_dest)
         else
             error("Unhandled terminator $t")
         end
+        # `enter` has a catch edge in addition to its implicit fallthrough.
+        return _append_catch_edge(normal_succs, blk)
     end
     return Dict{ID,Vector{ID}}((b.id, succ) for (b, succ) in zip(blks, succs))
 end
 
+function _append_catch_edge(succs::Vector{ID}, blk::BBlock)
+    enter = blk.insts[end].stmt
+    if enter isa IDEnterNode && enter.catch_id isa ID && !in(enter.catch_id, succs)
+        push!(succs, enter.catch_id)
+    end
+    return succs
+end
+
 function compute_all_predecessors(ir::BBCode)::Dict{ID,Vector{ID}}
     return _compute_all_predecessors(ir.blocks)
 end
@@ -288,8 +331,17 @@ function __line_numbers_to_block_numbers!(insts::Vector{Any}, cfg::CC.CFG)
                 Int32[CC.block_for_inst(cfg, Int(edge)) for edge in stmt.edges], stmt.values
             )
         elseif Meta.isexpr(stmt, :enter)
-            stmt.args[1] = CC.block_for_inst(cfg, stmt.args[1]::Int)
+            # `catch_dest == 0` marks a scope-only `enter` with no catch block; leave it.
+            if stmt.args[1]::Int != 0
+                stmt.args[1] = CC.block_for_inst(cfg, stmt.args[1]::Int)
+            end
             insts[i] = stmt
+        else
+            @static if isdefined(Core, :EnterNode)
+                if isa(stmt, Core.EnterNode) && stmt.catch_dest != 0
+                    insts[i] = Core.EnterNode(stmt, CC.block_for_inst(cfg, stmt.catch_dest))
+                end
+            end
         end
     end
     return insts
@@ -349,6 +401,12 @@ function _ssa_to_ids(d::SSAToIdDict, x::PhiNode)
 end
 _ssa_to_ids(d::SSAToIdDict, x::GotoNode) = x
 _ssa_to_ids(d::SSAToIdDict, x::GotoIfNot) = GotoIfNot(get(d, x.cond, x.cond), x.dest)
+@static if isdefined(Core, :EnterNode)
+    function _ssa_to_ids(d::SSAToIdDict, x::Core.EnterNode)
+        isdefined(x, :scope) || return x
+        return Core.EnterNode(x.catch_dest, get(d, x.scope, x.scope))
+    end
+end
 
 function _block_nums_to_ids(insts::InstVector, cfg::CC.CFG)::Tuple{Vector{ID},InstVector}
     ids = map(_ -> ID(), cfg.blocks)
@@ -364,7 +422,23 @@ function _block_num_to_ids(d::BlockNumToIdDict, x::PhiNode)
 end
 _block_num_to_ids(d::BlockNumToIdDict, x::GotoNode) = IDGotoNode(d[x.label])
 _block_num_to_ids(d::BlockNumToIdDict, x::GotoIfNot) = IDGotoIfNot(x.cond, d[x.dest])
+function _block_num_to_ids(d::BlockNumToIdDict, x::Expr)
+    # On Julia 1.10 a `try` block begins with `Expr(:enter, catch_block_number)`; a
+    # `catch_block_number` of 0 is a scope-only `enter` with no catch block.
+    if Meta.isexpr(x, :enter)
+        catch_id = x.args[1]::Int == 0 ? nothing : d[x.args[1]::Int]
+        return IDEnterNode(catch_id)
+    end
+    return x
+end
 _block_num_to_ids(d::BlockNumToIdDict, x) = x
+@static if isdefined(Core, :EnterNode)
+    function _block_num_to_ids(d::BlockNumToIdDict, x::Core.EnterNode)
+        # `catch_dest == 0` is a scope-only `enter` (no catch block); record `nothing`.
+        catch_id = x.catch_dest == 0 ? nothing : d[x.catch_dest]
+        return isdefined(x, :scope) ? IDEnterNode(catch_id, x.scope) : IDEnterNode(catch_id)
+    end
+end
 
 # A map from IDs to IDs; useful when generically replacing IDs in BBCode statements
 const IdToIdDict = Dict{ID,ID}
@@ -393,6 +467,9 @@ replace_ids(d::IdToIdDict, x::IDGotoNode) = x
 function replace_ids(d::IdToIdDict, x::IDGotoIfNot)
     return IDGotoIfNot(get(d, x.cond, x.cond), get(d, x.dest, x.dest))
 end
+function replace_ids(d::IdToIdDict, x::IDEnterNode)
+    return remake_enter(x, get(d, x.catch_id, x.catch_id), s -> get(d, s, s))
+end
 function replace_ids(d::IdToIdDict, x::Switch)
     new_conds = Vector{Any}(undef, length(x.conds))
     for n in eachindex(x.conds)
@@ -510,6 +587,24 @@ function _to_ssas(d::Dict, x::IDPhiNode)
 end
 _to_ssas(d::Dict, x::IDGotoNode) = GotoNode(d[x.label].id)
 _to_ssas(d::Dict, x::IDGotoIfNot) = GotoIfNot(get(d, x.cond, x.cond), d[x.dest].id)
+@static if isdefined(Core, :EnterNode)
+    # `d[x.catch_id].id` is the SSAValue index of the catch block's first statement; it is
+    # mapped to a block number later by `__line_numbers_to_block_numbers!`, exactly as for
+    # the destination of an `IDGotoNode`. A scope-only `enter` (`catch_id === nothing`) maps
+    # back to `catch_dest == 0`.
+    function _to_ssas(d::Dict, x::IDEnterNode)
+        dest = x.catch_id === nothing ? 0 : d[x.catch_id].id
+        return if isdefined(x, :scope)
+            Core.EnterNode(dest, get(d, x.scope, x.scope))
+        else
+            Core.EnterNode(dest)
+        end
+    end
+else
+    function _to_ssas(d::Dict, x::IDEnterNode)
+        return Expr(:enter, x.catch_id === nothing ? 0 : d[x.catch_id].id)
+    end
+end
 
 function _remove_double_edges(ir::BBCode)
     new_blks = map(enumerate(ir.blocks)) do (n, blk)
@@ -607,6 +702,9 @@ end
 function _find_id_uses!(d::Dict{ID,Bool}, x::ReturnNode)
     return isdefined(x, :val) && in(x.val, keys(d)) && setindex!(d, true, x.val)
 end
+function _find_id_uses!(d::Dict{ID,Bool}, x::IDEnterNode)
+    return isdefined(x, :scope) && in(x.scope, keys(d)) && setindex!(d, true, x.scope)
+end
 _find_id_uses!(d::Dict{ID,Bool}, x::QuoteNode) = nothing
 _find_id_uses!(d::Dict{ID,Bool}, x) = nothing
 
@@ -644,6 +742,14 @@ function Base.show(io::IO, node::IDGotoNode)
     return print(io, "goto ", _block_str(node.label))
 end
 
+function Base.show(io::IO, node::IDEnterNode)
+    print(
+        io, "enter ", node.catch_id === nothing ? "(no catch)" : _block_str(node.catch_id)
+    )
+    isdefined(node, :scope) && print(io, " (scope ", _val_str(node.scope), ")")
+    return nothing
+end
+
 function Base.show(io::IO, node::IDGotoIfNot)
     return print(io, "goto ", _block_str(node.dest), " if not ", _val_str(node.cond))
 end

src/copyable_task.jl

@@ -138,7 +138,7 @@ function generate_ir(stage::Symbol, fargs...; kwargs...)
     seed_id!()
     original_bb = BBCode(original_ir)
     stage == :input_bb && return original_bb
-    transformed_bb, refs, _ = derive_copyable_task_ir(BBCode(original_ir))
+    transformed_bb, refs, _ = derive_copyable_task_ir(original_bb)
     stage == :transformed_bb && return transformed_bb
     topt_bb, refs = eliminate_refs(transformed_bb, refs)
     stage == :optimised_bb && return topt_bb
@@ -363,6 +363,23 @@ julia> consume(t)
 `Int`s have been used here, but it is permissible to set the value returned by
 [`Libtask.get_taped_globals`](@ref) to anything you like.
 
+## Exception handling and dynamic scopes
+
+Ordinary control flow -- branches, loops, and so on -- is fully supported inside a
+`TapedTask`. Exception handling and `Base.ScopedValues.@with` are supported with some
+restrictions because their runtime state does not survive when [`Libtask.produce`](@ref)
+suspends the task:
+
+  - `try` / `catch` / `finally` and `Base.ScopedValues.@with` are supported on Julia 1.12 and
+    later. A function containing one is rejected on earlier versions.
+  - [`Libtask.produce`](@ref) must not be reachable while an exception handler or dynamic
+    scope is active. Capture the required value within the region and call `produce` after it.
+  - Some catch handlers that read local values lower to `UpsilonNode` / `PhiCNode`. These IR
+    nodes are not yet supported and are rejected.
+
+In every unsupported case a clear `ArgumentError` is thrown when the `TapedTask` is
+constructed, rather than producing incorrect results.
+
 # Implementation Notes
 
 Under the hood, we implement a `TapedTask` by obtaining the `IRCode` associated to the