Remove try-finally from @with_pool for inlining performance

docs/src/architecture/macro-internals.md

@@ -6,10 +6,23 @@ This page explains the internal mechanics of the `@with_pool` macro for advanced
 
 The `@with_pool` macro provides automatic lifecycle management with three key optimizations:
 
-1. **Try-Finally Safety** — Guarantees cleanup even on exceptions
+1. **Direct Rewind (no `try-finally`)** — Enables compiler inlining for ~35-73% less overhead
 2. **Typed Checkpoint/Rewind** — Only saves/restores used types (~77% faster)
 3. **Untracked Acquire Detection** — Safely handles `acquire!` calls outside macro visibility
 
+## Why No `try-finally`?
+
+Julia's compiler cannot inline functions containing `try-finally`. For `@inline @with_pool`
+functions called in hot loops, this means every call pays an exception handler frame cost
+(~20-40ns on modern hardware, worse on Julia 1.10 LTS).
+
+`@with_pool` avoids this by inserting `rewind!` directly at every exit point instead:
+
+| Macro | Strategy | Inlinable | Use case |
+|-------|----------|:---------:|----------|
+| `@with_pool` | Direct rewind at each exit point | Yes | Default — hot paths |
+| `@safe_with_pool` | `try-finally` wrapper | No | Exception safety required |
+
 ## Basic Lifecycle Flow
 
 ```
@@ -28,24 +41,101 @@ The `@with_pool` macro provides automatic lifecycle management with three key op
 ┌─────────────────────────────────────────────────────────────┐
 │  function foo(x)                                            │
 │      pool = get_task_local_pool()                           │
+│      _entry_depth = pool._current_depth                    │
 │      checkpoint!(pool, Float64)     # ← Type-specific       │
-│      try                                                    │
-│          A = _acquire_impl!(pool, Float64, 100)             │
-│          B = _similar_impl!(pool, A)                        │
-│          return sum(A) + sum(B)                             │
-│      finally                                                │
-│          rewind!(pool, Float64)     # ← Type-specific       │
+│                                                             │
+│      A = _acquire_impl!(pool, Float64, 100)                │
+│      B = _similar_impl!(pool, A)                           │
+│      _result = sum(A) + sum(B)                             │
+│                                                             │
+│      # Entry depth guard (cleans up leaked inner scopes)    │
+│      while pool._current_depth > _entry_depth + 1          │
+│          rewind!(pool)                                      │
 │      end                                                    │
+│      rewind!(pool, Float64)     # ← Type-specific           │
+│      return _result                                         │
 │  end                                                        │
 └─────────────────────────────────────────────────────────────┘
 ```
 
 ### Key Points
 
-- **`try-finally`** ensures `rewind!` executes even if an exception occurs
+- `rewind!` is inserted at **every exit point**: implicit return, explicit `return`, `break`, `continue`
 - `acquire!` → `_acquire_impl!` transformation bypasses untracked marking overhead
 - Type-specific `checkpoint!(pool, Float64)` is ~77% faster than full checkpoint
 
+### Exit Point Coverage
+
+| Exit type | Handling |
+|-----------|----------|
+| Implicit return (end of body) | `rewind!` appended before result |
+| Explicit `return` | `rewind!` inserted before each `return` statement |
+| `break` / `continue` | `rewind!` inserted before each (block form only) |
+| `@goto` (internal) | Allowed — stays within pool scope |
+| `@goto` (external) | Hard error at macro expansion time |
+| Uncaught exception | **Not handled** — use `@safe_with_pool` or `reset!(pool)` |
+
+## Exception Behavior
+
+### `@with_pool` (direct rewind)
+
+Without `try-finally`, uncaught exceptions skip `rewind!`. This is an intentional trade-off:
+
+```julia
+# Uncaught exception → pool state invalid
+try
+    @with_pool pool begin
+        acquire!(pool, Float64, 10)
+        error("boom")       # rewind! never called
+    end
+catch
+end
+# pool._current_depth is wrong here → call reset!(pool)
+```
+
+### Entry Depth Guard (nested catch recovery)
+
+When an inner `@with_pool` throws and the outer scope catches, the outer's exit
+automatically cleans up leaked inner scopes:
+
+```julia
+@with_pool pool function outer()
+    v = acquire!(pool, Float64, 10)
+    result = try
+        @with_pool pool begin
+            acquire!(pool, UInt8, 5)
+            error("inner boom")   # inner rewind! skipped
+        end
+    catch
+        42                        # pool depth is wrong HERE
+    end
+    sum(v) + result
+    # Entry depth guard runs here → cleans up leaked inner scope
+    # Own rewind! runs → outer scope cleaned up
+end
+```
+
+!!! warning "Catch block limitation"
+    Between the inner throw and the outer scope's exit, pool depth is incorrect.
+    Do not use pool operations inside the `catch` block.
+
+### `@safe_with_pool` (try-finally)
+
+For code that may throw and needs guaranteed cleanup:
+
+```julia
+@safe_with_pool pool begin
+    acquire!(pool, Float64, 10)
+    risky_operation()          # if this throws, rewind! still runs
+end
+```
+
+This prevents inlining but guarantees pool cleanup regardless of exceptions.
+Use it when:
+- The pool body calls functions that may throw
+- You need the pool to remain valid after a caught exception
+- A custom macro inside the body might generate hidden `return`/`break`/`continue`
+
 ## Type Extraction: Static Analysis at Compile Time
 
 The macro analyzes the AST to extract types used in `acquire!` calls:
@@ -60,11 +150,8 @@ end
 
 # Generated code uses typed checkpoint/rewind:
 checkpoint!(pool, Float64, ComplexF64)
-try
-    ...
-finally
-    rewind!(pool, Float64, ComplexF64)
-end
+# ... body with rewind! at each exit ...
+rewind!(pool, Float64, ComplexF64)
 ```
 
 ### Type Extraction Rules
@@ -227,6 +314,7 @@ end
 # OUTPUT (simplified)
 function compute(data)
     pool = get_task_local_pool()
+    _entry_depth = pool._current_depth
 
     # Bitmask subset check: can typed path handle any untracked acquires?
     if _can_use_typed_path(pool, _tracked_mask_for_types(Float64))
@@ -235,17 +323,22 @@ function compute(data)
         checkpoint!(pool)                    # Full checkpoint (safe)
     end
 
-    try
-        A = _acquire_impl!(pool, Float64, length(data))
-        result = helper!(pool, A)
-        return result
-    finally
-        if _can_use_typed_path(pool, _tracked_mask_for_types(Float64))
-            rewind!(pool, Float64)           # Typed rewind (fast)
-        else
-            rewind!(pool)                    # Full rewind (safe)
-        end
+    A = _acquire_impl!(pool, Float64, length(data))
+    _result = helper!(pool, A)
+
+    # Entry depth guard: clean up any leaked inner scopes
+    while pool._current_depth > _entry_depth + 1
+        rewind!(pool)
     end
+
+    # Own scope rewind
+    if _can_use_typed_path(pool, _tracked_mask_for_types(Float64))
+        rewind!(pool, Float64)              # Typed rewind (fast)
+    else
+        rewind!(pool)                       # Full rewind (safe)
+    end
+
+    return _result
 end
 ```
 
@@ -256,6 +349,9 @@ end
 | `_extract_acquire_types(expr, pool_name)` | AST walk to find types |
 | `_filter_static_types(types, local_vars)` | Filter out locally-defined types |
 | `_transform_acquire_calls(expr, pool_name)` | Replace `acquire!` → `_acquire_impl!` |
+| `_transform_return_stmts(expr, ...)` | Insert `rewind!` before each `return` |
+| `_transform_break_continue(expr, ...)` | Insert `rewind!` before `break`/`continue` |
+| `_check_unsafe_goto(expr)` | Hard error on `@goto` that exits pool scope |
 | `_record_type_touch!(pool, T)` | Record type touch in bitmask for current depth |
 | `_can_use_typed_path(pool, mask)` | Bitmask subset check for typed vs full path |
 | `_tracked_mask_for_types(T...)` | Compile-time bitmask for tracked types |

src/AdaptiveArrayPools.jl

@@ -8,7 +8,7 @@ export acquire_view!, acquire_array!  # Explicit naming aliases
 export zeros!, ones!, trues!, falses!, similar!, reshape!, default_eltype  # Convenience functions
 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 @with_pool, @maybe_with_pool, @safe_with_pool, @safe_maybe_with_pool
 export STATIC_POOLING, MAYBE_POOLING, RUNTIME_CHECK
 export PoolEscapeError, EscapePoint
 export checkpoint!, rewind!, reset!

src/debug.jl

@@ -223,6 +223,25 @@ _validate_pool_return(val, ::DisabledPool) = nothing
 # No-op fallback for pool types without specific validation (overridden by CUDA extension)
 _validate_pool_return(val, ::AbstractArrayPool) = nothing
 
+# ==============================================================================
+# Leaked Scope Warning (direct-rewind path, RUNTIME_CHECK >= 1)
+# ==============================================================================
+#
+# Detects when entry depth guard fires (inner scope didn't rewind properly).
+# @noinline to keep it out of the inlined hot path — only called on error.
+
+@noinline function _warn_leaked_scope(pool::AbstractArrayPool, entry_depth::Int)
+    return @error(
+        "Leaked @with_pool scope detected! " *
+            "Pool depth is $(pool._current_depth), expected $(entry_depth + 1). " *
+            "A macro inside @with_pool may have generated an unseen `return`/`break`, " *
+            "or an inner scope threw without try-finally protection. " *
+            "Consider using @safe_with_pool for exception safety.",
+        current_depth = pool._current_depth,
+        expected_depth = entry_depth + 1,
+    )
+end
+
 # ==============================================================================
 # Poisoning: Fill released vectors with sentinel values (S >= 1)
 # ==============================================================================