API Reference

@maybe_with_pool pool_name expr
@maybe_with_pool expr

Conditionally enables pooling based on MAYBE_POOLING[]. If disabled, pool_name is bound to a DisabledPool sentinel (e.g. DISABLED_CPU on CPU), and acquire! falls back to standard allocation.

Useful for libraries that want to let users control pooling behavior at runtime.

Like @with_pool, does not use try-finally — see @with_pool for exception behavior details. For exception safety, use @safe_maybe_with_pool.

Function Definition

Like @with_pool, wrap function definitions:

@maybe_with_pool pool function process_data(data)
    tmp = acquire!(pool, Float64, length(data))  # Conditionally pooled
    tmp .= data
    sum(tmp)
end

Block Usage

MAYBE_POOLING[] = false
@maybe_with_pool pool begin
    v = acquire!(pool, Float64, 100)  # Falls back to Vector{Float64}(undef, 100)
end

@safe_maybe_with_pool pool_name expr
@safe_maybe_with_pool expr
@safe_maybe_with_pool :backend pool_name expr
@safe_maybe_with_pool :backend expr

Like @maybe_with_pool but uses try-finally for exception safety. Combines the runtime pooling toggle of @maybe_with_pool with the exception guarantees of @safe_with_pool.

See also: @maybe_with_pool, @safe_with_pool

@safe_with_pool pool_name expr
@safe_with_pool expr
@safe_with_pool :backend pool_name expr
@safe_with_pool :backend expr

Like @with_pool but uses try-finally to guarantee pool cleanup even when exceptions are thrown. Use this when code inside the pool scope may throw and you need the pool to remain in a valid state afterward.

Performance note: try-finally prevents Julia's compiler from inlining the pool scope, resulting in ~35-73% overhead compared to @with_pool. Prefer @with_pool for hot paths and use @safe_with_pool only when exception safety is required.

See also: @with_pool, @safe_maybe_with_pool

@with_pool pool_name expr
@with_pool expr
@with_pool :backend pool_name expr
@with_pool :backend expr

Executes code within a pooling scope with automatic lifecycle management. Calls checkpoint! on entry and inserts rewind! at every exit point (implicit return, explicit return, break, continue).

If pool_name is omitted, a hidden variable is used (useful when you don't need to reference the pool directly).

Backend Selection

Use a symbol to specify the pool backend:

• :cpu - CPU pools (default)
• :cuda - GPU pools (requires using CUDA)

# CPU (default)
@with_pool pool begin ... end

# GPU via CUDA
@with_pool :cuda pool begin ... end

Function Definition

Wrap function definitions to inject pool lifecycle into the body:

# Long form function
@with_pool pool function compute_stats(data)
    tmp = acquire!(pool, Float64, length(data))
    tmp .= data
    mean(tmp), std(tmp)
end

# Short form function
@with_pool pool fast_sum(data) = begin
    tmp = acquire!(pool, eltype(data), length(data))
    tmp .= data
    sum(tmp)
end

Block Usage

# With explicit pool name
@with_pool pool begin
    v = acquire!(pool, Float64, 100)
    v .= 1.0
    sum(v)
end

# Without pool name (for simple blocks)
@with_pool begin
    inner_function()  # inner function can use get_task_local_pool()
end

Nesting

Nested @with_pool blocks work correctly - each maintains its own checkpoint.

@with_pool p1 begin
    v1 = acquire!(p1, Float64, 10)
    inner = @with_pool p2 begin
        v2 = acquire!(p2, Float64, 5)
        sum(v2)
    end
    # v1 is still valid here
    sum(v1) + inner
end

Exception Behavior

@with_pool does not use try-finally (for inlining performance). Implications:

1. Uncaught exceptions: If an exception propagates out of all @with_pool scopes, pool state is invalid. Call reset!(pool) or use a fresh pool.
2. Caught exceptions (nested): If an inner @with_pool throws and an outer scope catches, the outer scope's exit will clean up leaked inner scopes automatically (deferred recovery). Do not use pool operations inside the catch block.
3. PoolRuntimeEscapeError: After this error fires, the pool is poisoned. Fix the bug in your code and restart.
4. For full exception safety (try-finally guarantee), use @safe_with_pool.

_acquire_call_kind(expr, target_pool) -> Union{Symbol, Nothing}

Return the classification of an acquire call: :pool_view, :pool_array, :pool_bitarray, or nothing if not an acquire call.

_acquire_impl!(pool, Type{T}, n) -> Array{T,1}
_acquire_impl!(pool, Type{T}, dims...) -> Array{T,N}

Internal implementation of acquire!. Called directly by macro-transformed code (no type touch recording). User code calls acquire! which adds recording.

Returns raw Array{T,N} via cached wrapper reuse (setfield!-based on Julia 1.11+).

_acquire_view_impl!(pool, Type{T}, n) -> SubArray{T,1,Vector{T},...}
_acquire_view_impl!(pool, Type{T}, dims...) -> ReshapedArray{T,N,...}

Internal implementation of acquire_view!. Called directly by macro-transformed code (no type touch recording). Returns views into pool-managed vectors.

_can_use_typed_path(pool::AbstractArrayPool, tracked_mask::UInt16) -> Bool

Check if the typed (fast) checkpoint/rewind path is safe to use.

Returns true when all touched types at the current depth are a subset of the tracked types (bitmask subset check) AND no non-fixed-slot types were touched.

The subset check: (touched_mask & ~tracked_mask) == 0 means every bit set in touched_mask is also set in tracked_mask.

_check_compile_time_escape(expr, pool_name, source)

Compile-time (macro expansion time) escape detection.

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.

All detected escapes are errors — bare symbol (v), return v, and container patterns ((v, w), [v], (key=v,)).

Skipped when STATIC_POOLING = false (pooling disabled, acquire returns normal arrays).

_check_structural_mutation(expr, pool_name, source)

Compile-time (macro expansion time) structural mutation detection.

Checks if any pool-backed variable is passed to resize!, push!, pop!, or other functions that change array structure. These operations may cause pooling benefits (zero-alloc reuse) to be lost and temporary extra memory retention, because pool-backed arrays share memory with backing storage.

Skipped when STATIC_POOLING = false (pooling disabled, acquire returns normal arrays).

_check_unsafe_goto(expr)

Hard error if the body contains any @goto that targets a label NOT defined within the same body. Such jumps would bypass rewind! insertion.

Internal jumps (@goto label where @label label exists in the body) are safe and allowed — they don't exit the pool scope.

_claim_slot!(tp::TypedPool{T}) -> Int

Claim the next slot without provisioning memory (zero-length backing vector). Used by reshape! which only needs the slot index for nd_wrapper caching — the wrapper points to a different array's memory via setfield!(:ref).

_claim_slot!(tp::TypedPool{T}, n::Int) -> Int

Claim the next slot, ensuring the backing vector exists and has capacity >= n. Returns the slot index. This is the shared primitive for all acquisition paths (get_view!, get_array!).

_classify_escaped_vars(expr, target_pool, escaped, acquired)

Classify each escaped variable by its origin for better error messages:

• :pool_view — from acquire_view! (returns SubArray)
• :pool_array — from acquire!, zeros!, etc. (returns Array)
• :pool_bitarray — from trues!, falses! (returns BitArray)
• :alias — alias of another acquired variable (e.g., d = v)
• :container — wraps acquired variables in a literal (e.g., d = [v, 1])

Returns Dict{Symbol, Tuple{Symbol, Vector{Symbol}}} mapping var → (kind, source_vars).

Collect acquired variable names contained in a literal expression (symbol, tuple, vect).

_collect_all_return_values(expr) -> Vector{Tuple{Any, Union{Int,Nothing}}}

Collect all (expression, line) pairs that could be returned from a block/function body:

• Explicit return expr statements anywhere in the body (recursive, skips nested functions)
• Implicit returns: the last expression, recursing into if/else/elseif branches

Walk AST to find all explicit return expr statements with line numbers. Tracks LineNumberNodes through blocks. Skips nested function definitions.

Expand implicit return values by recursing into if/elseif/else branches. Non-branch expressions are collected as (expr, line) pairs.

_collect_local_gotos_and_labels(expr) -> (gotos::Set{Symbol}, labels::Set{Symbol})

Walk the body AST and collect all @goto target symbols and @label names. Skips :function/:-> bodies (inner functions have their own scope).

At macro expansion time, @goto/@label are :macrocall nodes, not :symbolicgoto.

_contains_acquire_call(expr, pool_name) -> Bool

Detect if expr (or any sub-expression) contains a pool acquire call. Matches both transformed (GlobalRef-based _*_impl!) and original (acquire!, zeros!, etc.) call forms.

_disabled_pool_expr(backend::Symbol) -> Expr

Generate expression for DisabledPool singleton based on backend. Used when pooling is disabled to preserve backend context.

_ensure_body_has_toplevel_lnn(body, source)

Ensure body has a LineNumberNode pointing to user source at the top level.

• Scans first few args to handle Expr(:meta, ...) from @inline etc.
• If first LNN points to user file (same as source.file), preserve it
• If first LNN points elsewhere (e.g., macros.jl), replace with source LNN
• If no LNN exists, prepend source LNN
• If source.file === :none (REPL/eval), don't clobber valid file LNNs

Returns a new Expr to avoid mutating the original AST.

_extract_acquire_types(expr, target_pool) -> Set{Any}

Extract type arguments from acquire/convenience function calls in an expression. Only extracts types from calls where the first argument matches target_pool. This prevents AST pollution when multiple pools are used in the same block.

Supported functions:

• acquire! and its alias acquire_array!
• acquire_view!
• zeros!, ones!, similar!

Handles various forms:

• acquire!(pool, Type, dims...): extracts Type directly
• acquire!(pool, x): generates eltype(x) expression
• zeros!(pool, dims...) / ones!(pool, dims...): Float64 (default)
• zeros!(pool, Type, dims...) / ones!(pool, Type, dims...): extracts Type
• similar!(pool, x): generates eltype(x) expression
• similar!(pool, x, Type, ...): extracts Type

_extract_acquired_vars(expr, target_pool) -> Set{Symbol}

Walk AST to find variable names assigned from acquire/convenience calls. Returns the set of symbols that hold pool-backed arrays.

Only top-level assignments in a block are tracked (not inside branches). Handles both simple assignment (v = acquire!(...)) and tuple destructuring ((v, w) = (acquire!(...), expr)).

_extract_declaration_sites(expr, escaped)

Walk the AST to find assignment sites for escaped variables. Returns a Vector{DeclarationSite} sorted by line number.

_extract_local_assignments(expr, locals=Set{Symbol}()) -> Set{Symbol}

Find all symbols that are assigned locally in the expression body. These cannot be used for typed checkpoint since they're defined after checkpoint!.

Detects patterns like: T = eltype(x), local T = ..., etc.

_filter_static_types(types, local_vars=Set{Symbol}()) -> (static_types, has_dynamic)

Filter types for typed checkpoint/rewind generation.

• Symbols NOT in local_vars are passed through (type parameters, global types)
• Symbols IN local_vars trigger fallback (defined after checkpoint!)
• Parametric types like Vector{T} trigger fallback
• eltype(x) expressions: usable if x does NOT reference a local variable

Type parameters (T, S from where clause) resolve to concrete types at runtime. Local variables (T = eltype(x)) are defined after checkpoint! and cannot be used.

_find_acquire_call_expr(expr, pool_name) -> Union{Expr, Nothing}

Find the first acquire call expression in expr targeting pool_name. Returns the original call Expr (e.g., :(zeros!(pool, Float64, 10))) or nothing. Used to capture the user's source expression for debug display.

_find_first_lnn_index(args) -> Union{Int, Nothing}

Find the index of the first LineNumberNode in the leading prefix of args.

Scans sequentially, skipping Expr(:meta, ...) nodes (inserted by @inline, @inbounds, etc.). Returns nothing as soon as a non-meta, non-LNN expression is encountered—this prevents matching LNNs deeper in the AST.

Example AST prefix patterns

• [Expr(:meta,:inline), LNN, ...] → returns 2
• [LNN, ...] → returns 1
• [Expr(:meta,:inline), Expr(:call,...), LNN, ...] → returns nothing (stopped at call)

_find_mutation_calls(expr, acquired) -> Vector{MutationPoint}

Walk AST to find all structural mutation calls on pool-backed variables. Tracks LineNumberNodes for precise error locations. Skips nested function definitions (lambdas within @with_pool are separate scopes).

_fix_generated_lnn!(expr, source)

Replace all macro-generated LineNumberNodes with user source location.

When quote...end blocks in macros.jl generate code, every line gets a LNN pointing to macros.jl. This causes Julia's stack traces to show macros.jl instead of the user's call site. This function walks the entire AST and replaces every LNN whose file differs from source.file with the source LNN.

User code (inserted via esc()) retains its own LNNs since those already point to the user's file and won't be replaced.

If source.file === :none (REPL/eval), don't clobber valid file LNNs. Modifies expr in-place and returns it.

_generate_block_inner(pool_name, expr, safe, source) -> Expr

Generate the inner body for block-form @with_pool. Handles both safe (try-finally) and direct-rewind paths. Used by both CPU and backend-specific code generators.

Does NOT handle the outer dispatch wrapper or MAYBE_POOLING branching — callers handle those after receiving the inner body.

_generate_function_inner(pool_name, expr, safe, source)

Shared helper for function-form code generation (both CPU and backend variants). Like _generate_block_inner but does NOT apply _transform_break_continue — break/continue cannot exit a function scope.

_generate_function_pool_code_with_backend(backend, pool_name, func_def, force_enable, disable_pooling)

Generate function code for a specific backend (e.g., :cuda). Wraps the function body with pool getter, checkpoint, and rewind.

When disable_pooling=true (STATICPOOLING=false), generates DisabledPool binding. When `forceenable=true(@with_pool), always uses the real pool. Whenforceenable=false` (@maybewithpool), generates MAYBEPOOLING[] runtime check.

_generate_lazy_checkpoint_call(pool_expr)

Generate a depth-only checkpoint call for dynamic-selective mode (use_typed=false). Much lighter than full checkpoint!: only increments depth and pushes bitmask sentinels.

_generate_lazy_rewind_call(pool_expr)

Generate selective rewind code for dynamic-selective mode (use_typed=false). Delegates to _lazy_rewind! — a single function call, symmetric with _lazy_checkpoint! for checkpoint. This avoids let-block overhead in finally clauses (which can impair Julia's type inference and cause boxing).

_generate_pool_code_with_backend(backend, pool_name, expr, force_enable)

Generate pool code for a specific backend (e.g., :cuda, :cpu). Uses _get_pool_for_backend(Val{backend}()) for zero-overhead dispatch.

Includes type-specific checkpoint/rewind optimization (same as regular @with_pool).

_generate_raw_entry_depth_guard(pool_name, entry_depth_var) -> Expr

Generate un-escaped entry depth guard for cleaning up leaked inner scopes.

Produces: while pool._current_depth > _entry_depth + 1; rewind!(pool); end Uses full rewind!(pool) (not typed/lazy) because leaked inner scope may have touched types outside this scope's static type set.

_generate_raw_rewind_call(pool_name, use_typed, static_types) -> Expr

Generate un-escaped rewind call for embedding in AST transforms. Uses GlobalRef function references and raw pool_name symbol.

_generate_typed_checkpoint_call(pool_expr, types)

Generate bitmask-aware checkpoint call. When types are known at compile time, emits a conditional:

• if touched types ⊆ tracked types → typed checkpoint (fast path)
• otherwise → _typed_lazy_checkpoint! (typed checkpoint + set bit 14 for lazy first-touch checkpointing of extra types touched by helpers)

_generate_typed_rewind_call(pool_expr, types)

Generate bitmask-aware rewind call. When types are known at compile time, emits a conditional:

• if touched types ⊆ tracked types → typed rewind (fast path)
• otherwise → _typed_lazy_rewind! (rewinds tracked | touched mask; all touched types have Case A checkpoints via bit 14 lazy mode)

_get_last_expression(expr) -> Any

Return the last non-LineNumberNode expression from a block. For non-block expressions, returns the expression itself.

Return the last non-LineNumberNode expression from a block, together with its line. Recurses into nested blocks.

_get_pool_for_backend(::Val{:cpu}) -> AdaptiveArrayPool

Get task-local pool for the specified backend.

Extensions add methods for their backends (e.g., Val{:cuda}). Using Val{Symbol} enables compile-time dispatch and full inlining, achieving zero overhead compared to Dict-based registry.

Example (in CUDA extension)

@inline AdaptiveArrayPools._get_pool_for_backend(::Val{:cuda}) = get_task_local_cuda_pool()

_inject_pending_callsite(expr, pool_name, original_expr=expr) -> Expr

Walk block-level statements, track LineNumberNodes, and insert _runtime_check(pool) && (pool._pending_callsite = "file:line\nexpr") before each statement containing a pool acquire call.

When original_expr differs from expr (i.e., after _transform_acquire_calls), the original untransformed AST is used to extract the user's source expression (e.g., zeros!(pool, Float64, 10) instead of _zeros_impl!(pool, Float64, 10)).

Only processes :block expressions. Non-block expressions are recursed into to find nested blocks.

_is_acquire_call(expr, target_pool) -> Bool

Check if an expression is a call to any pool acquire/convenience function targeting target_pool.

Check if a call target is identity or Base.identity.

_is_mutation_call(expr, acquired) -> Union{Tuple{Symbol, Symbol}, Nothing}

Check if expr is a call to a structural mutation function with a pool-backed variable as the first argument. Returns (func_name, var_name) or nothing.

_lazy_checkpoint!(pool::AdaptiveArrayPool)

Lightweight checkpoint for lazy mode (use_typed=false macro path).

Increments _current_depth and pushes bitmask sentinels — but does not save n_active for any fixed-slot typed pool. The _LAZY_MODE_BIT (bit 15) in _touched_type_masks marks this depth as lazy mode so that _record_type_touch! can trigger lazy first-touch checkpoints.

Existing others entries are eagerly checkpointed since there is no per-type tracking for non-fixed-slot pools; Case B in _rewind_typed_pool! handles any new others entries created during the scope (n_active starts at 0 = sentinel).

Performance: ~2ns vs ~540ns for full checkpoint!.

_lazy_rewind!(pool::AdaptiveArrayPool)

Complete rewind for lazy mode (use_typed=false macro path).

Reads the combined mask at the current depth, rewinds only the fixed-slot pools whose bits are set, handles any others entries, then pops the depth metadata.

Called directly from the macro-generated finally clause as a single function call (matching the structure of _lazy_checkpoint! for symmetry and performance).

Check if a literal expression (symbol, identity call, tuple/vect) transitively contains any acquired var.

_looks_like_type(expr) -> Bool

Heuristic to check if an expression looks like a type. Returns true for: uppercase Symbols (Float64, Int), curly expressions (Vector{T}), GlobalRef to types.

_lookup_borrow_callsite(pool, v) -> Union{Nothing, String}

Look up the callsite string for a pool backing vector. Returns nothing if no borrow was recorded (S=0 or non-macro path without callsite info).

_make_pool(level, old::AdaptiveArrayPool) -> AdaptiveArrayPool

Create a new pool, transferring cached arrays and scope state from old. Only reference copies — no memory allocation for the underlying buffers.

Transferred: all TypedPool/BitTypedPool slots, others, depth & touch tracking. Reset: _pending_callsite/return_site (transient macro state), _borrow_log (created fresh when S >= 1).

_make_pool(level) -> AdaptiveArrayPool

Function barrier: converts runtime check level to concrete AdaptiveArrayPool{S}. Accepts Bool (true→1, false→0) or Int (used directly as S).

_maybe_record_borrow!(pool, tp::AbstractTypedPool)

Flush the pending callsite into the borrow log (S=1). Delegates to _record_borrow_from_pending! (defined in debug.jl). Compiles to no-op when S=0.

_poison_released_vectors!(tp::AbstractTypedPool, old_n_active)

Fill released backing vectors (indices n_active+1:old_n_active) with sentinel values. Called from _invalidate_released_slots! when S >= 1, before resize! zeroes the lengths.

_pool_type_for_backend(::Val{B}) -> Type

Returns the concrete pool type for a given backend, used at macro expansion time to generate direct type assertions. Extensions override this for their backends.

CPU returns AdaptiveArrayPool, CUDA extension returns CuAdaptiveArrayPool.

_record_borrow_from_pending!(pool, tp)

Record the pending callsite for the most recently claimed slot in tp. Called from _acquire_impl! / _acquire_view_impl! when S >= 1.

_record_type_touch!(pool::AbstractArrayPool, ::Type{T})

Record that type T was touched (acquired) at the current checkpoint depth. Called by acquire! and convenience wrappers; macro-transformed calls use _acquire_impl! directly (bypassing this for zero overhead).

For fixed-slot types, sets the corresponding bit in _touched_type_masks. For non-fixed-slot types, sets _touched_has_others flag.

_remove_flat_reassigned!(expr, acquired, target_pool)

Walk top-level statements in order and remove variables from acquired if they are reassigned to a non-acquire call. Handles both simple assignment (v = expr) and tuple destructuring ((a, v) = expr). Only handles flat (non-branching) reassignment — conditional is conservatively kept.

Render an expression with escaped variable names highlighted in red. Handles return, tuple, NamedTuple, array literal; falls back to print for others.

_reshape_impl!(pool::AdaptiveArrayPool, A::Array{T,M}, dims::NTuple{N,Int}) -> Array{T,N}

Zero-allocation reshape using setfield!-based wrapper reuse (Julia 1.11+).

• Same dimensionality (M == N): setfield!(A, :size, dims) — no pool interaction
• Different dimensionality (M ≠ N): Claims a pool slot via _claim_slot!, reuses cached Array{T,N} wrapper with setfield!(:ref, :size) pointing to A's memory. Automatically reclaimed on rewind! via n_active restoration.

_runtime_check(pool) -> Bool

Return whether runtime safety checks are enabled for pool (i.e., S >= 1). Compile-time constant for AdaptiveArrayPool{S} — dead-code eliminated when S = 0.

_selective_rewind_fixed_slots!(pool::AdaptiveArrayPool, mask::UInt16)

Rewind only the fixed-slot typed pools whose bits are set in mask.

Each of the 8 fixed-slot pools maps to bits 0–7 (same encoding as _fixed_slot_bit). Bits 8–15 (mode flags) are not checked here — callers must strip them before passing the mask (e.g. mask & _TYPE_BITS_MASK).

Unset bits are skipped entirely: for pools that were acquired without a matching checkpoint, _rewind_typed_pool! Case B safely restores from the parent checkpoint.

_set_pending_callsite!(pool, msg::String)

Record a pending callsite string for borrow tracking (S=1). Only sets the callsite if no prior callsite is pending (macro-injected ones take priority). Compiles to no-op when S=0.

_store_arr_wrapper!(tp::AbstractTypedPool, N::Int, slot::Int, wrapper)

Store a cached N-D wrapper for the given slot. Creates the per-N Vector if needed.

_tracked_mask_for_types(types::Type...) -> UInt16

Compute compile-time bitmask for the types tracked by a typed checkpoint/rewind. Uses @generated for zero-overhead constant folding.

Returns UInt16(0) when called with no arguments. Non-fixed-slot types contribute UInt16(0) (their bit is 0).

_transform_break_continue(expr, rewind_call, entry_depth_guard) -> Expr

Walk AST and insert entry depth guard + rewind before break/continue statements that would exit the pool scope. Only used for block-form @with_pool (not function form).

Skips :for, :while bodies (break/continue there are for those loops). Skips :function, :-> bodies (inner function scope boundary).

_transform_return_stmts(expr, pool_name; rewind_call=nothing, entry_depth_guard=nothing) -> Expr

Walk AST and wrap explicit return value statements with escape validation. Generates: local _ret = value; if _runtime_check(pool) validate(_ret, pool); end; return _ret

When rewind_call and entry_depth_guard are provided (direct-rewind path, safe=false), they are inserted after validation but before return: local _ret = value; validate; entry_depth_guard; rewind_call; return _ret

When nothing (safe path / try-finally), behavior is unchanged — rewind happens in the finally clause instead.

Does NOT recurse into nested :function or :-> expressions (inner functions have their own return semantics).

_typed_lazy_checkpoint!(pool::AdaptiveArrayPool, types::Type...)

Typed checkpoint that enables lazy first-touch checkpointing for extra types touched by helpers (use_typed=true, _can_use_typed_path=false path).

Calls checkpoint!(pool, types...) (checkpoints only the statically-known types), then sets _TYPED_LAZY_BIT (bit 14) in _touched_type_masks[depth] to signal typed lazy mode.

_record_type_touch! checks (mask & _MODE_BITS_MASK) != 0 (bit 14 OR bit 15) to trigger a lazy first-touch checkpoint for each extra type on first acquire, ensuring Case A (not Case B) applies at rewind and parent n_active is preserved correctly.

_typed_lazy_rewind!(pool::AdaptiveArrayPool, tracked_mask::UInt16)

Selective rewind for typed mode (use_typed=true) fallback path.

Called when _can_use_typed_path returns false (helpers touched types beyond the statically-tracked set). Rewinds only pools whose bits are set in tracked_mask | touched_mask. All touched types have Case A checkpoints, guaranteed by the _TYPED_LAZY_BIT mode set in _typed_lazy_checkpoint!.

_uses_local_var(expr, local_vars) -> Bool

Check if an expression uses any local variable (recursively). Handles field access (x.y.z) and indexing (x[i]) by checking the base variable.

This is used to detect cases like acquire!(pool, cp1d.t_i_average) where cp1d is defined locally - the eltype expression can't be evaluated at checkpoint time since cp1d doesn't exist yet.

_wrap_with_dispatch(pool_name_esc, pool_getter, inner_body; backend=:cpu)

Direct type assertion: generates let pool = getter::PoolType{RUNTIME_CHECK}.

Since RUNTIME_CHECK is a compile-time const Int, the pool type parameter S is resolved at compile time. _runtime_check(pool) returns a compile-time Bool, enabling dead-code elimination of all safety branches when RUNTIME_CHECK = 0.

The pool type is resolved at macro expansion time via _pool_type_for_backend, which extensions override (e.g., CUDA adds CuAdaptiveArrayPool).

acquire!(pool, x::AbstractArray) -> Array

Acquire an array with the same element type and size as x (similar to similar(x)).

Example

A = rand(10, 10)
@with_pool pool begin
    B = acquire!(pool, A)  # Same type and size as A
    B .= A .* 2
end

acquire!(pool, Type{T}, n) -> Array{T,1}
acquire!(pool, Type{T}, dims...) -> Array{T,N}
acquire!(pool, Type{T}, dims::NTuple{N,Int}) -> Array{T,N}

Acquire a pooled array of type T with size n or dimensions dims.

Returns:

• CPU: Array{T,N} (zero-alloc via cached wrapper reuse on Julia 1.11+)
• Bit (T === Bit): BitVector / BitArray{N} (chunks-sharing, SIMD optimized)
• CUDA: CuArray{T,N} (unified N-way cache)

The returned array is only valid within the @with_pool scope.

Example

@with_pool pool begin
    v = acquire!(pool, Float64, 100)      # Vector{Float64}
    m = acquire!(pool, Float64, 10, 10)   # Matrix{Float64}
    v .= 1.0
    m .= 2.0
    sum(v) + sum(m)
end

See also: acquire_view! for view-based access.

acquire_array!(pool, Type{T}, dims...)

Alias for acquire!.

Explicit name emphasizing the return type is a raw Array. Use when you prefer symmetric naming with acquire_view!.

acquire_view!(pool, Type{T}, n) -> SubArray{T,1,Vector{T},...}
acquire_view!(pool, Type{T}, dims...) -> ReshapedArray{T,N,...}
acquire_view!(pool, Type{T}, dims::NTuple{N,Int}) -> view type

Acquire a view into pool-managed memory.

Returns a view (SubArray/ReshapedArray) instead of a raw Array. Useful when you need stack-allocated wrappers via SROA.

See also: acquire! for the default array-returning API.

checkpoint!(pool::AdaptiveArrayPool, types::Type...)

Save state for multiple specific types. Uses @generated for zero-overhead compile-time unrolling. Increments currentdepth once for all types.

checkpoint!(pool::AdaptiveArrayPool)

Save the current pool state (n_active counters) to internal stacks.

This is called automatically by @with_pool and related macros. After warmup, this function has zero allocation.

See also: rewind!, @with_pool

checkpoint!(pool::AdaptiveArrayPool, ::Type{T})

Save state for a specific type only. Used by optimized macros that know which types will be used at compile time.

Also updates currentdepth and bitmask state for type touch tracking.

~77% faster than full checkpoint! when only one type is used.

default_eltype(pool) -> Type

Default element type for convenience functions when type is not specified. CPU pools default to Float64, CUDA pools to Float32.

Backends can override this to provide appropriate defaults.

default_eltype(::DisabledPool{:cpu}) -> Float64

Default element type for disabled CPU pools (matches Julia's zeros() default).

falses!(pool, dims...) -> BitArray
falses!(pool, dims::Tuple) -> BitArray

Acquire a bit-packed boolean array filled with false from the pool.

Equivalent to Julia's falses(dims...) but using pooled memory. Uses ~8x less memory than zeros!(pool, Bool, dims...).

Example

@with_pool pool begin
    bv = falses!(pool, 100)       # BitVector, all false
    bm = falses!(pool, 10, 10)    # BitMatrix, all false
end

See also: trues!, zeros!, acquire!

foreach_fixed_slot(f, pool::AdaptiveArrayPool)

Apply f to each fixed slot TypedPool. Zero allocation via compile-time unrolling.

get_array!(tp::AbstractTypedPool{T,Vector{T}}, dims::NTuple{N,Int}) -> Array{T,N}

Get an N-dimensional Array from the pool with setfield!-based wrapper reuse.

Uses _claim_slot! directly for slot management (independent of view path). Cache hit: setfield!(arr, :ref/size) — 0 allocation. Cache miss: creates wrapper via setfield! pattern, then cached forever.

get_bitarray!(tp::BitTypedPool, dims::NTuple{N,Int}) -> BitArray{N}

Get a BitArray{N} that shares chunks with the pooled BitVector.

Uses setfield!-based wrapper reuse — unlimited dim patterns, 0-alloc after warmup.

Implementation Notes

• BitVector (N=1): size() uses len field, dims is ignored
• BitArray{N>1}: size() uses dims field
• All BitArrays share chunks with the pool's backing BitVector

Safety

The returned BitArray is only valid within the @with_pool scope. Do NOT use after the scope ends (use-after-free risk).

get_task_local_cuda_pool() -> CuAdaptiveArrayPool

Retrieves (or creates) the CUDA pool for the current Task and current GPU device.

Requires CUDA.jl to be loaded. Throws an error if CUDA extension is not available.

See also: get_task_local_pool for CPU pools.

get_task_local_cuda_pools() -> Dict{Int, CuAdaptiveArrayPool}

Returns the dictionary of all CUDA pools for the current task (one per device).

Requires CUDA.jl to be loaded. Throws an error if CUDA extension is not available.

get_task_local_metal_pool() -> MetalAdaptiveArrayPool

Retrieves (or creates) the Metal pool for the current Task and current Metal device.

Requires Metal.jl to be loaded. Throws an error if Metal extension is not available.

See also: get_task_local_pool for CPU pools.

get_task_local_metal_pools() -> Dict{UInt64, MetalAdaptiveArrayPool}

Returns the dictionary of all Metal pools for the current task (one per device).

Requires Metal.jl to be loaded. Throws an error if Metal extension is not available.

get_task_local_pool() -> AdaptiveArrayPool

Retrieves (or creates) the AdaptiveArrayPool for the current Task.

Each Task gets its own pool instance via task_local_storage(), ensuring thread safety without locks.

The pool type is AdaptiveArrayPool{S} where S is determined by the compile-time constant RUNTIME_CHECK::Int. Macro-generated code type-asserts directly to AdaptiveArrayPool{RUNTIME_CHECK}.

get_view!(tp::TypedPool{T}, n::Int) -> SubArray{T,1}
get_view!(tp::TypedPool{T}, dims::NTuple{N,Int}) -> ReshapedArray{T,N}

Get a pooled view from the typed pool.

• 1D: Returns a fresh SubArray (stack-allocated via SROA in compiled code).
• N-D: Returns a ReshapedArray wrapping a 1D view (zero creation cost).

Always creates fresh views — caching is unnecessary since both SubArray and ReshapedArray are small structs that SROA can stack-allocate.

Dispatches on TypedPool{T} (not AbstractTypedPool) because _claim_slot! is only defined for TypedPool{T}. Other subtypes override get_view! directly (e.g., CuTypedPool) or use a separate path (e.g., BitTypedPool → get_bitarray!).

ones!(pool, dims...) -> Array
ones!(pool, T, dims...) -> Array
ones!(pool, dims::Tuple) -> Array
ones!(pool, T, dims::Tuple) -> Array

Acquire a one-initialized array from the pool.

Equivalent to acquire!(pool, T, dims...) followed by fill!(arr, one(T)). Default element type depends on pool backend (CPU: Float64, CUDA: Float32). See default_eltype.

Example

@with_pool pool begin
    v = ones!(pool, 100)              # Vector{Float64}, all ones
    m = ones!(pool, Float32, 10, 10)  # Matrix{Float32}, all ones
end

See also: zeros!, similar!, acquire!

pool_stats(tp::AbstractTypedPool; io::IO=stdout, indent::Int=0, name::String="")

Print statistics for a TypedPool or BitTypedPool.

pool_stats(pool::AdaptiveArrayPool; io::IO=stdout)

Print detailed statistics about pool usage with colored output.

Example

pool = AdaptiveArrayPool()
@with_pool pool begin
    v = acquire!(pool, Float64, 100)
    pool_stats(pool)
end

pool_stats(:cpu; io::IO=stdout)

Print statistics for the CPU task-local pool only.

pool_stats(:cuda; io::IO=stdout)

Print statistics for CUDA task-local pools. Requires CUDA.jl to be loaded.

pool_stats(; io::IO=stdout)

Print statistics for all task-local pools (CPU and CUDA if loaded).

Example

@with_pool begin
    v = acquire!(pool, Float64, 100)
    pool_stats()  # Shows all pool stats
end

pooling_enabled(pool) -> Bool

Returns true if pool is an active pool, false if pooling is disabled.

Examples

@maybe_with_pool pool begin
    if pooling_enabled(pool)
        # Using pooled memory
    else
        # Using standard allocation
    end
end

See also: DisabledPool

reset!(tp::AbstractTypedPool)

Reset state without clearing allocated storage. Sets n_active = 0 and restores checkpoint stacks to sentinel state.

reset!(pool::AdaptiveArrayPool)

Reset pool state without clearing allocated storage.

This function:

• Resets all n_active counters to 0
• Restores all checkpoint stacks to sentinel state
• Resets _current_depth and type touch tracking state

Unlike empty!, this preserves all allocated vectors and N-D wrapper caches for reuse, avoiding reallocation costs.

Use Case

When functions that acquire from the pool are called without proper checkpoint!/rewind! management, n_active can grow indefinitely. Use reset! to cleanly restore the pool to its initial state while keeping allocated memory available.

Example

pool = AdaptiveArrayPool()

# Some function that acquires without checkpoint management
function compute!(pool)
    v = acquire!(pool, Float64, 100)
    # ... use v ...
    # No rewind! called
end

for _ in 1:1000
    compute!(pool)  # n_active grows each iteration
end

reset!(pool)  # Restore state, keep allocated memory
# Now pool.n_active == 0, but vectors are still available for reuse

See also: empty!, rewind!

reset!(pool::AdaptiveArrayPool, types::Type...)

Reset state for multiple specific types. Uses @generated for zero-overhead compile-time unrolling.

See also: reset!(::AdaptiveArrayPool), rewind!

reset!(pool::AdaptiveArrayPool, ::Type{T})

Reset state for a specific type only. Clears n_active and checkpoint stacks to sentinel state while preserving allocated vectors.

See also: reset!(::AdaptiveArrayPool), rewind!

reshape!(pool, A, dims...) -> reshaped array
reshape!(pool, A, dims::Tuple) -> reshaped array

Reshape array A to dimensions dims using the pool's wrapper cache.

The returned array shares memory with A — mutations are visible in both. The pool provides cached wrapper objects to reduce allocation on repeated calls.

On Julia 1.11+:

• If ndims(A) == length(dims) (same dimensionality), reshape! mutates A in-place by changing its size. This differs from Base.reshape, which always returns a new wrapper.
• For cross-dimensional reshapes (ndims(A) != length(dims)), the returned Array wrapper is taken from the pool's internal cache and may be reused after rewind! or pool scope exit.

As with all pool-backed objects, the reshaped result must not escape the surrounding @with_pool scope.

On Julia 1.10 and CUDA, falls back to Base.reshape.

Throws DimensionMismatch if prod(dims) != length(A).

Example

A = collect(1.0:12.0)
@with_pool pool begin
    B = reshape!(pool, A, 3, 4)   # 12-element vector → 3×4 matrix
    B[1,1] = 999.0                # A[1] is now 999.0
end

See also: acquire!, similar!

rewind!(pool::AdaptiveArrayPool)

Restore the pool state (nactive counters) from internal stacks. Uses _checkpointdepths to accurately determine which entries to pop vs restore.

Only the counters are restored; allocated memory remains for reuse. Handles touched types by checking checkpointdepths for accurate restoration.

Safety: If called at global scope (depth=1, no pending checkpoints), automatically delegates to reset! to safely clear all n_active counters.

See also: checkpoint!, reset!, @with_pool

rewind!(pool::AdaptiveArrayPool, types::Type...)

Restore state for multiple specific types in reverse order. Decrements currentdepth once after all types are rewound.

rewind!(pool::AdaptiveArrayPool, ::Type{T})

Restore state for a specific type only. Also updates currentdepth and bitmask state.

safe_prod(dims::NTuple{N, Int}) -> Int

Compute the product of dimensions with overflow checking.

Throws OverflowError if the product exceeds typemax(Int), preventing memory corruption from integer overflow in array creation operations.

Rationale

Without overflow checking, large dimensions like (10^10, 10^10) would wrap around to a small value, causing array creation to index beyond allocated memory.

Performance

Adds ~0.3-1.2 ns overhead (<1%) compared to unchecked prod(), which is negligible relative to the 100-200 ns cost of the full allocation path.

similar!(pool, array) -> Array
similar!(pool, array, T) -> Array
similar!(pool, array, dims...) -> Array
similar!(pool, array, T, dims...) -> Array

Acquire an uninitialized array from the pool, using a template array for defaults.

• similar!(pool, A): same element type and size as A
• similar!(pool, A, T): element type T, same size as A
• similar!(pool, A, dims...): same element type as A, specified dimensions
• similar!(pool, A, T, dims...): element type T, specified dimensions

Example

A = rand(10, 10)
@with_pool pool begin
    B = similar!(pool, A)              # Same type and size
    C = similar!(pool, A, Float32)     # Float32, same size
    D = similar!(pool, A, 5, 5)        # Same type, different size
    E = similar!(pool, A, Int, 20)     # Int, 1D
end

See also: zeros!, ones!, acquire!

trues!(pool, dims...) -> BitArray
trues!(pool, dims::Tuple) -> BitArray

Acquire a bit-packed boolean array filled with true from the pool.

Equivalent to Julia's trues(dims...) but using pooled memory. Uses ~8x less memory than ones!(pool, Bool, dims...).

Example

@with_pool pool begin
    bv = trues!(pool, 100)        # BitVector, all true
    bm = trues!(pool, 10, 10)     # BitMatrix, all true
end

See also: falses!, ones!, acquire!

zeros!(pool, dims...) -> Array
zeros!(pool, T, dims...) -> Array
zeros!(pool, dims::Tuple) -> Array
zeros!(pool, T, dims::Tuple) -> Array

Acquire a zero-initialized array from the pool.

Equivalent to acquire!(pool, T, dims...) followed by fill!(arr, zero(T)). Default element type depends on pool backend (CPU: Float64, CUDA: Float32). See default_eltype.

Example

@with_pool pool begin
    v = zeros!(pool, 100)              # Vector{Float64}, all zeros
    m = zeros!(pool, Float32, 10, 10)  # Matrix{Float32}, all zeros
end

See also: ones!, similar!, acquire!

empty!(tp::BitTypedPool)

Clear all internal storage for BitTypedPool, releasing all memory. Restores sentinel values for 1-based sentinel pattern.

empty!(tp::TypedPool)

Clear all internal storage for TypedPool, releasing all memory. Restores sentinel values for 1-based sentinel pattern.

empty!(pool::AdaptiveArrayPool)

Completely clear the pool, releasing all stored vectors and resetting all state.

This is useful when you want to free memory or start fresh without creating a new pool instance.

Example

pool = AdaptiveArrayPool()
v = acquire!(pool, Float64, 1000)
# ... use v ...
empty!(pool)  # Release all memory

Warning

Any SubArrays previously acquired from this pool become invalid after empty!.

AbstractArrayPool

Abstract base for multi-type array pools.

AbstractTypedPool{T, V<:AbstractVector{T}}

Abstract base for type-specific memory pools.

AdaptiveArrayPool{S}

Multi-type memory pool with fixed slots for common types and IdDict fallback for others. Zero allocation after warmup. NOT thread-safe - use one pool per Task.

The type parameter S::Int encodes the runtime check mode (0 = off, 1 = on). Inside @inline call chains, S is a compile-time constant — safety checks are eliminated by dead-code elimination when S = 0, achieving true zero overhead.

See also: [_runtime_check], [_make_pool], [RUNTIME_CHECK]

Create pool with the default RUNTIME_CHECK level.

BackendNotLoadedError <: Exception

Error thrown when a backend-specific operation is attempted but the backend package is not loaded.

Example

@maybe_with_pool :cuda pool begin
    zeros!(pool, 10)  # Throws if CUDA.jl not loaded
end

Bit

Sentinel type for bit-packed boolean storage via BitVector.

Use Bit instead of Bool in pool operations to get memory-efficient bit-packed arrays (1 bit per element vs 1 byte for Vector{Bool}).

Usage

@with_pool pool begin
    # BitVector (1 bit per element, ~8x memory savings)
    bv = acquire!(pool, Bit, 1000)

    # vs Vector{Bool} (1 byte per element)
    vb = acquire!(pool, Bool, 1000)

    # Convenience functions work too
    mask = falses!(pool, 100)       # BitVector filled with false
    flags = trues!(pool, 100)       # BitVector filled with true
end

Return Types (Unified for Performance)

Unlike other types, Bit always returns native BitVector/BitArray:

• 1D: BitVector (both acquire! and acquire_view!)
• N-D: BitArray{N} (reshaped, preserves SIMD optimization)

This design ensures users always get SIMD-optimized performance.

Performance

BitVector operations like count(), sum(), and bitwise operations are ~(10x ~ 100x) faster than equivalent operations on SubArray{Bool} because they use SIMD-optimized algorithms on packed 64-bit chunks.

@with_pool pool begin
    bv = acquire!(pool, Bit, 10000)
    fill!(bv, true)
    count(bv)  # Uses fast SIMD path automatically
end

Memory Safety

The returned BitVector shares its internal chunks array with the pool. It is only valid within the @with_pool scope - using it after the scope ends leads to undefined behavior (use-after-free risk).

See also: trues!, falses!, BitTypedPool

BitTypedPool <: AbstractTypedPool{Bool, BitVector}

Specialized pool for BitVector arrays with memory reuse.

Unlike TypedPool{Bool} which stores Vector{Bool} (1 byte per element), this pool stores BitVector (1 bit per element, ~8x memory efficiency).

Unified API (Always Returns BitVector)

Both acquire! and acquire_view! return BitVector for the Bit type. This design ensures users always get SIMD-optimized performance.

• acquire!(pool, Bit, n) → BitVector (SIMD optimized)
• acquire_view!(pool, Bit, n) → BitVector (same behavior)
• trues!(pool, n) → BitVector filled with true
• falses!(pool, n) → BitVector filled with false

Fields

• vectors: Backing BitVector storage
• arr_wrappers: Vector{Union{Nothing, Vector{Any}}} — setfield!-based cache (Julia 1.11+)
• n_active: Count of currently active arrays
• _checkpoint_*: State management stacks (1-based sentinel pattern)

Performance

Operations like count(), sum(), and bitwise operations are ~(10x ~ 100x) faster than equivalent operations on SubArray{Bool} because BitVector uses SIMD-optimized algorithms on packed 64-bit chunks.

See also: trues!, falses!, Bit

Per-variable declaration site: where an escaping variable was assigned.

DisabledPool{Backend}

Sentinel type for disabled pooling that preserves backend context. When STATIC_POOLING=false (compile-time) or MAYBE_POOLING[]=false (runtime), macros return DisabledPool{backend}() instead of nothing.

Backend symbols:

• :cpu - Standard Julia arrays
• :cuda - CUDA.jl CuArrays (defined in extension)

This enables @with_pool :cuda to return correct array types even when pooling is off.

Example

# When STATIC_POOLING=false:
@with_pool :cuda pool begin
    v = zeros!(pool, 10)  # Returns CuArray{Float32}, not Array{Float64}!
end

See also: pooling_enabled, DISABLED_CPU

Per-return-point escape detail: which expression, at which line, leaks which vars.

Per-mutation-site detail: which expression, at which line, mutates which var.

PoolEscapeError <: Exception

Thrown at macro expansion time when pool-backed variables are detected in return position of @with_pool / @maybe_with_pool blocks.

This is a compile-time check with zero runtime cost.

PoolMutationError <: Exception

Thrown at macro expansion time when structural mutation functions (resize!, push!, pop!, etc.) are called on pool-backed variables within @with_pool / @maybe_with_pool blocks.

Pool-backed arrays share memory with the pool's backing storage. Structural mutations may cause pooling benefits (zero-alloc reuse) to be lost and temporary extra memory retention until the next acquire! at the same slot.

This is a compile-time check with zero runtime cost.

PoolRuntimeEscapeError <: Exception

Thrown at runtime when _validate_pool_return detects a pool-backed array escaping from an @with_pool scope (requires RUNTIME_CHECK >= 1).

This is the runtime counterpart of PoolEscapeError (compile-time).

TypedPool{T} <: AbstractTypedPool{T, Vector{T}}

Internal structure managing pooled vectors for a specific element type T.

Fields

Storage

• vectors: Backing Vector{T} storage (actual memory allocation)

N-D Wrapper Cache (Julia 1.11+, setfield!-based reuse)

• arr_wrappers: Vector{Union{Nothing, Vector{Any}}} — indexed by N (dimensionality), each entry is a per-slot cached Array{T,N} wrapper. Uses setfield!(wrapper, :size, dims) and setfield!(wrapper, :ref, parent) for zero-allocation reuse of unlimited dim patterns.

State Management (1-based sentinel pattern)

• n_active: Count of currently active (checked-out) arrays
• _checkpoint_n_active: Saved n_active values at each checkpoint (sentinel: [0])
• _checkpoint_depths: Depth of each checkpoint entry (sentinel: [0])

Design Notes

• acquire! returns Array{T,N} via setfield! wrapper reuse — unlimited dim patterns, 0-alloc after warmup.
• acquire_view! returns 1D SubArray (created fresh, stack-allocated via SROA) or N-D ReshapedArray.
• Slot management is unified via _claim_slot! — the shared primitive for all acquisition paths.

DISABLED_CPU

Singleton instance for disabled CPU pooling. Used by macros when STATIC_POOLING=false without backend specification.

FIXED_SLOT_FIELDS

Field names for fixed slot TypedPools. Single source of truth for foreach_fixed_slot.

When modifying, also update: struct definition, get_typed_pool! dispatches, constructor. Tests verify synchronization automatically.

MAYBE_POOLING::Ref{Bool}

Runtime toggle for @maybe_with_pool macro only (Tier 2 control). When false, @maybe_with_pool will use DisabledPool, causing acquire! to allocate normally.

Note: This only affects @maybe_with_pool. @with_pool ignores this flag (always uses pooling).

For complete removal of pooling overhead at compile time, use STATIC_POOLING instead.

Default: true

MAYBE_POOLING[] = false   # disable @maybe_with_pool at runtime
MAYBE_POOLING[] = true    # re-enable

RUNTIME_CHECK::Int

Compile-time constant controlling the runtime safety check level.

• 0 — off (zero overhead, default)
• 1 — full checks (invalidation, poisoning, escape detection, borrow tracking)

Set via LocalPreferences.toml:

[AdaptiveArrayPools]
runtime_check = true   # or 1

STATIC_POOLING::Bool

Compile-time constant (master switch) to completely disable pooling. When false, all macros (@with_pool, @maybe_with_pool) generate code that uses DisabledPool, causing acquire! to fall back to normal allocation with zero overhead.

This is the Tier 1 control: when disabled, MAYBE_POOLING has no effect.

Configuration via Preferences.jl

Set in your project's LocalPreferences.toml:

[AdaptiveArrayPools]
use_pooling = false

Or programmatically (requires restart):

using Preferences
Preferences.set_preferences!(AdaptiveArrayPools, "use_pooling" => false)

Default: true

_ALL_ACQUIRE_NAMES

Set of all function names that return pool-backed arrays. Used by _extract_acquired_vars to identify assignments like v = acquire!(pool, ...).

Function names that return raw Arrays backed by pool memory.

Function names that return BitArrays from pool memory.

Set of all transformed _*_impl! function names (GlobalRef targets).

Set of function names that structurally mutate arrays (resize, grow, shrink).

Function names that return views (SubArray) from pool memory.