(feat): add 2-tier pool safety — compile-time escape detection + runtime validation

src/AdaptiveArrayPools.jl

@@ -9,7 +9,8 @@ export zeros!, ones!, trues!, falses!, similar!, reshape!, default_eltype  # Con
 export unsafe_zeros!, unsafe_ones!, unsafe_similar!  # Unsafe convenience functions
 export Bit  # Sentinel type for BitArray (use with acquire!, trues!, falses!)
 export @with_pool, @maybe_with_pool
-export STATIC_POOLING, MAYBE_POOLING, POOL_DEBUG
+export STATIC_POOLING, MAYBE_POOLING, POOL_DEBUG, POOL_SAFETY_LV, STATIC_POOL_CHECKS
+export PoolEscapeError, EscapePoint
 export USE_POOLING, MAYBE_POOLING_ENABLED  # Deprecated aliases (backward compat)
 export checkpoint!, rewind!, reset!
 export get_task_local_cuda_pool, get_task_local_cuda_pools  # CUDA (stubs, overridden by extension)
@@ -28,6 +29,7 @@ export DisabledPool, DISABLED_CPU, pooling_enabled  # Disabled pool support
     include("convenience.jl")
     include("state.jl")
     include("task_local_pool.jl")
+    include("debug.jl")
     include("macros.jl")
 else
     export CACHE_WAYS, set_cache_ways!  # N-way cache configuration (legacy only)
@@ -38,6 +40,7 @@ else
     include("convenience.jl")
     include("legacy/state.jl")
     include("task_local_pool.jl")
+    include("debug.jl")
     include("macros.jl")
 end
 

src/acquire.jl

@@ -327,12 +327,16 @@ Internal implementation of acquire!. Called directly by macro-transformed code
 """
 @inline function _acquire_impl!(pool::AbstractArrayPool, ::Type{T}, n::Int) where {T}
     tp = get_typed_pool!(pool, T)
-    return get_view!(tp, n)
+    result = get_view!(tp, n)
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 @inline function _acquire_impl!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     tp = get_typed_pool!(pool, T)
-    return get_view!(tp, dims)
+    result = get_view!(tp, dims)
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 @inline function _acquire_impl!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
@@ -349,17 +353,23 @@ Internal implementation of unsafe_acquire!. Called directly by macro-transformed
 """
 @inline function _unsafe_acquire_impl!(pool::AbstractArrayPool, ::Type{T}, n::Int) where {T}
     tp = get_typed_pool!(pool, T)
-    return get_array!(tp, (n,))
+    result = get_array!(tp, (n,))
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 @inline function _unsafe_acquire_impl!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     tp = get_typed_pool!(pool, T)
-    return get_array!(tp, dims)
+    result = get_array!(tp, dims)
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 @inline function _unsafe_acquire_impl!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
     tp = get_typed_pool!(pool, T)
-    return get_array!(tp, dims)
+    result = get_array!(tp, dims)
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 # Similar-style
@@ -403,18 +413,21 @@ See also: [`unsafe_acquire!`](@ref) for native array access.
 """
 @inline function acquire!(pool::AbstractArrayPool, ::Type{T}, n::Int) where {T}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct acquire! call>")
     return _acquire_impl!(pool, T, n)
 end
 
 # Multi-dimensional support (zero-allocation with N-D cache)
 @inline function acquire!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct acquire! call>")
     return _acquire_impl!(pool, T, dims...)
 end
 
 # Tuple support: allows acquire!(pool, T, size(A)) where size(A) returns NTuple{N,Int}
 @inline function acquire!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct acquire! call>")
     return _acquire_impl!(pool, T, dims...)
 end
 
@@ -435,6 +448,7 @@ end
 """
 @inline function acquire!(pool::AbstractArrayPool, x::AbstractArray)
     _record_type_touch!(pool, eltype(x))
+    _set_pending_callsite!(pool, "<direct acquire! call>")
     return _acquire_impl!(pool, eltype(x), size(x))
 end
 
@@ -490,17 +504,20 @@ See also: [`acquire!`](@ref) for view-based access.
 """
 @inline function unsafe_acquire!(pool::AbstractArrayPool, ::Type{T}, n::Int) where {T}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_acquire! call>")
     return _unsafe_acquire_impl!(pool, T, n)
 end
 
 @inline function unsafe_acquire!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_acquire! call>")
     return _unsafe_acquire_impl!(pool, T, dims...)
 end
 
 # Tuple support
 @inline function unsafe_acquire!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_acquire! call>")
     return _unsafe_acquire_impl!(pool, T, dims)
 end
 
@@ -521,6 +538,7 @@ end
 """
 @inline function unsafe_acquire!(pool::AbstractArrayPool, x::AbstractArray)
     _record_type_touch!(pool, eltype(x))
+    _set_pending_callsite!(pool, "<direct unsafe_acquire! call>")
     return _unsafe_acquire_impl!(pool, eltype(x), size(x))
 end
 

src/bitarray.jl

@@ -139,17 +139,23 @@ end
 # Bit type: returns BitArray{N} with shared chunks (SIMD optimized, N-D cached)
 @inline function _unsafe_acquire_impl!(pool::AbstractArrayPool, ::Type{Bit}, n::Int)
     tp = get_typed_pool!(pool, Bit)::BitTypedPool
-    return get_bitarray!(tp, n)
+    result = get_bitarray!(tp, n)
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 @inline function _unsafe_acquire_impl!(pool::AbstractArrayPool, ::Type{Bit}, dims::Vararg{Int, N}) where {N}
     tp = get_typed_pool!(pool, Bit)::BitTypedPool
-    return get_bitarray!(tp, dims)
+    result = get_bitarray!(tp, dims)
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 @inline function _unsafe_acquire_impl!(pool::AbstractArrayPool, ::Type{Bit}, dims::NTuple{N, Int}) where {N}
     tp = get_typed_pool!(pool, Bit)::BitTypedPool
-    return get_bitarray!(tp, dims)
+    result = get_bitarray!(tp, dims)
+    _maybe_record_borrow!(pool, tp)
+    return result
 end
 
 # ==============================================================================
@@ -193,19 +199,24 @@ end
 # ==============================================================================
 
 # Check if BitArray chunks overlap with the pool's BitTypedPool storage
-function _check_bitchunks_overlap(arr::BitArray, pool::AdaptiveArrayPool)
+function _check_bitchunks_overlap(arr::BitArray, pool::AdaptiveArrayPool, original_val = arr)
     arr_chunks = arr.chunks
     arr_ptr = UInt(pointer(arr_chunks))
     arr_len = length(arr_chunks) * sizeof(UInt64)
     arr_end = arr_ptr + arr_len
 
+    return_site = let rs = pool._pending_return_site
+        isempty(rs) ? nothing : rs
+    end
+
     for v in pool.bits.vectors
         v_chunks = v.chunks
         v_ptr = UInt(pointer(v_chunks))
         v_len = length(v_chunks) * sizeof(UInt64)
         v_end = v_ptr + v_len
         if !(arr_end <= v_ptr || v_end <= arr_ptr)
-            error("Safety Violation: The function returned a BitArray backed by pool memory. This is unsafe as the memory will be reclaimed. Please return a copy (copy) or a scalar.")
+            callsite = _lookup_borrow_callsite(pool, v)
+            _throw_pool_escape_error(original_val, Bit, callsite, return_site)
         end
     end
     return nothing

src/convenience.jl

@@ -44,21 +44,25 @@ See also: [`ones!`](@ref), [`similar!`](@ref), [`acquire!`](@ref)
 """
 @inline function zeros!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct zeros! call>")
     return _zeros_impl!(pool, T, dims...)
 end
 
 @inline function zeros!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct zeros! call>")
     return _zeros_impl!(pool, default_eltype(pool), dims...)
 end
 
 @inline function zeros!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct zeros! call>")
     return _zeros_impl!(pool, T, dims...)
 end
 
 @inline function zeros!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct zeros! call>")
     return _zeros_impl!(pool, default_eltype(pool), dims...)
 end
 
@@ -117,21 +121,25 @@ See also: [`zeros!`](@ref), [`similar!`](@ref), [`acquire!`](@ref)
 """
 @inline function ones!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct ones! call>")
     return _ones_impl!(pool, T, dims...)
 end
 
 @inline function ones!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct ones! call>")
     return _ones_impl!(pool, default_eltype(pool), dims...)
 end
 
 @inline function ones!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct ones! call>")
     return _ones_impl!(pool, T, dims...)
 end
 
 @inline function ones!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct ones! call>")
     return _ones_impl!(pool, default_eltype(pool), dims...)
 end
 
@@ -187,10 +195,12 @@ See also: [`falses!`](@ref), [`ones!`](@ref), [`acquire!`](@ref)
 """
 @inline function trues!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, Bit)
+    _set_pending_callsite!(pool, "<direct trues! call>")
     return _trues_impl!(pool, dims...)
 end
 @inline function trues!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N}
     _record_type_touch!(pool, Bit)
+    _set_pending_callsite!(pool, "<direct trues! call>")
     return _trues_impl!(pool, dims...)
 end
 
@@ -227,10 +237,12 @@ See also: [`trues!`](@ref), [`zeros!`](@ref), [`acquire!`](@ref)
 """
 @inline function falses!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, Bit)
+    _set_pending_callsite!(pool, "<direct falses! call>")
     return _falses_impl!(pool, dims...)
 end
 @inline function falses!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N}
     _record_type_touch!(pool, Bit)
+    _set_pending_callsite!(pool, "<direct falses! call>")
     return _falses_impl!(pool, dims...)
 end
 
@@ -274,21 +286,25 @@ See also: [`zeros!`](@ref), [`ones!`](@ref), [`acquire!`](@ref)
 """
 @inline function similar!(pool::AbstractArrayPool, x::AbstractArray)
     _record_type_touch!(pool, eltype(x))
+    _set_pending_callsite!(pool, "<direct similar! call>")
     return _similar_impl!(pool, x)
 end
 
 @inline function similar!(pool::AbstractArrayPool, x::AbstractArray, ::Type{T}) where {T}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct similar! call>")
     return _similar_impl!(pool, x, T)
 end
 
 @inline function similar!(pool::AbstractArrayPool, x::AbstractArray, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, eltype(x))
+    _set_pending_callsite!(pool, "<direct similar! call>")
     return _similar_impl!(pool, x, dims...)
 end
 
 @inline function similar!(pool::AbstractArrayPool, x::AbstractArray, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct similar! call>")
     return _similar_impl!(pool, x, T, dims...)
 end
 
@@ -398,21 +414,25 @@ See also: [`unsafe_ones!`](@ref), [`zeros!`](@ref), [`unsafe_acquire!`](@ref)
 """
 @inline function unsafe_zeros!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_zeros! call>")
     return _unsafe_zeros_impl!(pool, T, dims...)
 end
 
 @inline function unsafe_zeros!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct unsafe_zeros! call>")
     return _unsafe_zeros_impl!(pool, default_eltype(pool), dims...)
 end
 
 @inline function unsafe_zeros!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_zeros! call>")
     return _unsafe_zeros_impl!(pool, T, dims...)
 end
 
 @inline function unsafe_zeros!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct unsafe_zeros! call>")
     return _unsafe_zeros_impl!(pool, default_eltype(pool), dims...)
 end
 
@@ -465,21 +485,25 @@ See also: [`unsafe_zeros!`](@ref), [`ones!`](@ref), [`unsafe_acquire!`](@ref)
 """
 @inline function unsafe_ones!(pool::AbstractArrayPool, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_ones! call>")
     return _unsafe_ones_impl!(pool, T, dims...)
 end
 
 @inline function unsafe_ones!(pool::AbstractArrayPool, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct unsafe_ones! call>")
     return _unsafe_ones_impl!(pool, default_eltype(pool), dims...)
 end
 
 @inline function unsafe_ones!(pool::AbstractArrayPool, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_ones! call>")
     return _unsafe_ones_impl!(pool, T, dims...)
 end
 
 @inline function unsafe_ones!(pool::AbstractArrayPool, dims::NTuple{N, Int}) where {N}
     _record_type_touch!(pool, default_eltype(pool))
+    _set_pending_callsite!(pool, "<direct unsafe_ones! call>")
     return _unsafe_ones_impl!(pool, default_eltype(pool), dims...)
 end
 
@@ -535,21 +559,25 @@ See also: [`similar!`](@ref), [`unsafe_acquire!`](@ref)
 """
 @inline function unsafe_similar!(pool::AbstractArrayPool, x::AbstractArray)
     _record_type_touch!(pool, eltype(x))
+    _set_pending_callsite!(pool, "<direct unsafe_similar! call>")
     return _unsafe_similar_impl!(pool, x)
 end
 
 @inline function unsafe_similar!(pool::AbstractArrayPool, x::AbstractArray, ::Type{T}) where {T}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_similar! call>")
     return _unsafe_similar_impl!(pool, x, T)
 end
 
 @inline function unsafe_similar!(pool::AbstractArrayPool, x::AbstractArray, dims::Vararg{Int, N}) where {N}
     _record_type_touch!(pool, eltype(x))
+    _set_pending_callsite!(pool, "<direct unsafe_similar! call>")
     return _unsafe_similar_impl!(pool, x, dims...)
 end
 
 @inline function unsafe_similar!(pool::AbstractArrayPool, x::AbstractArray, ::Type{T}, dims::Vararg{Int, N}) where {T, N}
     _record_type_touch!(pool, T)
+    _set_pending_callsite!(pool, "<direct unsafe_similar! call>")
     return _unsafe_similar_impl!(pool, x, T, dims...)
 end