From b30d8541488e75d3e8b7c6a3e1897cf20e5528cf Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 19:05:37 -0800
Subject: [PATCH 1/9] feat: add reset! function for state-only pool reset

Add reset!(pool) that clears all n_active counters and checkpoint stacks
to sentinel state while preserving allocated vectors for reuse.

Unlike empty! which releases all memory, reset! allows continued use of
pre-allocated storage - useful when functions acquire without proper
checkpoint/rewind management.
---
 src/AdaptiveArrayPools.jl |  4 +-
 src/state.jl              | 85 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/src/AdaptiveArrayPools.jl b/src/AdaptiveArrayPools.jl
index e51e810..24bc999 100644
--- a/src/AdaptiveArrayPools.jl
+++ b/src/AdaptiveArrayPools.jl
@@ -6,7 +6,7 @@ export AdaptiveArrayPool, acquire!, unsafe_acquire!, pool_stats, get_task_local_
 export acquire_view!, acquire_array!  # Explicit naming aliases
 export @with_pool, @maybe_with_pool
 export USE_POOLING, MAYBE_POOLING_ENABLED, POOL_DEBUG
-export checkpoint!, rewind!
+export checkpoint!, rewind!, reset!
 export CACHE_WAYS, set_cache_ways!  # N-way cache configuration
 
 # Core data structures
@@ -18,7 +18,7 @@ include("utils.jl")
 # Acquisition operations: get_view!, acquire!, unsafe_acquire!, aliases
 include("acquire.jl")
 
-# State management: checkpoint!, rewind!, empty!
+# State management: checkpoint!, rewind!, reset!, empty!
 include("state.jl")
 
 # Task-local pool
diff --git a/src/state.jl b/src/state.jl
index f842020..0a48776 100644
--- a/src/state.jl
+++ b/src/state.jl
@@ -291,3 +291,88 @@ function Base.empty!(pool::AdaptiveArrayPool)
 end
 
 Base.empty!(::Nothing) = nothing
+
+# ==============================================================================
+# Pool Reset (preserve storage, reset state)
+# ==============================================================================
+
+"""
+    reset!(tp::TypedPool)
+
+Reset TypedPool state without clearing allocated storage.
+
+Sets `n_active = 0` and restores checkpoint stacks to sentinel state.
+All vectors, views, and N-D arrays are preserved for reuse.
+
+This is useful when you want to "start fresh" without reallocating memory.
+"""
+function reset!(tp::TypedPool)
+    tp.n_active = 0
+    # Restore sentinel values (1-based sentinel pattern)
+    empty!(tp._checkpoint_n_active)
+    push!(tp._checkpoint_n_active, 0)   # Sentinel: n_active=0 at depth=0
+    empty!(tp._checkpoint_depths)
+    push!(tp._checkpoint_depths, 0)     # Sentinel: depth=0 = no checkpoint
+    return tp
+end
+
+"""
+    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 `_untracked_flags`
+
+Unlike `empty!`, this **preserves** all allocated vectors, views, and N-D arrays
+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
+```julia
+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!`](@ref), [`rewind!`](@ref)
+"""
+function reset!(pool::AdaptiveArrayPool)
+    # Fixed slots - zero allocation via @generated iteration
+    foreach_fixed_slot(pool) do tp
+        reset!(tp)
+    end
+
+    # Others - reset all TypedPools
+    for tp in values(pool.others)
+        reset!(tp)
+    end
+
+    # Reset untracked detection state (1-based sentinel pattern)
+    pool._current_depth = 1                   # 1 = global scope (sentinel)
+    empty!(pool._untracked_flags)
+    push!(pool._untracked_flags, false)       # Sentinel: global scope starts with false
+
+    return pool
+end
+
+reset!(::Nothing) = nothing

From d64370d3fee8e88735ff6ebd3193d46d6d7174f3 Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 19:08:28 -0800
Subject: [PATCH 2/9] test: add comprehensive tests for reset! function

Cover 10 test cases including:
- Basic n_active reset to zero
- Vector preservation after reset
- Checkpoint stack restoration to sentinel
- Fallback types (others dict) handling
- nothing compatibility
- Pool usability after reset
- A/B scenario (unmanaged acquires then reset)
- reset! vs empty! comparison
- TypedPool direct reset
---
 test/test_state.jl | 194 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 194 insertions(+)

diff --git a/test/test_state.jl b/test/test_state.jl
index 2955040..a143c65 100644
--- a/test/test_state.jl
+++ b/test/test_state.jl
@@ -243,6 +243,200 @@
         rewind!(pool)
     end
 
+    @testset "reset! (state-only reset)" begin
+        import AdaptiveArrayPools: reset!
+
+        @testset "basic reset! - n_active to zero" begin
+            pool = AdaptiveArrayPool()
+
+            # Acquire some arrays
+            v1 = acquire!(pool, Float64, 100)
+            v2 = acquire!(pool, Float32, 50)
+            v3 = acquire!(pool, Int64, 30)
+            @test pool.float64.n_active == 1
+            @test pool.float32.n_active == 1
+            @test pool.int64.n_active == 1
+
+            # Reset
+            result = reset!(pool)
+            @test result === pool  # Returns self
+
+            # All n_active should be 0
+            @test pool.float64.n_active == 0
+            @test pool.float32.n_active == 0
+            @test pool.int64.n_active == 0
+        end
+
+        @testset "reset! preserves vectors" begin
+            pool = AdaptiveArrayPool()
+
+            # Acquire arrays
+            v1 = acquire!(pool, Float64, 100)
+            v2 = acquire!(pool, Float64, 200)
+            v3 = acquire!(pool, Float64, 50)
+            @test length(pool.float64.vectors) == 3
+
+            # Reset - should preserve vectors
+            reset!(pool)
+            @test pool.float64.n_active == 0
+            @test length(pool.float64.vectors) == 3  # Vectors preserved!
+            @test length(pool.float64.vectors[1]) >= 100
+            @test length(pool.float64.vectors[2]) >= 200
+            @test length(pool.float64.vectors[3]) >= 50
+        end
+
+        @testset "reset! restores checkpoint stacks to sentinel" begin
+            pool = AdaptiveArrayPool()
+
+            # Nested checkpoints
+            checkpoint!(pool)
+            acquire!(pool, Float64, 10)
+            checkpoint!(pool)
+            acquire!(pool, Float64, 20)
+            checkpoint!(pool)
+            acquire!(pool, Float64, 30)
+
+            @test pool._current_depth == 4
+            @test length(pool.float64._checkpoint_n_active) > 1
+            @test length(pool.float64._checkpoint_depths) > 1
+
+            # Reset - should restore sentinel state
+            reset!(pool)
+
+            @test pool._current_depth == 1
+            @test pool._untracked_flags == [false]
+            @test pool.float64._checkpoint_n_active == [0]  # Sentinel only
+            @test pool.float64._checkpoint_depths == [0]    # Sentinel only
+        end
+
+        @testset "reset! with fallback types" begin
+            pool = AdaptiveArrayPool()
+
+            # Use fallback type (not in fixed slots)
+            v1 = acquire!(pool, UInt8, 100)
+            v2 = acquire!(pool, UInt16, 50)
+            @test pool.others[UInt8].n_active == 1
+            @test pool.others[UInt16].n_active == 1
+            @test length(pool.others[UInt8].vectors) == 1
+
+            # Reset
+            reset!(pool)
+
+            # n_active reset but vectors preserved
+            @test pool.others[UInt8].n_active == 0
+            @test pool.others[UInt16].n_active == 0
+            @test length(pool.others[UInt8].vectors) == 1  # Preserved!
+            @test length(pool.others[UInt16].vectors) == 1
+        end
+
+        @testset "reset!(nothing) compatibility" begin
+            @test reset!(nothing) === nothing
+        end
+
+        @testset "pool usable after reset!" begin
+            pool = AdaptiveArrayPool()
+
+            # First use
+            v1 = acquire!(pool, Float64, 100)
+            v1 .= 42.0
+            backing1 = parent(v1)
+
+            # Reset
+            reset!(pool)
+
+            # Should be usable and reuse existing vector
+            checkpoint!(pool)
+            v2 = acquire!(pool, Float64, 100)
+            @test parent(v2) === backing1  # Same backing vector reused
+            @test pool.float64.n_active == 1
+            rewind!(pool)
+            @test pool.float64.n_active == 0
+        end
+
+        @testset "A/B scenario - unmanaged then reset" begin
+            # Simulates: inner function acquires without management,
+            # outer function calls reset! to clean up
+
+            pool = AdaptiveArrayPool()
+
+            # Function that acquires without checkpoint/rewind
+            function unmanaged_compute!(p)
+                v = acquire!(p, Float64, 100)
+                v .= 1.0
+                # No rewind!
+            end
+
+            # Call multiple times - n_active grows
+            for _ in 1:10
+                unmanaged_compute!(pool)
+            end
+            @test pool.float64.n_active == 10
+            @test length(pool.float64.vectors) == 10
+
+            # Reset - clean slate but vectors preserved
+            reset!(pool)
+            @test pool.float64.n_active == 0
+            @test length(pool.float64.vectors) == 10  # All preserved for reuse
+
+            # Next use reuses existing vectors
+            checkpoint!(pool)
+            for _ in 1:5
+                acquire!(pool, Float64, 100)
+            end
+            @test pool.float64.n_active == 5
+            @test length(pool.float64.vectors) == 10  # No new allocations
+            rewind!(pool)
+        end
+
+        @testset "reset! vs empty! comparison" begin
+            # Verify reset! preserves while empty! clears
+
+            pool1 = AdaptiveArrayPool()
+            pool2 = AdaptiveArrayPool()
+
+            # Both acquire same arrays
+            for _ in 1:5
+                acquire!(pool1, Float64, 100)
+                acquire!(pool2, Float64, 100)
+            end
+            @test length(pool1.float64.vectors) == 5
+            @test length(pool2.float64.vectors) == 5
+
+            # reset! preserves
+            reset!(pool1)
+            @test pool1.float64.n_active == 0
+            @test length(pool1.float64.vectors) == 5  # Preserved
+
+            # empty! clears
+            empty!(pool2)
+            @test pool2.float64.n_active == 0
+            @test length(pool2.float64.vectors) == 0  # Cleared
+        end
+
+        @testset "TypedPool reset!" begin
+            import AdaptiveArrayPools: get_typed_pool!
+
+            pool = AdaptiveArrayPool()
+            tp = get_typed_pool!(pool, Float64)
+
+            # Acquire and checkpoint
+            checkpoint!(tp)
+            acquire!(pool, Float64, 100)
+            checkpoint!(tp)
+            acquire!(pool, Float64, 200)
+            @test tp.n_active == 2
+            @test length(tp._checkpoint_n_active) > 1
+
+            # Reset TypedPool directly
+            result = reset!(tp)
+            @test result === tp
+            @test tp.n_active == 0
+            @test tp._checkpoint_n_active == [0]
+            @test tp._checkpoint_depths == [0]
+            @test length(tp.vectors) == 2  # Vectors preserved
+        end
+    end
+
     @testset "Typed checkpoint!/rewind! (generated functions)" begin
         pool = AdaptiveArrayPool()
 

From bbcd9442633273d335844ed435eb36c6751ee26e Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 19:19:02 -0800
Subject: [PATCH 3/9] feat: add safety guard to rewind! at global scope

When rewind! is called at depth=1 (no pending checkpoints), it now
safely delegates to reset! instead of breaking.

This enables:
- Safe rewind! without matching checkpoint
- Safe rewind! after reset!
- Multiple consecutive rewind! calls
- A/B scenario cleanup with just rewind!
---
 src/state.jl       | 12 ++++++++-
 test/test_state.jl | 67 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/state.jl b/src/state.jl
index 0a48776..049d7cc 100644
--- a/src/state.jl
+++ b/src/state.jl
@@ -49,11 +49,21 @@ Uses _checkpoint_depths to accurately determine which entries to pop vs restore.
 Only the counters are restored; allocated memory remains for reuse.
 Handles untracked acquires by checking _checkpoint_depths for accurate restoration.
 
-See also: [`checkpoint!`](@ref), [`@with_pool`](@ref)
+**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!`](@ref), [`reset!`](@ref), [`@with_pool`](@ref)
 """
 function rewind!(pool::AdaptiveArrayPool)
     cur_depth = pool._current_depth
 
+    # Safety guard: at global scope (depth=1), no checkpoint to rewind to
+    # Delegate to reset! which safely clears all n_active counters
+    if cur_depth == 1
+        reset!(pool)
+        return nothing
+    end
+
     # Fixed slots - zero allocation via @generated iteration
     foreach_fixed_slot(pool) do tp
         _rewind_typed_pool!(tp, cur_depth)
diff --git a/test/test_state.jl b/test/test_state.jl
index a143c65..2a4e748 100644
--- a/test/test_state.jl
+++ b/test/test_state.jl
@@ -437,6 +437,73 @@
         end
     end
 
+    @testset "Safe rewind! at depth=1" begin
+        @testset "rewind! without checkpoint (depth=1)" begin
+            pool = AdaptiveArrayPool()
+
+            # Acquire without checkpoint
+            v1 = acquire!(pool, Float64, 100)
+            v2 = acquire!(pool, Float64, 200)
+            @test pool.float64.n_active == 2
+            @test pool._current_depth == 1
+
+            # rewind! at depth=1 should be safe (delegates to reset!)
+            rewind!(pool)
+            @test pool.float64.n_active == 0
+            @test pool._current_depth == 1
+            @test pool._untracked_flags == [false]
+        end
+
+        @testset "rewind! after reset!" begin
+            pool = AdaptiveArrayPool()
+
+            checkpoint!(pool)
+            acquire!(pool, Float64, 100)
+            @test pool._current_depth == 2
+
+            # Reset clears everything
+            reset!(pool)
+            @test pool._current_depth == 1
+
+            # rewind! after reset should be safe
+            rewind!(pool)
+            @test pool._current_depth == 1  # Still at global scope
+        end
+
+        @testset "checkpoint/rewind cycle after reset!" begin
+            pool = AdaptiveArrayPool()
+
+            # Initial acquire and reset
+            acquire!(pool, Float64, 50)
+            reset!(pool)
+
+            # Normal checkpoint/rewind should work
+            checkpoint!(pool)
+            @test pool._current_depth == 2
+            acquire!(pool, Float64, 100)
+            @test pool.float64.n_active == 1
+
+            rewind!(pool)
+            @test pool.float64.n_active == 0
+            @test pool._current_depth == 1
+        end
+
+        @testset "multiple rewind! at depth=1 is safe" begin
+            pool = AdaptiveArrayPool()
+
+            acquire!(pool, Float64, 100)
+            @test pool.float64.n_active == 1
+
+            # Multiple rewind! calls should all be safe
+            rewind!(pool)
+            @test pool.float64.n_active == 0
+            rewind!(pool)
+            @test pool.float64.n_active == 0
+            rewind!(pool)
+            @test pool._current_depth == 1
+        end
+    end
+
     @testset "Typed checkpoint!/rewind! (generated functions)" begin
         pool = AdaptiveArrayPool()
 

From e8ef12c99afc61b20be47f5b789231b986bdaa52 Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 19:50:21 -0800
Subject: [PATCH 4/9] feat: add type-specific reset! and safety guards for
 typed rewind!

- Add reset!(pool, Type) and reset!(pool, Type...) for API consistency
- Add depth=1 safety guard to rewind!(pool, Type) and rewind!(pool, Type...)
- Add tests for type-specific reset! and safe typed rewind! behavior
---
 src/state.jl       | 42 ++++++++++++++++++++++++
 test/test_state.jl | 82 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 124 insertions(+)

diff --git a/src/state.jl b/src/state.jl
index 049d7cc..6151aeb 100644
--- a/src/state.jl
+++ b/src/state.jl
@@ -188,6 +188,11 @@ Restore state for a specific type only.
 Also updates _current_depth and _untracked_flags.
 """
 @inline function rewind!(pool::AdaptiveArrayPool, ::Type{T}) where T
+    # Safety guard: at global scope (depth=1), delegate to reset!
+    if pool._current_depth == 1
+        reset!(get_typed_pool!(pool, T))
+        return nothing
+    end
     rewind!(get_typed_pool!(pool, T))
     pop!(pool._untracked_flags)
     pool._current_depth -= 1
@@ -222,7 +227,13 @@ Decrements _current_depth once after all types are rewound.
 @generated function rewind!(pool::AdaptiveArrayPool, types::Type...)
     # Reverse order for proper stack unwinding, rewind TypedPools directly
     rewind_exprs = [:(rewind!(get_typed_pool!(pool, types[$i]))) for i in length(types):-1:1]
+    reset_exprs = [:(reset!(get_typed_pool!(pool, types[$i]))) for i in 1:length(types)]
     quote
+        # Safety guard: at global scope (depth=1), delegate to reset!
+        if pool._current_depth == 1
+            $(reset_exprs...)
+            return nothing
+        end
         $(rewind_exprs...)
         pop!(pool._untracked_flags)
         pool._current_depth -= 1
@@ -386,3 +397,34 @@ function reset!(pool::AdaptiveArrayPool)
 end
 
 reset!(::Nothing) = nothing
+
+"""
+    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)`](@ref), [`rewind!`](@ref)
+"""
+@inline function reset!(pool::AdaptiveArrayPool, ::Type{T}) where T
+    reset!(get_typed_pool!(pool, T))
+end
+
+"""
+    reset!(pool::AdaptiveArrayPool, types::Type...)
+
+Reset state for multiple specific types. Uses @generated for zero-overhead
+compile-time unrolling.
+
+See also: [`reset!(::AdaptiveArrayPool)`](@ref), [`rewind!`](@ref)
+"""
+@generated function reset!(pool::AdaptiveArrayPool, types::Type...)
+    reset_exprs = [:(reset!(get_typed_pool!(pool, types[$i]))) for i in 1:length(types)]
+    quote
+        $(reset_exprs...)
+        nothing
+    end
+end
+
+reset!(::Nothing, ::Type) = nothing
+reset!(::Nothing, types::Type...) = nothing
diff --git a/test/test_state.jl b/test/test_state.jl
index 2a4e748..2f7a28e 100644
--- a/test/test_state.jl
+++ b/test/test_state.jl
@@ -435,6 +435,46 @@
             @test tp._checkpoint_depths == [0]
             @test length(tp.vectors) == 2  # Vectors preserved
         end
+
+        @testset "type-specific reset!(pool, Type)" begin
+            pool = AdaptiveArrayPool()
+
+            # Acquire multiple types
+            acquire!(pool, Float64, 100)
+            acquire!(pool, Int64, 50)
+            @test pool.float64.n_active == 1
+            @test pool.int64.n_active == 1
+
+            # Reset only Float64
+            reset!(pool, Float64)
+            @test pool.float64.n_active == 0
+            @test pool.int64.n_active == 1  # Int64 unchanged
+            @test pool.float64._checkpoint_n_active == [0]
+            @test length(pool.float64.vectors) == 1  # Vector preserved
+        end
+
+        @testset "type-specific reset!(pool, Type...)" begin
+            pool = AdaptiveArrayPool()
+
+            # Acquire multiple types
+            acquire!(pool, Float64, 100)
+            acquire!(pool, Int64, 50)
+            acquire!(pool, Float32, 25)
+            @test pool.float64.n_active == 1
+            @test pool.int64.n_active == 1
+            @test pool.float32.n_active == 1
+
+            # Reset Float64 and Int64, but not Float32
+            reset!(pool, Float64, Int64)
+            @test pool.float64.n_active == 0
+            @test pool.int64.n_active == 0
+            @test pool.float32.n_active == 1  # Float32 unchanged
+        end
+
+        @testset "reset!(nothing, Type) compatibility" begin
+            @test reset!(nothing, Float64) === nothing
+            @test reset!(nothing, Float64, Int64) === nothing
+        end
     end
 
     @testset "Safe rewind! at depth=1" begin
@@ -502,6 +542,48 @@
             rewind!(pool)
             @test pool._current_depth == 1
         end
+
+        @testset "type-specific rewind!(pool, Type) at depth=1" begin
+            pool = AdaptiveArrayPool()
+
+            # Acquire without checkpoint
+            v1 = acquire!(pool, Float64, 100)
+            @test pool.float64.n_active == 1
+            @test pool._current_depth == 1
+
+            # Type-specific rewind! at depth=1 should be safe
+            rewind!(pool, Float64)
+            @test pool.float64.n_active == 0
+            @test pool._current_depth == 1  # Should not go to 0
+            @test pool.float64._checkpoint_n_active == [0]  # Sentinel preserved
+
+            # Multiple calls should be safe
+            rewind!(pool, Float64)
+            @test pool._current_depth == 1
+        end
+
+        @testset "type-specific rewind!(pool, Type...) at depth=1" begin
+            pool = AdaptiveArrayPool()
+
+            # Acquire multiple types without checkpoint
+            v1 = acquire!(pool, Float64, 100)
+            v2 = acquire!(pool, Int64, 50)
+            @test pool.float64.n_active == 1
+            @test pool.int64.n_active == 1
+            @test pool._current_depth == 1
+
+            # Multi-type rewind! at depth=1 should be safe
+            rewind!(pool, Float64, Int64)
+            @test pool.float64.n_active == 0
+            @test pool.int64.n_active == 0
+            @test pool._current_depth == 1  # Should not go to 0
+            @test pool.float64._checkpoint_n_active == [0]
+            @test pool.int64._checkpoint_n_active == [0]
+
+            # Multiple calls should be safe
+            rewind!(pool, Float64, Int64)
+            @test pool._current_depth == 1
+        end
     end
 
     @testset "Typed checkpoint!/rewind! (generated functions)" begin

From 4a4eaef4002ef33fa17108a27f7b58481bd87395 Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 19:53:36 -0800
Subject: [PATCH 5/9] test: verify reset! preserves 1D and N-D caches

---
 test/test_state.jl | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/test/test_state.jl b/test/test_state.jl
index 2f7a28e..f50f947 100644
--- a/test/test_state.jl
+++ b/test/test_state.jl
@@ -267,22 +267,24 @@
             @test pool.int64.n_active == 0
         end
 
-        @testset "reset! preserves vectors" begin
+        @testset "reset! preserves vectors and caches" begin
             pool = AdaptiveArrayPool()
 
-            # Acquire arrays
+            # Acquire arrays (populates vectors)
             v1 = acquire!(pool, Float64, 100)
             v2 = acquire!(pool, Float64, 200)
-            v3 = acquire!(pool, Float64, 50)
-            @test length(pool.float64.vectors) == 3
+            # Use unsafe_acquire! to populate 1D cache
+            v3 = unsafe_acquire!(pool, Float64, 50)
+            # Use N-D acquire to populate nd cache
+            m1 = acquire!(pool, Float64, 10, 10)
+            @test length(pool.float64.vectors) >= 3
 
-            # Reset - should preserve vectors
+            # Reset - should preserve everything
             reset!(pool)
             @test pool.float64.n_active == 0
-            @test length(pool.float64.vectors) == 3  # Vectors preserved!
-            @test length(pool.float64.vectors[1]) >= 100
-            @test length(pool.float64.vectors[2]) >= 200
-            @test length(pool.float64.vectors[3]) >= 50
+            @test length(pool.float64.vectors) >= 3    # Vectors preserved
+            @test length(pool.float64.views) >= 1      # 1D cache preserved
+            @test length(pool.float64.nd_arrays) >= 1  # N-D cache preserved
         end
 
         @testset "reset! restores checkpoint stacks to sentinel" begin

From abc1a736c7fbc3011b5d5a6d02df37b6e095d16f Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 19:57:16 -0800
Subject: [PATCH 6/9] docs: add Internal API warning to reset!(TypedPool)

---
 src/state.jl | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/state.jl b/src/state.jl
index 6151aeb..27b0825 100644
--- a/src/state.jl
+++ b/src/state.jl
@@ -320,12 +320,20 @@ Base.empty!(::Nothing) = nothing
 """
     reset!(tp::TypedPool)
 
-Reset TypedPool state without clearing allocated storage.
+Internal method for resetting TypedPool state without clearing storage.
 
 Sets `n_active = 0` and restores checkpoint stacks to sentinel state.
 All vectors, views, and N-D arrays are preserved for reuse.
 
-This is useful when you want to "start fresh" without reallocating memory.
+!!! warning "Internal API"
+    This is an internal implementation detail. For manual pool management,
+    use the public API instead:
+    ```julia
+    reset!(pool)          # Reset entire pool
+    reset!(pool, Float64) # Reset specific type
+    ```
+
+See also: [`reset!(::AdaptiveArrayPool)`](@ref), [`reset!(::AdaptiveArrayPool, ::Type)`](@ref)
 """
 function reset!(tp::TypedPool)
     tp.n_active = 0

From 21d92740bbe20d7d2cb2320203f7dc24044a1d0a Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 19:58:35 -0800
Subject: [PATCH 7/9] Revert "docs: add Internal API warning to
 reset!(TypedPool)"

This reverts commit abc1a736c7fbc3011b5d5a6d02df37b6e095d16f.
---
 src/state.jl | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/src/state.jl b/src/state.jl
index 27b0825..6151aeb 100644
--- a/src/state.jl
+++ b/src/state.jl
@@ -320,20 +320,12 @@ Base.empty!(::Nothing) = nothing
 """
     reset!(tp::TypedPool)
 
-Internal method for resetting TypedPool state without clearing storage.
+Reset TypedPool state without clearing allocated storage.
 
 Sets `n_active = 0` and restores checkpoint stacks to sentinel state.
 All vectors, views, and N-D arrays are preserved for reuse.
 
-!!! warning "Internal API"
-    This is an internal implementation detail. For manual pool management,
-    use the public API instead:
-    ```julia
-    reset!(pool)          # Reset entire pool
-    reset!(pool, Float64) # Reset specific type
-    ```
-
-See also: [`reset!(::AdaptiveArrayPool)`](@ref), [`reset!(::AdaptiveArrayPool, ::Type)`](@ref)
+This is useful when you want to "start fresh" without reallocating memory.
 """
 function reset!(tp::TypedPool)
     tp.n_active = 0

From d6f07efc7651c68f478c78bd44e9a294ece877be Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 21:11:23 -0800
Subject: [PATCH 8/9] refactor: clean up state.jl and unify rewind logic

- Remove redundant TypedPool public overloads (checkpoint!/rewind!)
- All rewind variants now use _rewind_typed_pool! with orphan cleanup
- Reorganize state.jl into sections: checkpoint!, rewind!, empty!, reset!
- Update tests to use internal helpers (_checkpoint_typed_pool!, _rewind_typed_pool!)
---
 src/state.jl       | 208 +++++++++++++++++----------------------------
 test/test_state.jl |  31 ++++---
 2 files changed, 91 insertions(+), 148 deletions(-)

diff --git a/src/state.jl b/src/state.jl
index 6151aeb..a9bef38 100644
--- a/src/state.jl
+++ b/src/state.jl
@@ -1,5 +1,5 @@
 # ==============================================================================
-# State Management
+# State Management - checkpoint!
 # ==============================================================================
 
 """
@@ -31,14 +31,52 @@ function checkpoint!(pool::AdaptiveArrayPool)
     return nothing
 end
 
-# Internal helper for full checkpoint
+"""
+    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 _current_depth and _untracked_flags for untracked acquire detection.
+
+~77% faster than full checkpoint! when only one type is used.
+"""
+@inline function checkpoint!(pool::AdaptiveArrayPool, ::Type{T}) where T
+    pool._current_depth += 1
+    push!(pool._untracked_flags, false)
+    _checkpoint_typed_pool!(get_typed_pool!(pool, T), pool._current_depth)
+end
+
+"""
+    checkpoint!(pool::AdaptiveArrayPool, types::Type...)
+
+Save state for multiple specific types. Uses @generated for zero-overhead
+compile-time unrolling. Increments _current_depth once for all types.
+"""
+@generated function checkpoint!(pool::AdaptiveArrayPool, types::Type...)
+    checkpoint_exprs = [:(_checkpoint_typed_pool!(get_typed_pool!(pool, types[$i]), pool._current_depth)) for i in 1:length(types)]
+    quote
+        pool._current_depth += 1
+        push!(pool._untracked_flags, false)
+        $(checkpoint_exprs...)
+        nothing
+    end
+end
+
+checkpoint!(::Nothing) = nothing
+checkpoint!(::Nothing, ::Type) = nothing
+checkpoint!(::Nothing, types::Type...) = nothing
+
+# Internal helper for checkpoint
 @inline function _checkpoint_typed_pool!(tp::TypedPool, depth::Int)
     push!(tp._checkpoint_n_active, tp.n_active)
     push!(tp._checkpoint_depths, depth)
     nothing
 end
 
-checkpoint!(::Nothing) = nothing
+# ==============================================================================
+# State Management - rewind!
+# ==============================================================================
 
 """
     rewind!(pool::AdaptiveArrayPool)
@@ -80,107 +118,6 @@ function rewind!(pool::AdaptiveArrayPool)
     return nothing
 end
 
-# Internal helper for full rewind with _checkpoint_depths
-# Uses 1-based sentinel pattern: no isempty checks needed (sentinel [0] guarantees non-empty)
-@inline function _rewind_typed_pool!(tp::TypedPool, current_depth::Int)
-    # 1. Orphaned Checkpoints Cleanup
-    # If there are checkpoints from deeper scopes (depth > current), pop them first.
-    # This happens when a nested scope did full checkpoint but typed rewind,
-    # leaving orphaned checkpoints that must be cleaned before finding current state.
-    while @inbounds tp._checkpoint_depths[end] > current_depth
-        pop!(tp._checkpoint_depths)
-        pop!(tp._checkpoint_n_active)
-    end
-
-    # 2. Normal Rewind Logic (Sentinel Pattern)
-    # Now the stack top is guaranteed to be at depth <= current depth.
-    if @inbounds tp._checkpoint_depths[end] == current_depth
-        # Checkpointed at current depth: pop and restore
-        pop!(tp._checkpoint_depths)
-        tp.n_active = pop!(tp._checkpoint_n_active)
-    else
-        # No checkpoint at current depth (this type was excluded from typed checkpoint)
-        # MUST restore n_active from parent checkpoint value!
-        # - Untracked acquire may have modified n_active
-        # - If sentinel (_checkpoint_n_active=[0]), restores to n_active=0
-        tp.n_active = @inbounds tp._checkpoint_n_active[end]
-    end
-    nothing
-end
-
-rewind!(::Nothing) = nothing
-
-# ==============================================================================
-# Type-Specific State Management (for optimized macros)
-# ==============================================================================
-
-"""
-    checkpoint!(tp::TypedPool)
-
-Internal method for saving TypedPool state (legacy, uses depth=0).
-
-!!! warning "Internal API"
-    This is an internal implementation detail. For manual pool management,
-    use the public API instead:
-    ```julia
-    checkpoint!(pool, Float64)  # Type-specific checkpoint
-    ```
-
-See also: [`checkpoint!(::AdaptiveArrayPool, ::Type)`](@ref), [`rewind!`](@ref)
-"""
-@inline function checkpoint!(tp::TypedPool)
-    push!(tp._checkpoint_n_active, tp.n_active)
-    push!(tp._checkpoint_depths, 0)  # Legacy depth
-    nothing
-end
-
-"""
-    checkpoint!(tp::TypedPool, depth::Int)
-
-Internal method for saving TypedPool state with depth tracking.
-"""
-@inline function checkpoint!(tp::TypedPool, depth::Int)
-    push!(tp._checkpoint_n_active, tp.n_active)
-    push!(tp._checkpoint_depths, depth)
-    nothing
-end
-
-"""
-    rewind!(tp::TypedPool)
-
-Internal method for restoring TypedPool state (pops both stacks).
-
-!!! warning "Internal API"
-    This is an internal implementation detail. For manual pool management,
-    use the public API instead:
-    ```julia
-    rewind!(pool, Float64)  # Type-specific rewind
-    ```
-
-See also: [`rewind!(::AdaptiveArrayPool, ::Type)`](@ref), [`checkpoint!`](@ref)
-"""
-@inline function rewind!(tp::TypedPool)
-    pop!(tp._checkpoint_depths)
-    tp.n_active = pop!(tp._checkpoint_n_active)
-    nothing
-end
-
-"""
-    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 _current_depth and _untracked_flags for untracked acquire detection.
-
-~77% faster than full checkpoint! when only one type is used.
-"""
-@inline function checkpoint!(pool::AdaptiveArrayPool, ::Type{T}) where T
-    pool._current_depth += 1
-    push!(pool._untracked_flags, false)
-    checkpoint!(get_typed_pool!(pool, T), pool._current_depth)
-end
-
 """
     rewind!(pool::AdaptiveArrayPool, ::Type{T})
 
@@ -193,31 +130,11 @@ Also updates _current_depth and _untracked_flags.
         reset!(get_typed_pool!(pool, T))
         return nothing
     end
-    rewind!(get_typed_pool!(pool, T))
+    _rewind_typed_pool!(get_typed_pool!(pool, T), pool._current_depth)
     pop!(pool._untracked_flags)
     pool._current_depth -= 1
 end
 
-checkpoint!(::Nothing, ::Type) = nothing
-rewind!(::Nothing, ::Type) = nothing
-
-"""
-    checkpoint!(pool::AdaptiveArrayPool, types::Type...)
-
-Save state for multiple specific types. Uses @generated for zero-overhead
-compile-time unrolling. Increments _current_depth once for all types.
-"""
-@generated function checkpoint!(pool::AdaptiveArrayPool, types::Type...)
-    # First increment depth, then checkpoint each type with that depth
-    checkpoint_exprs = [:(checkpoint!(get_typed_pool!(pool, types[$i]), pool._current_depth)) for i in 1:length(types)]
-    quote
-        pool._current_depth += 1
-        push!(pool._untracked_flags, false)
-        $(checkpoint_exprs...)
-        nothing
-    end
-end
-
 """
     rewind!(pool::AdaptiveArrayPool, types::Type...)
 
@@ -225,8 +142,7 @@ Restore state for multiple specific types in reverse order.
 Decrements _current_depth once after all types are rewound.
 """
 @generated function rewind!(pool::AdaptiveArrayPool, types::Type...)
-    # Reverse order for proper stack unwinding, rewind TypedPools directly
-    rewind_exprs = [:(rewind!(get_typed_pool!(pool, types[$i]))) for i in length(types):-1:1]
+    rewind_exprs = [:(_rewind_typed_pool!(get_typed_pool!(pool, types[$i]), pool._current_depth)) for i in length(types):-1:1]
     reset_exprs = [:(reset!(get_typed_pool!(pool, types[$i]))) for i in 1:length(types)]
     quote
         # Safety guard: at global scope (depth=1), delegate to reset!
@@ -241,11 +157,40 @@ Decrements _current_depth once after all types are rewound.
     end
 end
 
-checkpoint!(::Nothing, types::Type...) = nothing
+rewind!(::Nothing) = nothing
+rewind!(::Nothing, ::Type) = nothing
 rewind!(::Nothing, types::Type...) = nothing
 
+# Internal helper for rewind with orphan cleanup
+# Uses 1-based sentinel pattern: no isempty checks needed (sentinel [0] guarantees non-empty)
+@inline function _rewind_typed_pool!(tp::TypedPool, current_depth::Int)
+    # 1. Orphaned Checkpoints Cleanup
+    # If there are checkpoints from deeper scopes (depth > current), pop them first.
+    # This happens when a nested scope did full checkpoint but typed rewind,
+    # leaving orphaned checkpoints that must be cleaned before finding current state.
+    while @inbounds tp._checkpoint_depths[end] > current_depth
+        pop!(tp._checkpoint_depths)
+        pop!(tp._checkpoint_n_active)
+    end
+
+    # 2. Normal Rewind Logic (Sentinel Pattern)
+    # Now the stack top is guaranteed to be at depth <= current depth.
+    if @inbounds tp._checkpoint_depths[end] == current_depth
+        # Checkpointed at current depth: pop and restore
+        pop!(tp._checkpoint_depths)
+        tp.n_active = pop!(tp._checkpoint_n_active)
+    else
+        # No checkpoint at current depth (this type was excluded from typed checkpoint)
+        # MUST restore n_active from parent checkpoint value!
+        # - Untracked acquire may have modified n_active
+        # - If sentinel (_checkpoint_n_active=[0]), restores to n_active=0
+        tp.n_active = @inbounds tp._checkpoint_n_active[end]
+    end
+    nothing
+end
+
 # ==============================================================================
-# Pool Clearing
+# State Management - empty!
 # ==============================================================================
 
 """
@@ -314,7 +259,7 @@ end
 Base.empty!(::Nothing) = nothing
 
 # ==============================================================================
-# Pool Reset (preserve storage, reset state)
+# State Management - reset!
 # ==============================================================================
 
 """
@@ -396,8 +341,6 @@ function reset!(pool::AdaptiveArrayPool)
     return pool
 end
 
-reset!(::Nothing) = nothing
-
 """
     reset!(pool::AdaptiveArrayPool, ::Type{T})
 
@@ -426,5 +369,6 @@ See also: [`reset!(::AdaptiveArrayPool)`](@ref), [`rewind!`](@ref)
     end
 end
 
+reset!(::Nothing) = nothing
 reset!(::Nothing, ::Type) = nothing
 reset!(::Nothing, types::Type...) = nothing
diff --git a/test/test_state.jl b/test/test_state.jl
index f50f947..0ba05de 100644
--- a/test/test_state.jl
+++ b/test/test_state.jl
@@ -416,15 +416,15 @@
         end
 
         @testset "TypedPool reset!" begin
-            import AdaptiveArrayPools: get_typed_pool!
+            import AdaptiveArrayPools: get_typed_pool!, _checkpoint_typed_pool!
 
             pool = AdaptiveArrayPool()
             tp = get_typed_pool!(pool, Float64)
 
-            # Acquire and checkpoint
-            checkpoint!(tp)
+            # Acquire and checkpoint using internal helper
+            _checkpoint_typed_pool!(tp, 1)
             acquire!(pool, Float64, 100)
-            checkpoint!(tp)
+            _checkpoint_typed_pool!(tp, 2)
             acquire!(pool, Float64, 200)
             @test tp.n_active == 2
             @test length(tp._checkpoint_n_active) > 1
@@ -653,47 +653,46 @@
         @test rewind!(nothing, Float64, Int64) === nothing
     end
 
-    @testset "Direct TypedPool checkpoint!/rewind!" begin
-        import AdaptiveArrayPools: get_typed_pool!
+    @testset "Internal TypedPool helpers" begin
+        import AdaptiveArrayPools: get_typed_pool!, _checkpoint_typed_pool!, _rewind_typed_pool!
         pool = AdaptiveArrayPool()
 
         # Get TypedPool directly
         tp = get_typed_pool!(pool, Float64)
         @test tp.n_active == 0
 
-        # Direct TypedPool checkpoint and rewind
-        checkpoint!(tp)
+        # Direct TypedPool checkpoint and rewind using internal helpers
+        _checkpoint_typed_pool!(tp, 1)
         v1 = acquire!(pool, Float64, 100)
         @test tp.n_active == 1
         v2 = acquire!(pool, Float64, 200)
         @test tp.n_active == 2
-        rewind!(tp)
+        _rewind_typed_pool!(tp, 1)
         @test tp.n_active == 0
 
         # Nested checkpoint/rewind on TypedPool
-        checkpoint!(tp)
+        _checkpoint_typed_pool!(tp, 1)
         v1 = acquire!(pool, Float64, 10)
         @test tp.n_active == 1
 
-        checkpoint!(tp)
+        _checkpoint_typed_pool!(tp, 2)
         v2 = acquire!(pool, Float64, 20)
         @test tp.n_active == 2
 
-        checkpoint!(tp)
+        _checkpoint_typed_pool!(tp, 3)
         v3 = acquire!(pool, Float64, 30)
         @test tp.n_active == 3
 
-        rewind!(tp)
+        _rewind_typed_pool!(tp, 3)
         @test tp.n_active == 2
 
-        rewind!(tp)
+        _rewind_typed_pool!(tp, 2)
         @test tp.n_active == 1
 
-        rewind!(tp)
+        _rewind_typed_pool!(tp, 1)
         @test tp.n_active == 0
 
         # Verify type-specific checkpoint delegates to TypedPool
-        # (This tests the refactored implementation)
         checkpoint!(pool, Float64)
         v = acquire!(pool, Float64, 50)
         @test tp.n_active == 1

From 4305ac31636a06b0d92aed2a2a7737ccc9698ffb Mon Sep 17 00:00:00 2001
From: Min-Gu Yoo <mgyoo86@gmail.com>
Date: Wed, 10 Dec 2025 21:28:36 -0800
Subject: [PATCH 9/9] fix: ensure consistent return values for type-specific
 methods

- checkpoint!(pool, Type): explicitly return nothing
- rewind!(pool, Type): explicitly return nothing
- reset!(pool, Type): return pool (not TypedPool)
- reset!(pool, Type...): return pool (not nothing)
---
 src/state.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/state.jl b/src/state.jl
index a9bef38..708770c 100644
--- a/src/state.jl
+++ b/src/state.jl
@@ -45,6 +45,7 @@ Also updates _current_depth and _untracked_flags for untracked acquire detection
     pool._current_depth += 1
     push!(pool._untracked_flags, false)
     _checkpoint_typed_pool!(get_typed_pool!(pool, T), pool._current_depth)
+    nothing
 end
 
 """
@@ -133,6 +134,7 @@ Also updates _current_depth and _untracked_flags.
     _rewind_typed_pool!(get_typed_pool!(pool, T), pool._current_depth)
     pop!(pool._untracked_flags)
     pool._current_depth -= 1
+    nothing
 end
 
 """
@@ -351,6 +353,7 @@ See also: [`reset!(::AdaptiveArrayPool)`](@ref), [`rewind!`](@ref)
 """
 @inline function reset!(pool::AdaptiveArrayPool, ::Type{T}) where T
     reset!(get_typed_pool!(pool, T))
+    pool
 end
 
 """
@@ -365,7 +368,7 @@ See also: [`reset!(::AdaptiveArrayPool)`](@ref), [`rewind!`](@ref)
     reset_exprs = [:(reset!(get_typed_pool!(pool, types[$i]))) for i in 1:length(types)]
     quote
         $(reset_exprs...)
-        nothing
+        pool
     end
 end