From 0d6dddd072ffe3fa6a015d32a0830096b1a0e59d Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Mon, 19 Jan 2026 23:41:31 +0000
Subject: [PATCH] Fix memory leak in ShapeArray.Remove

- Set the element in the underlying array to nil after swapping/removing it from the slice.
- Add regression test TestShapeArray_Remove_MemoryLeak.
---
 spacepartition/spacemap.go      |  1 +
 spacepartition/spacemap_test.go | 34 +++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/spacepartition/spacemap.go b/spacepartition/spacemap.go
index 142c5ab..b2c9784 100644
--- a/spacepartition/spacemap.go
+++ b/spacepartition/spacemap.go
@@ -263,6 +263,7 @@ func (a ShapeArray) Remove(shape shared.Shape) ([]shared.Shape, int) {
 		for a[i] == shape && len(a)-shrink > i {
 			shrink++
 			a[i] = a[len(a)-shrink]
+			a[len(a)-shrink] = nil
 		}
 	}
 	return a[:len(a)-shrink], shrink
diff --git a/spacepartition/spacemap_test.go b/spacepartition/spacemap_test.go
index fd37f0b..69c07e1 100644
--- a/spacepartition/spacemap_test.go
+++ b/spacepartition/spacemap_test.go
@@ -343,3 +343,37 @@ func TestZIndex(t *testing.T) {
 		})
 	}
 }
+
+func TestShapeArray_Remove_MemoryLeak(t *testing.T) {
+	s1 := shared.NewRectangle(0, 0, 10, 10)
+	s2 := shared.NewRectangle(10, 10, 20, 20)
+	s3 := shared.NewRectangle(20, 20, 30, 30)
+
+	// Initialize ShapeArray with capacity = length
+	arr := make(ShapeArray, 3)
+	arr[0] = s1
+	arr[1] = s2
+	arr[2] = s3
+
+	// Remove s2 (middle element)
+	// Current implementation: swap s2 with s3, shrink by 1.
+	// Expected result slice: [s1, s3]
+	// Expected underlying array: [s1, s3, s3] (Leak)
+	// Desired underlying array: [s1, s3, nil]
+
+	newArr, _ := arr.Remove(s2)
+
+	if len(newArr) != 2 {
+		t.Fatalf("Expected length 2, got %d", len(newArr))
+	}
+
+	// Access the underlying array up to capacity
+	underlying := newArr[:cap(newArr)]
+
+	// The element at index 2 (the one that was "removed" from the slice)
+	leaked := underlying[2]
+
+	if leaked != nil {
+		t.Errorf("Memory leak detected: element at index 2 is not nil: %v", leaked)
+	}
+}