From 5c3d1e316f9c93475be81c5d3ad76971d64b9219 Mon Sep 17 00:00:00 2001
From: leburgel <lander.burgelman@ugent.be>
Date: Thu, 10 Sep 2026 16:52:04 +0200
Subject: [PATCH] Insert allocator checkpoints for `@planar` calls

---
 src/planar/macros.jl   |  2 ++
 test/tensors/planar.jl | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/src/planar/macros.jl b/src/planar/macros.jl
index 16f7f1c0a..1fc30ea06 100644
--- a/src/planar/macros.jl
+++ b/src/planar/macros.jl
@@ -44,6 +44,8 @@ function planarparser(planarexpr, kwargs...)
             push!(parser.postprocessors, ex -> insertplanarallocator(ex, allocator))
             # the alloc/free calls are still `GlobalRef(TensorOperations, ...)`
             push!(parser.postprocessors, ex -> TO.insertallocator(ex, allocator))
+            # add checkpoints around the block
+            push!(parser.postprocessors, ex -> TO.insertcheckpoints(ex, allocator))
             break
         end
     end
diff --git a/test/tensors/planar.jl b/test/tensors/planar.jl
index 99ae49ebd..b5d0bbd45 100644
--- a/test/tensors/planar.jl
+++ b/test/tensors/planar.jl
@@ -128,6 +128,34 @@ end
         @test occursin("DefaultBackend", string(ex))
     end
 
+    @testset "allocator is rewound" begin
+        # A `BufferAllocator` hands out slices of a single buffer and reclaims them only
+        # by rewinding its offset -- `tensorfree!` is a no-op for it. The temporaries a
+        # block creates for intermediate results are released that way, so without a
+        # checkpoint/reset pair around the block their space is never reclaimed. A buffer
+        # that is not fully drained also never resizes itself, so it would stay pinned at
+        # whatever size it first grew to, and every later temporary would fall back on the
+        # garbage collector.
+        for W in (ℂ^4, Vect[FermionParity](0 => 2, 1 => 2))
+            A = rand(T, W ← W ⊗ W)
+            B = rand(T, W ⊗ W ← W)
+            @planar Cref[i; j] := A[i; k l] * τ[k l; m n] * B[m n; j]
+
+            # three tensors, so the first contraction is an intermediate temporary
+            buffer = TensorOperations.BufferAllocator(; sizehint = 1 << 16)
+            @planar allocator = buffer C[i; j] := A[i; k l] * τ[k l; m n] * B[m n; j]
+            @test isempty(buffer)
+            @test C ≈ Cref
+
+            # the result must not live in the buffer: the next block hands out the same
+            # memory again, and `C` has to survive that
+            @planar allocator = buffer C2[i; j] := A[i; k l] * τ[k l; m n] * B[m n; j]
+            @test isempty(buffer)
+            @test C ≈ Cref
+            @test C2 ≈ Cref
+        end
+    end
+
     @testset "contractcheck" begin
         V = ℂ^2
         A = rand(T, V ⊗ V ← V)