diff --git a/docs/src/safety/compile-time.md b/docs/src/safety/compile-time.md index fa978418..07390b88 100644 --- a/docs/src/safety/compile-time.md +++ b/docs/src/safety/compile-time.md @@ -67,10 +67,10 @@ in expression starting at myfile.jl:1 ### Incidental Tail Patterns -Beyond the direct-return patterns above, three additional tail shapes are +Beyond the direct-return patterns above, four additional tail shapes are checked — they don't *look* like a `return`, but the scope's last expression still evaluates to a pool-backed array, which then escapes as the scope's -value: +value (also when wrapped in `identity(...)`): ```julia @with_pool pool begin @@ -87,6 +87,10 @@ end v = acquire!(pool, Float64, 100) y = v # ← assignment tail end + +@with_pool pool begin + (acquire!(pool, Float64, 100), 1) # ← literal holding a direct acquire call +end ``` Fix: end the block with `nothing` if the value is meant to be discarded — @@ -145,3 +149,26 @@ end ``` That's what the [runtime layer](runtime.md) is for. + +Loops and comprehensions bind *elements*, not the iterable. Reading an element of +a pool-backed array is a plain value, so this is silent and correct: + +```julia +@with_pool pool begin + order = acquire!(pool, Int, n) + contact = 0 + for k in order # `k` is an element (e.g. an Int), not pool memory + k == target && (contact = k) + end + (; contact) # fine — no warning, no guard +end +``` + +Iterating a container *literal* of pool arrays (`for c in (v, w)`) is tracked: +`c` is a pool-backed array inside the loop. Iterating a container held in a +*variable* (`t = (v, w); for c in t`) and destructuring iteration +(`for (a, b) in ...`) are not tracked at compile time. With the default +`runtime_check = 0`, nothing guards those shapes. With `runtime_check = 1`, the +[runtime layer](runtime.md) inspects the scope's return value at exit and throws +`PoolRuntimeEscapeError` if it still references pool memory — including through +a `Tuple` or `NamedTuple` — on every supported Julia version. diff --git a/ext/AdaptiveArrayPoolsCUDAExt/debug.jl b/ext/AdaptiveArrayPoolsCUDAExt/debug.jl index 39324dcc..221ab3c8 100644 --- a/ext/AdaptiveArrayPoolsCUDAExt/debug.jl +++ b/ext/AdaptiveArrayPoolsCUDAExt/debug.jl @@ -310,11 +310,11 @@ end ) calls = [ :( - _check_tp_cuda_overlap( - getfield(pool, $(QuoteNode(f))), arr_ptr, arr_end, - current_depth, pool, return_site, original_val - ) + _check_tp_cuda_overlap( + getfield(pool, $(QuoteNode(f))), arr_ptr, arr_end, + current_depth, pool, return_site, original_val ) + ) for f in GPU_FIXED_SLOT_FIELDS ] return quote diff --git a/ext/AdaptiveArrayPoolsMetalExt/debug.jl b/ext/AdaptiveArrayPoolsMetalExt/debug.jl index a1ddc707..3194a5c9 100644 --- a/ext/AdaptiveArrayPoolsMetalExt/debug.jl +++ b/ext/AdaptiveArrayPoolsMetalExt/debug.jl @@ -310,11 +310,11 @@ end ) calls = [ :( - _check_tp_metal_overlap( - getfield(pool, $(QuoteNode(f))), arr_buf, arr_off, arr_end, - current_depth, pool, return_site, original_val - ) + _check_tp_metal_overlap( + getfield(pool, $(QuoteNode(f))), arr_buf, arr_off, arr_end, + current_depth, pool, return_site, original_val ) + ) for f in METAL_FIXED_SLOT_FIELDS ] return quote diff --git a/src/debug.jl b/src/debug.jl index 636441de..90907a3f 100644 --- a/src/debug.jl +++ b/src/debug.jl @@ -193,11 +193,11 @@ end ) calls = [ :( - _check_tp_pointer_overlap( - getfield(pool, $(QuoteNode(f))), arr_ptr, arr_end, - current_depth, pool, return_site, original_val - ) + _check_tp_pointer_overlap( + getfield(pool, $(QuoteNode(f))), arr_ptr, arr_end, + current_depth, pool, return_site, original_val ) + ) for f in FIXED_SLOT_FIELDS ] return quote diff --git a/src/macros.jl b/src/macros.jl index 206964aa..fad86e43 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -15,7 +15,8 @@ const _MACRO_TYPED_UPGRADES = VERSION >= v"1.12-" """Per-return-point escape detail: which expression, at which line, leaks which vars. `incidental` carries the `(kind, detail)` pair for a Stage-2 incidental-tail escape -(a direct acquire call / broadcast-assign / assignment tail — see `_incidental_exposure`) +(a direct acquire call / broadcast-assign / assignment tail, or a literal holding a +direct acquire call — see `_incidental_exposure`) and is `nothing` for a Stage-1 intentional-return escape. Storing it at throw time lets `showerror` render the message without re-classifying the expression, and is what distinguishes the two error layouts (no `vars`-emptiness sentinel).""" @@ -393,9 +394,9 @@ first actual use throws an `EscapedPoolUseError` carrying the variable name, shape, and scope location. Detected patterns cover the direct-return shapes (a bare variable, `return v`, -container literals like `(v, w)` / `[v]`), and three additional "incidental" -tail patterns that don't *look* like a `return` but still expose a pool-backed -array as the scope's value: +container literals like `(v, w)` / `[v]` / `(; v)`), and four additional +"incidental" tail patterns that don't *look* like a `return` but still expose a +pool-backed array as the scope's value (also when wrapped in `identity(...)`): ```julia @with_pool pool begin @@ -412,6 +413,10 @@ end v = acquire!(pool, Float64, 100) y = v # ← assignment tail (evaluates to its RHS) end + +@with_pool pool begin + (acquire!(pool, Float64, 100), 1) # ← literal holding a direct acquire call +end ``` If the scope's last expression is meant to be discarded (a "run it for its @@ -2031,8 +2036,9 @@ function _transform_return_stmts( ) expr isa Expr || return expr - # Don't recurse into nested function definitions or quoted AST - if expr.head in (:function, :->, :quote) + # Don't recurse into nested function definitions (long- or short-form) or quoted + # AST — a `return` there exits the nested function, never this scope + if expr.head in (:function, :->, :quote) || (expr.head == :(=) && _is_function_def(expr)) return expr end @@ -2114,9 +2120,9 @@ function _transform_return_stmts( # Other expressions: recurse with current_lnn new_args = Any[ _transform_return_stmts( - arg, pool_name, current_lnn; - rewind_call, entry_depth_guard - ) for arg in expr.args + arg, pool_name, current_lnn; + rewind_call, entry_depth_guard + ) for arg in expr.args ] return Expr(expr.head, new_args...) end @@ -2352,6 +2358,9 @@ function _extract_acquired_vars(expr, target_pool, vars = Set{Symbol}()) for arg in expr.args _extract_acquired_vars(arg, target_pool, vars) end + elseif expr.head == :(=) && _is_function_def(expr) + # Short-form `f(x) = …` is a nested scope, skipped like long-form below + return vars elseif expr.head == :(=) && length(expr.args) >= 2 lhs = expr.args[1] rhs = expr.args[2] @@ -2376,6 +2385,38 @@ function _extract_acquired_vars(expr, target_pool, vars = Set{Symbol}()) end # Recurse into RHS (for nested blocks with acquire calls) _extract_acquired_vars(rhs, target_pool, vars) + elseif expr.head == :for && length(expr.args) >= 2 + # `for k in x` parses as `Expr(:for, :(k = x), body)` — the spec binds + # an element, it is not an alias assignment + shadowed = _bind_iteration_spec!(vars, expr.args[1], target_pool) + _extract_acquired_vars(expr.args[2], target_pool, vars) + union!(vars, shadowed) + elseif expr.head == :generator && length(expr.args) >= 2 + # Specs first (they bind before the body runs; for `:flatten` the inner + # generator IS this body), then filter conditions, then the body + shadowed = Symbol[] + for spec in expr.args[2:end] + if Meta.isexpr(spec, :filter) && length(spec.args) >= 2 + for s in spec.args[2:end] + append!(shadowed, _bind_iteration_spec!(vars, s, target_pool)) + end + _extract_acquired_vars(spec.args[1], target_pool, vars) + else + append!(shadowed, _bind_iteration_spec!(vars, spec, target_pool)) + end + end + _extract_acquired_vars(expr.args[1], target_pool, vars) + union!(vars, shadowed) + elseif expr.head == :tuple + # A `name = value` child of a tuple is a NamedTuple field (lowering rejects + # `((x = 1), 2)`), never an assignment — only its value is walked + for arg in expr.args + if Meta.isexpr(arg, :(=)) && length(arg.args) >= 2 + _extract_acquired_vars(arg.args[2], target_pool, vars) + else + _extract_acquired_vars(arg, target_pool, vars) + end + end else # Skip scope-introducing expressions whose `return` does NOT # return from the enclosing function (function, ->, macro). @@ -2392,6 +2433,50 @@ function _extract_acquired_vars(expr, target_pool, vars = Set{Symbol}()) return vars end +""" + _bind_iteration_spec!(vars, spec, target_pool) -> Vector{Symbol} + +Apply one `for`/generator iteration spec (`lhs = iter`, or a `:block` of them for +`for a in x, b in y`) to the taint set. Iterating yields the iterable's ELEMENTS: +an element of a pool-backed array is a value, so a Symbol iterable never taints +the loop variable. Only a container literal exposing a tainted var +(`for c in (v, w)`) yields pool-backed elements and taints it. + +A non-`outer` loop variable is a fresh local, so an outer binding of the same +name is un-tainted for the body and returned, for the caller to restore after +the body. `outer` never un-taints (a zero-trip loop keeps the old value), and a +destructuring `lhs` is not tracked. + +A literal-tainted loop variable deliberately stays tainted after the loop: the +escape and mutation checkers are flow-insensitive over this one name set, so +un-tainting it would hide a `return c` or `resize!(c, …)` INSIDE the loop. The +cost is a conservative (loud) false positive when a same-named outer variable is +used after the loop. +""" +function _bind_iteration_spec!(vars, spec, target_pool) + if Meta.isexpr(spec, :block) + shadowed = Symbol[] + for s in spec.args + s isa LineNumberNode && continue + append!(shadowed, _bind_iteration_spec!(vars, s, target_pool)) + end + return shadowed + end + (Meta.isexpr(spec, :(=)) && length(spec.args) >= 2) || return Symbol[] + lhs, iter = spec.args[1], spec.args[2] + _extract_acquired_vars(iter, target_pool, vars) # nested acquires in the iterable + outer = Meta.isexpr(lhs, :outer) + name = outer ? lhs.args[1] : lhs + name isa Symbol || return Symbol[] + if Meta.isexpr(iter, (:tuple, :vect)) && _literal_contains_acquired(iter, vars) + push!(vars, name) + elseif !outer && name in vars + delete!(vars, name) + return [name] + end + return Symbol[] +end + """ _get_last_expression(expr) -> Any @@ -2432,6 +2517,7 @@ Tracks LineNumberNodes through blocks. Skips nested function definitions.""" function _collect_explicit_returns!(values, expr, current_line::Union{Int, Nothing}) expr isa Expr || return expr.head in (:function, :(->)) && return + expr.head == :(=) && _is_function_def(expr) && return # short-form `f(x) = …` if expr.head == :return push!(values, (expr, current_line)) return @@ -2504,7 +2590,9 @@ pass over top-level block statements, correctly handling statement ordering. For top-level assignments, taint is added/removed as statements are scanned in source order. For nested blocks (if/for/while), recursion adds taint -conservatively (never removes inside branches — can't resolve control flow). +conservatively (never removes inside branches — can't resolve control flow); +`for`/generator iteration specs bind elements, not aliases (see +`_bind_iteration_spec!`). Tuple destructuring is handled atomically: all RHS values are evaluated against the current taint set before any LHS updates are applied. @@ -2517,6 +2605,8 @@ function _extract_ordered_acquired(expr, target_pool) tainted = Set{Symbol}() for arg in expr.args arg isa LineNumberNode && continue + # Short-form `f(x) = …` is a nested scope: its body taints nothing here + Meta.isexpr(arg, :(=)) && _is_function_def(arg) && continue if arg isa Expr && arg.head == :(=) && length(arg.args) >= 2 lhs = arg.args[1] rhs = arg.args[2] @@ -2574,12 +2664,12 @@ Only catches high-confidence patterns: - Bare Symbol: `v` - Explicit return: `return v` - Tuple/array literal containing a var: `(v, w)`, `[v, w]` -- NamedTuple-style kw: `(a=v,)` +- NamedTuple-style kw: `(a=v,)`, `(; a=v)`, and the shorthand `(; v)` Does NOT recurse into function calls (can't know what `f(v)` returns). """ -"""Check if a literal expression (symbol, identity call, tuple/vect) transitively contains any acquired var.""" +"""Check if a literal expression (symbol, identity call, tuple/vect, `(; …)` parameters) transitively contains any acquired var.""" function _literal_contains_acquired(expr, acquired) expr isa Symbol && return expr in acquired if expr isa Expr @@ -2587,7 +2677,7 @@ function _literal_contains_acquired(expr, acquired) if expr.head == :call && length(expr.args) >= 2 && expr.args[1] === :identity return _literal_contains_acquired(expr.args[2], acquired) end - if expr.head in (:tuple, :vect) + if expr.head in (:tuple, :vect, :parameters) for arg in expr.args if Meta.isexpr(arg, :(=)) && length(arg.args) >= 2 _literal_contains_acquired(arg.args[2], acquired) && return true @@ -2606,6 +2696,14 @@ end _is_identity_call(x) = x === :identity || (x isa Expr && x.head == :. && x.args == [:Base, QuoteNode(:identity)]) +"""Strip any `identity(...)` / `Base.identity(...)` wrappers — they are transparent.""" +function _unwrap_identity(x) + while x isa Expr && x.head == :call && length(x.args) >= 2 && _is_identity_call(x.args[1]) + x = x.args[2] + end + return x +end + function _find_direct_exposure(expr, acquired) found = Set{Symbol}() if expr isa Symbol @@ -2630,10 +2728,12 @@ function _find_direct_exposure(expr, acquired) end end elseif expr.head == :parameters - # (; a=v) style named tuple parameters + # (; a=v) style named tuple parameters, and the (; v) shorthand for arg in expr.args if Meta.isexpr(arg, :kw) && length(arg.args) >= 2 union!(found, _find_direct_exposure(arg.args[2], acquired)) + elseif arg isa Symbol + union!(found, _find_direct_exposure(arg, acquired)) end end elseif expr.head == :call && length(expr.args) >= 2 && _is_identity_call(expr.args[1]) @@ -2654,7 +2754,7 @@ end """ _incidental_exposure(expr, tainted, pool_name) -> Union{Nothing, Tuple{Symbol, Any}} -Detect the three incidental pool-backed tail patterns the direct-exposure ERROR +Detect the incidental pool-backed tail patterns the direct-exposure ERROR does not cover (their value escapes as the block's return value, but the syntax does not *look* like a return): @@ -2663,6 +2763,10 @@ does not *look* like a return): (also `x[...] .= v` on an acquired base) - `(:assign, var)` — tail is `x = ` or `x = ` (assignment evaluates to its RHS value) +- `(:literal_acquire_call, expr)` — tail is a tuple/vect/NamedTuple literal holding + a direct acquire call (`(acquire!(pool, T, n),)`); + a literal that also exposes an acquired var is left + to the direct-exposure stage, so it warns once `ret_expr` values arrive here either as the bare tail expression (implicit return) or as the whole `Expr(:return, ...)` node (explicit `return`, per @@ -2677,8 +2781,14 @@ function _incidental_exposure(expr, tainted, pool_name) if expr.head == :return && !isempty(expr.args) return _incidental_exposure(expr.args[1], tainted, pool_name) end + # identity(x) is transparent (as in `_find_direct_exposure` and the poisoner) + unwrapped = _unwrap_identity(expr) + unwrapped !== expr && return _incidental_exposure(unwrapped, tainted, pool_name) if _is_acquire_call(expr, pool_name) return (:acquire_call, expr) + elseif expr.head in (:tuple, :vect) && _tuple_contains_acquire_call(expr, pool_name) && + isempty(_find_direct_exposure(expr, tainted)) + return (:literal_acquire_call, expr) elseif _is_dotted_assign_head(expr.head) && length(expr.args) >= 1 lhs = expr.args[1] base = Meta.isexpr(lhs, :ref) ? lhs.args[1] : lhs @@ -2721,6 +2831,8 @@ function _lint_message(kind, detail, ret_expr; guarded::Bool = false) tail = sprint(Base.show_unquoted, ret_expr) what = if kind === :acquire_call "is a direct acquire call whose pool-backed array" + elseif kind === :literal_acquire_call + "is a container literal holding a direct acquire call, whose pool-backed array" elseif kind === :broadcast_assign "writes into the pool-backed array `$(detail)`, whose value" elseif kind === :bare_var @@ -2776,9 +2888,10 @@ assignment, acquire) still execute; only the *yielded value* changes. Explicit `return` nodes are never rewritten (they are hard errors in `_check_compile_time_escape`). Applied only under `ESCAPE_LINT == "warn"`, in lockstep with the checker's warn diagnostics: the tail shapes rewritten here -mirror `_find_direct_exposure` (bare var, tuple/vect/NamedTuple literals, -`identity` unwrap) and `_incidental_exposure` (acquire-call tail, -broadcast-assign tail, assignment tail). +mirror `_find_direct_exposure` (bare var, tuple/vect/NamedTuple literals incl. +the `(; v)` shorthand, `identity` unwrap) and `_incidental_exposure` (acquire-call +tail, broadcast-assign tail, assignment tail, literal holding an acquire call, +`identity` unwrap). The "checker ⇔ poisoner lockstep" test pins this agreement. """ function _poison_block_tails(expr, pool_name, source::Union{LineNumberNode, Nothing}) acquired = _extract_ordered_acquired(expr, pool_name) @@ -2865,6 +2978,9 @@ function _poison_literal(expr, acquired, pool_name, file, line) arg.head, arg.args[1], _poison_literal_value(arg.args[2], acquired, pool_name, file, line) ) + elseif expr.head === :parameters && arg isa Symbol && arg in acquired + # `(; v)` shorthand slots accept only a name, never a call: expand to `v = guard` + Expr(:kw, arg, _guard_expr(arg, arg, file, line)) else _poison_literal_value(arg, acquired, pool_name, file, line) end @@ -2903,7 +3019,7 @@ function _collect_acquired_in_literal!(found, expr, acquired_keys) elseif expr isa Expr if expr.head == :call && length(expr.args) >= 2 && expr.args[1] === :identity _collect_acquired_in_literal!(found, expr.args[2], acquired_keys) - elseif expr.head in (:tuple, :vect) + elseif expr.head in (:tuple, :vect, :parameters) for arg in expr.args if Meta.isexpr(arg, :(=)) && length(arg.args) >= 2 _collect_acquired_in_literal!(found, arg.args[2], acquired_keys) @@ -3024,17 +3140,19 @@ Checks if the block/function body's return expression directly contains a pool-backed variable. This catches the most common beginner mistake at zero runtime cost. -Stage 1 — bare symbol (`v`), `return v`, and container patterns (`(v, w)`, -`[v]`, `(key=v,)`) — always throws `PoolEscapeError`. Stage 2 — the three -incidental-tail patterns (direct acquire-call tail, broadcast-assign tail, -assignment tail) — resolves severity by form: with `function_form = true` -(the scope is a function definition, so the tail IS the function's return -value), or when the tail is an explicit `return` (which returns from the -enclosing function even in block form — a `begin` block is not a function -boundary), the escape is definite and always throws. A block-form implicit -tail — whose value may simply be discarded by the surrounding code — is -reported at the `ESCAPE_LINT` severity instead (default `"warn"`; `"error"` -promotes it, `"off"` disables Stage 2 entirely). +Stage 1 — direct exposure: bare symbol (`v`), `return v`, and container +patterns (`(v, w)`, `[v]`, `(key=v,)`, `(; v)`). Stage 2 — the incidental-tail +patterns (direct acquire-call tail, broadcast-assign tail, assignment tail, +literal holding a direct acquire call). Both stages look through `identity(...)` +and resolve severity by form: with `function_form = true` (the scope is a +function definition, so the tail IS the function's return value), or when the +tail is an explicit `return` (which returns from the enclosing function even in +block form — a `begin` block is not a function boundary), the escape is definite +and always throws `PoolEscapeError`. A block-form implicit tail — whose value +may simply be discarded by the surrounding code — is reported at the +`ESCAPE_LINT` severity instead (default `"warn"`: `@warn` + `EscapedPoolArray` +guard; `"error"` throws; `"off"` skips Stage-1 block tails and disables Stage 2 +entirely). Skipped when `STATIC_POOLING = false` (pooling disabled, acquire returns normal arrays). """ @@ -3129,25 +3247,21 @@ end """ _tuple_contains_acquire_call(expr, target_pool) -> Bool -Check if a tuple/vect literal contains any acquire calls as values. -Unlike `_literal_contains_acquired`, this checks the RHS calls directly, -so it works even when `vars` is empty (solves the chicken-and-egg problem). +Check if a tuple/vect/NamedTuple literal (`(a, b)`, `[a]`, `(k = a,)`, `(; k = a)`) +contains any direct acquire call as an element, looking through `identity(...)` +and nested literals. Unlike `_literal_contains_acquired`, this checks the RHS +calls directly, so it works even when `vars` is empty (solves the +chicken-and-egg problem). """ function _tuple_contains_acquire_call(expr, target_pool) expr isa Expr || return false - if expr.head in (:tuple, :vect) + if expr.head in (:tuple, :vect, :parameters) for arg in expr.args - if Meta.isexpr(arg, :(=)) && length(arg.args) >= 2 - # NamedTuple value: (wv=zeros!(pool, ...), ...) - _is_acquire_call(arg.args[2], target_pool) && return true - _tuple_contains_acquire_call(arg.args[2], target_pool) && return true - elseif Meta.isexpr(arg, :kw) && length(arg.args) >= 2 - _is_acquire_call(arg.args[2], target_pool) && return true - _tuple_contains_acquire_call(arg.args[2], target_pool) && return true - else - _is_acquire_call(arg, target_pool) && return true - _tuple_contains_acquire_call(arg, target_pool) && return true - end + # NamedTuple value: (wv=zeros!(pool, ...), ...) / (; wv=zeros!(pool, ...)) + val = Meta.isexpr(arg, (:(=), :kw)) && length(arg.args) >= 2 ? arg.args[2] : arg + val = _unwrap_identity(val) + _is_acquire_call(val, target_pool) && return true + _tuple_contains_acquire_call(val, target_pool) && return true end end return false @@ -3829,8 +3943,9 @@ end function _find_mutation_calls!(points, expr, acquired, current_line::Union{Int, Nothing}) expr isa Expr || return - # Skip nested function definitions (they have their own scope) + # Skip nested function definitions (they have their own scope), long- or short-form expr.head in (:function, :(->)) && return + expr.head == :(=) && _is_function_def(expr) && return # Track line numbers if expr.head == :block diff --git a/test/test_compile_escape.jl b/test/test_compile_escape.jl index 0f5f913a..10614a6e 100644 --- a/test/test_compile_escape.jl +++ b/test/test_compile_escape.jl @@ -13,7 +13,7 @@ import AdaptiveArrayPools: _extract_acquired_vars, _get_last_expression, _extract_container_vars, _is_dotted_assign_head, _incidental_exposure, _lint_message, ESCAPE_LINT, _report_incidental_escape, PoolEscapeError, - _poison_block_tails, EscapedPoolArray, EscapedPoolUseError + _poison_block_tails, EscapedPoolArray, EscapedPoolUseError, _is_function_def # Helper: expansion must throw PoolEscapeError (possibly LoadError-wrapped) — # the severity of function-form and explicit-return escapes. @@ -37,6 +37,19 @@ function _expansion_incidental_warns(ex) return expanded isa Expr end +# Lockstep helpers: did block-form expansion emit the escape @warn, and does the +# block-tail poisoner change the block? The two must always agree. +function _block_tail_warned(body) + logs, _ = Test.collect_test_logs() do + macroexpand(@__MODULE__, :(@with_pool pool $body)) + end + return any( + l -> l.level == Base.CoreLogging.Warn && + occursin("becomes the scope's return value", string(l.message)), logs + ) +end +_block_tail_rewritten(body) = _poison_block_tails(body, :pool, LineNumberNode(1, :lockstep)) != body + function _capture_stderr(f) tmpf = tempname() open(tmpf, "w") do io @@ -244,6 +257,13 @@ end @test :v in found @test :w in found + # NamedTuple shorthand: (; s, v) — the bare name IS the value + @test _find_direct_exposure(:((; s, v)), acquired) == Set([:v]) + @test _find_direct_exposure(:((; s, v = v, w)), acquired) == Set([:v, :w]) + # field-access shorthand and non-pool shorthand → not detected + @test isempty(_find_direct_exposure(:((; s, t.v)), acquired)) + @test isempty(_find_direct_exposure(:((; s)), acquired)) + # Scalar literal → not detected @test isempty(_find_direct_exposure(42, acquired)) @test isempty(_find_direct_exposure(3.14, acquired)) @@ -782,6 +802,18 @@ end sum(v) end + # Element read from a pool-array loop, then returned → safe (an element is a value) + @test_nowarn @macroexpand @with_pool pool begin + order = acquire!(pool, Int, 4) + contact = 0 + for k in order + if k == 2 + contact = k + end + end + contact + end + # Indexing → safe (element access, not array itself) @test_nowarn @macroexpand @with_pool pool begin v = acquire!(pool, Float64, 10) @@ -1212,6 +1244,286 @@ end @test :v in vars end + # ============================================================================== + # Binding positions: `:(=)` nodes that are NOT assignments (iteration specs, + # NamedTuple fields) and loop variables that shadow an outer name + # ============================================================================== + + @testset "binding positions: iteration specs, shadowing, NamedTuple fields" begin + taint(b) = sort!(collect(_extract_ordered_acquired(b, :pool))) + + # iterating a pool array yields ELEMENTS — the loop variable is not an alias + @test taint( + quote + order = acquire!(pool, Int, n); for k in order + end + end + ) == [:order] + @test taint( + quote + order = acquire!(pool, Int, n); contact = 0 + for k in order + if k == 2 + contact = k + end + end + contact + end + ) == [:order] + @test taint( + quote + order = acquire!(pool, Int, n); for i in 1:n, k in order + c = k + end + end + ) == [:order] + # generator / filtered / typed / flattened comprehension specs + @test taint( + quote + order = acquire!(pool, Int, n); k = 0; ks = [k for k in order if k > 0]; k + end + ) == [:order] + @test taint( + quote + order = acquire!(pool, Int, n); k = 0; g = sum(k for k in order); k + end + ) == [:order] + @test taint( + quote + order = acquire!(pool, Int, n); k = 0; ks = Int[k for k in order]; k + end + ) == [:order] + @test taint( + quote + order = acquire!(pool, Int, n); k = 0; ks = [k for k in order for j in 1:2]; k + end + ) == [:order] + # a `name = value` child of a tuple is a NamedTuple field, not an assignment + @test taint( + quote + order = acquire!(pool, Int, n); k = 0; s = sum((k = order,).k); k + end + ) == [:order] + @test taint( + quote + v = acquire!(pool, Int, n); nt = (a = v,) + end + ) == [:nt, :v] + # shadowing: the loop-local `k` is a scalar; the outer `k` stays tainted + @test taint( + quote + v = acquire!(pool, Int, n); order = acquire!(pool, Int, n); k = v; contact = 0 + for k in order + contact = k + end + end + ) == [:k, :order, :v] + # generator specs are bound before the body runs + @test taint( + quote + v = acquire!(pool, Int, n); best = nothing + s = sum((best = c) for c in (v,)) + end + ) == [:best, :c, :v] + # Known limitation (conservative, loud): a literal-tainted loop variable stays + # tainted after the loop, so a same-named OUTER variable used afterwards is + # flagged — the price of the in-loop checks below + @test taint( + quote + v = acquire!(pool, Int, n); c = 1.0 + for c in (v,) + sum(c) + end + c + end + ) == [:c, :v] + # the escape/mutation checkers are flow-insensitive over one taint set, so the + # loop variable must stay tainted for them: a `return c` or a structural + # mutation of `c` INSIDE a loop over a literal of pool arrays is still caught + @test _expansion_escape_error( + :( + @with_pool pool begin + v = acquire!(pool, Float64, 3) + for c in (v,) + return c + end + nothing + end + ) + ) + mut_err = try + macroexpand( + @__MODULE__, :( + @with_pool pool begin + v = acquire!(pool, Float64, 3) + for c in (v,) + resize!(c, 0) + end + nothing + end + ) + ) + nothing + catch e + e isa LoadError ? e.error : e + end + @test mut_err isa AdaptiveArrayPools.PoolMutationError + # `outer` loop variable over a literal of pool arrays outlives the loop + @test taint( + quote + v = acquire!(pool, Int, n); for outer c in (v,) + end + end + ) == [:c, :v] + # Known limitation: iterating a container VARIABLE is treated as element + # iteration, so `c` is not tracked here (a container LITERAL is, below) + @test taint( + quote + v = acquire!(pool, Int, n); t = (v,); for c in t + end + end + ) == [:t, :v] + + # true positives preserved + @test taint( + quote + v = acquire!(pool, Int, n); w = acquire!(pool, Int, n) + for c in (v, w) + best = c + end + end + ) == [:best, :c, :v, :w] + @test taint( + quote + v = acquire!(pool, Int, n); for c in [v] + d = c + end + end + ) == [:c, :d, :v] + @test taint( + quote + order = acquire!(pool, Int, n); for outer k in order + end + end + ) == [:order] + @test taint( + quote + v = acquire!(pool, Int, n); let a = v + a + end + end + ) == [:a, :v] + @test :w in taint( + quote + for k in (w = acquire!(pool, Int, n)) + end + end + ) + # Known limitation (pre-existing): destructuring iteration is not tracked + @test taint( + quote + v = acquire!(pool, Int, n); for (c,) in ((v,),) + best = c + end + end + ) == [:v] + end + + @testset "short-form nested definitions are skipped like long-form" begin + taint(b) = sort!(collect(_extract_ordered_acquired(b, :pool))) + + # a short-form definition's body is its own scope, not the enclosing block + @test taint( + quote + v = acquire!(pool, Int, n); f(x) = (a = v) + end + ) == [:v] + @test taint( + quote + v = acquire!(pool, Int, n); f(x::T) where {T} = (a = v) + end + ) == [:v] + @test taint( + quote + v = acquire!(pool, Int, n); f(x)::Int = (a = v; 1) + end + ) == [:v] + + # hard-error false positives removed (long-form `function g … end` already passes) + @test !_expansion_escape_error( + :( + @with_pool pool begin + v = acquire!(pool, Float64, 3) + g(x) = (return v) + nothing + end + ) + ) + @test ( + macroexpand( + @__MODULE__, :( + @with_pool pool begin + v = acquire!(pool, Float64, 3) + g() = resize!(v, 0) + nothing + end + ) + ); true + ) # no PoolMutationError + + # runtime: a `return` inside a short-form nested def must NOT carry the scope's + # rewind (regression: calling `g` rewound the pool mid-scope, so the next + # acquire reused `v`'s live memory) — block form, function form, and the + # `g() = (return v)` shape that used to be a hard PoolEscapeError + r = Core.eval( + @__MODULE__, :( + @with_pool pool begin + v = acquire!(pool, Float64, 3) + v .= 1 + g(x) = (return x + 1) + y = g(1) + w = acquire!(pool, Float64, 3) + w .= 2 + (sum(v), y) + end + ) + ) + @test r == (3.0, 2) + f = Core.eval( + @__MODULE__, :( + @with_pool pool function short_def_return_fn() + v = acquire!(pool, Float64, 3) + v .= 1 + g(x) = (return x + 1) + y = g(1) + w = acquire!(pool, Float64, 3) + w .= 2 + sum(v) + y + end + ) + ) + @test Base.invokelatest(f) == 5.0 + r = Core.eval( + @__MODULE__, :( + @with_pool pool begin + v = acquire!(pool, Float64, 3) + v .= 1 + g() = (return v) + u = g() + w = acquire!(pool, Float64, 3) + w .= 2 + sum(v) + end + ) + ) + @test r == 3.0 + + # typed / indexed assignments are not definitions + @test !_is_function_def(:(a::Vector{Int} = v)) + @test !_is_function_def(:(a[i] = v)) + @test _is_function_def(:((f::Foo)(x) = 1)) + end + # ============================================================================== # Order-aware escape detection through macro pipeline # ============================================================================== @@ -1313,6 +1625,20 @@ end v = acquire!(pool, Float64, n) return (v, n) end + + # NamedTuple shorthand tail → escape + @test_throws PoolEscapeError @macroexpand @with_pool pool function fn_nt_short(n) + v = acquire!(pool, Float64, n) + s = 0.0 + (; s, v) + end + + # Semicolon-NamedTuple container, then returned → escape + @test_throws PoolEscapeError @macroexpand @with_pool pool function fn_nt_container(n) + v = acquire!(pool, Float64, n) + d = (; a = v) + d + end end @testset "Function form: safe patterns (no false positives)" begin @@ -2182,6 +2508,11 @@ end @test !_literal_contains_acquired( Expr(:tuple, Expr(:kw, :a, :x)), acquired ) + + # Semicolon NamedTuple: shorthand (; v) and explicit (; a = v) + @test _literal_contains_acquired(:((; v)), acquired) + @test _literal_contains_acquired(:((; a = v)), acquired) + @test !_literal_contains_acquired(:((; a = 1, s)), acquired) end # ============================================================================== @@ -2207,6 +2538,10 @@ end ) @test :w in found + # Semicolon NamedTuple: shorthand + explicit field + found = _collect_acquired_in_literal(:((; s, v, a = w)), acquired) + @test sort!(found) == [:v, :w] + # Non-Expr/non-Symbol → empty found = _collect_acquired_in_literal(42, acquired) @test isempty(found) @@ -2489,11 +2824,20 @@ end end end @test !contains(warn_output2, "PoolContainerEscapeWarning") + + # semicolon-NamedTuple container of an acquire call is a container too + @test _extract_container_vars( + quote + vac = (; wv = zeros!(pool, Float64, 3)) + vac.wv + end, :pool + ) == Set([:vac]) end # ============================================================================== # Incidental-tail escape detection: direct acquire-call tails, broadcast-assign - # tails, and assignment tails. Form-based severity: a function-form tail or an + # tails, assignment tails, and literals holding a direct acquire call. + # Form-based severity: a function-form tail or an # explicit `return` (definite escape to the enclosing function's caller) always # errors; a block-form implicit tail (value may be discarded) reports at the # ESCAPE_LINT severity — "warn" by default. @@ -2599,6 +2943,24 @@ end ) ) + # function form, w4: acquire call wrapped in a container literal + @test _expansion_escape_error( + :( + @with_pool pool function f() + (acquire!(pool, Float64, 2),) + end + ) + ) + + # function form: identity-wrapped acquire-call tail + @test _expansion_escape_error( + :( + @with_pool pool function f() + identity(acquire!(pool, Float64, 2)) + end + ) + ) + # function form, w3: assignment tail @test _expansion_escape_error( :( @@ -2947,7 +3309,7 @@ end @test !_is_dotted_assign_head(1) # non-Symbol head end - @testset "_incidental_exposure: detects the three incidental patterns" begin + @testset "_incidental_exposure: detects the incidental patterns" begin # w1: direct acquire-family call tail — no tainted vars needed at all hit = _incidental_exposure(:(acquire!(pool, Float64, 4)), Set{Symbol}(), :pool) @test hit !== nothing @@ -2977,6 +3339,20 @@ end # w3': assignment tail whose RHS is an already-tainted var (alias) hit = _incidental_exposure(:(d = v), Set([:v]), :pool) @test hit == (:assign, :v) + + # w4: container literal holding a direct acquire call + @test _incidental_exposure(:((acquire!(pool, Float64, 2),)), Set{Symbol}(), :pool)[1] === :literal_acquire_call + @test _incidental_exposure(:((; a = acquire!(pool, Float64, 2))), Set{Symbol}(), :pool)[1] === :literal_acquire_call + @test _incidental_exposure(:([zeros!(pool, Float64, 2)]), Set{Symbol}(), :pool)[1] === :literal_acquire_call + # a literal that also exposes a tainted var is Stage 1's (one warning, not two) + @test _incidental_exposure(:((v, acquire!(pool, Float64, 2))), Set([:v]), :pool) === nothing + + # identity(x) is transparent, as in the direct-exposure stage and the poisoner + @test _incidental_exposure(:(identity(acquire!(pool, Float64, 2))), Set{Symbol}(), :pool)[1] === :acquire_call + @test _incidental_exposure(:(Base.identity(v .= 1.0)), Set([:v]), :pool) == (:broadcast_assign, :v) + @test _incidental_exposure(:((identity(acquire!(pool, Float64, 2)),)), Set{Symbol}(), :pool)[1] === :literal_acquire_call + @test _incidental_exposure(:(identity((acquire!(pool, Float64, 2),))), Set{Symbol}(), :pool)[1] === :literal_acquire_call + @test _incidental_exposure(:(identity(v)), Set([:v]), :pool) === nothing # Stage 1's end @testset "_incidental_exposure: safe tails return nothing" begin @@ -3012,6 +3388,12 @@ end msg = _lint_message(:assign, :v, assign_expr) @test occursin("last expression", msg) @test occursin("nothing", msg) + + lit_expr = :((acquire!(pool, Float64, 2),)) + msg = _lint_message(:literal_acquire_call, lit_expr, lit_expr) + @test occursin("holding a direct acquire call", msg) + @test occursin("last expression", msg) + @test occursin("nothing", msg) end @testset "_report_incidental_escape: both severity branches" begin @@ -3218,6 +3600,37 @@ end @test guard_call(thenlast) @test elselast === :nothing # the identifier `nothing` is a Symbol in the AST + # NamedTuple shorthand: (; s, v) → (; s, v = guard(v)) — a guard CALL is not + # a legal shorthand element, so the field is expanded to explicit `kw` form + out = _poison_block_tails( + :( + begin + v = acquire!(pool, Float64, 4) + s = 1.0 + (; s, v) + end + ), :pool, lnn + ) + tail = lastexpr(out) + @test Meta.isexpr(tail, :tuple) + p = tail.args[1] + @test Meta.isexpr(p, :parameters) + @test p.args[1] === :s # non-pool shorthand untouched + @test Meta.isexpr(p.args[2], :kw) && p.args[2].args[1] === :v && guard_call(p.args[2].args[2]) + + # field-access shorthand (; s, t.v) is not a bare name → untouched + out = _poison_block_tails( + :( + begin + v = acquire!(pool, Float64, 4) + t = (v = 1,) + s = 1.0 + (; s, t.v) + end + ), :pool, lnn + ) + @test lastexpr(out) == :((; s, t.v)) + # safe tail: structurally unchanged src = :( begin @@ -3228,6 +3641,60 @@ end @test _poison_block_tails(src, :pool, lnn) == src end + @testset "checker ⇔ poisoner lockstep (block tails)" begin + # Every block tail the poisoner rewrites must also be reported by the checker + # (and vice versa): a silent rewrite hides the escape, and a warning without a + # rewrite leaves the pool-backed value unguarded. + prelude = :( + begin + s = 1.0 + v = acquire!(pool, Float64, 3) + w = acquire!(pool, Float64, 3) + end + ) + body(tail) = Expr(:block, prelude.args..., tail) + escaping = [ + :v, :((s, v)), :((; s, v)), :((; s, v = v)), :((a = v,)), :([v, w]), + :(identity(v)), :(acquire!(pool, Float64, 2)), :(v .= 1.0), :(dest[:] = v), + :(u = acquire!(pool, Float64, 2)), :( + if s > 0 + v + else + nothing + end + ), + # acquire calls held directly in a container literal + :((acquire!(pool, Float64, 2),)), :((; a = acquire!(pool, Float64, 2))), + :([acquire!(pool, Float64, 2)]), + # identity(...) wrappers are transparent to both checker and poisoner + :(identity(acquire!(pool, Float64, 2))), :((identity(acquire!(pool, Float64, 2)),)), + :(identity((acquire!(pool, Float64, 2),))), :(identity(v .= 1.0)), + ] + safe = [:(sum(v)), :((; s)), :((; s, n = length(v))), :nothing] + for tail in escaping + b = body(tail) + @test _block_tail_warned(b) + @test _block_tail_warned(b) == _block_tail_rewritten(b) + end + for tail in safe + b = body(tail) + @test !_block_tail_warned(b) + @test _block_tail_warned(b) == _block_tail_rewritten(b) + end + + # reporter shape: element read from a pool-array loop, shorthand tail → safe + b = quote + order = acquire!(pool, Int, 4) + contact = 0 + for k in order + contact = k + end + (; contact) + end + @test !_block_tail_warned(b) + @test _block_tail_warned(b) == _block_tail_rewritten(b) + end + @testset "block-tail guard integration (runtime)" begin # NOTE: subjects are eval'd at RUNTIME so the expansion-time @warn is # captured by @test_logs (a literal @with_pool here would warn once at @@ -3313,6 +3780,65 @@ end @test u isa EscapedPoolArray @test u.var === :expression @test u.dims == (3,) + + # NamedTuple shorthand tail: must expand (regression: the guard call used to + # land in a shorthand slot → "invalid named tuple element") with names kept + r = @test_logs (:warn, r"becomes the scope's return value") match_mode = :any Core.eval( + @__MODULE__, :( + @with_pool pool begin + s = 1.5 + v = acquire!(pool, Float64, 3) + w = acquire!(pool, Float64, 2) + (; s, v = v, w) + end + ) + ) + @test r isa NamedTuple + @test keys(r) == (:s, :v, :w) + @test r.s === 1.5 + @test r.v isa EscapedPoolArray && r.v.var === :v + @test r.w isa EscapedPoolArray && r.w.var === :w + @test_throws EscapedPoolUseError r.w[1] + + # mixed literal (tainted var + direct acquire call): exactly ONE escape + # warning (Stage 1 owns it), and both elements are guarded + logs, r = Test.collect_test_logs() do + Core.eval( + @__MODULE__, :( + @with_pool pool begin + v = acquire!(pool, Float64, 2) + (v, acquire!(pool, Float64, 3)) + end + ) + ) + end + @test count(l -> occursin("becomes the scope's return value", string(l.message)), logs) == 1 + @test r[1] isa EscapedPoolArray && r[2] isa EscapedPoolArray + end + + @testset "reporter shape: loop element + NamedTuple shorthand (runtime)" begin + # Regression (FastInterpolations): an integer read from `for k in order` was + # tainted as an alias of `order`, and the `(; …)` tail then failed to compile. + r = @test_logs min_level = Base.CoreLogging.Warn Core.eval( + @__MODULE__, :( + let n = 4 + @with_pool pool begin + order = acquire!(pool, Int, n) + order .= (4, 2, 3, 1) + s = 0.0 + contact = 0 + for k in order + s += k + if k == 3 + contact = k + end + end + (; s, contact) + end + end + ) + ) + @test r === (s = 10.0, contact = 3) end end # Compile-Time Escape Detection