From 9ed3c89594677df8d04c53ad3e80b07414815ead Mon Sep 17 00:00:00 2001
From: Hajime Hoshi <hajimehoshi@gmail.com>
Date: Sun, 8 Oct 2023 19:38:09 +0900
Subject: [PATCH] internal/graphicscommand: remove an image from
 imageWithBuffers whenever possible

---
 image_test.go                     | 19 +++++++++++++++++++
 internal/graphicscommand/image.go | 21 ++++++++++++++++-----
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/image_test.go b/image_test.go
index 479b032ab..52ab46539 100644
--- a/image_test.go
+++ b/image_test.go
@@ -4231,3 +4231,22 @@ func Fragment(dstPos vec4, srcPos vec2, color vec4) vec4 {
 		t.Errorf("got: (%0.2f, %0.2f), want: (0, 0)", x, y)
 	}
 }
+
+func TestImageWritePixelAndDispose(t *testing.T) {
+	const (
+		w = 16
+		h = 16
+	)
+	img := ebiten.NewImage(w, h)
+	pix := make([]byte, 4*w*h)
+	for i := range pix {
+		pix[i] = 0xff
+	}
+	img.WritePixels(pix)
+	img.Dispose()
+
+	// Confirm that any pixel information is invalidated after Dispose is called.
+	if got, want := img.At(0, 0), (color.RGBA{}); got != want {
+		t.Errorf("got: %v, want: %v", got, want)
+	}
+}
diff --git a/internal/graphicscommand/image.go b/internal/graphicscommand/image.go
index 4ef3f330f..97b1108ff 100644
--- a/internal/graphicscommand/image.go
+++ b/internal/graphicscommand/image.go
@@ -55,21 +55,28 @@ func genNextID() int {
 }
 
 // imagesWithBuffers is the set of an image with buffers.
-var imagesWithBuffers []*Image
+var imagesWithBuffers = map[*Image]struct{}{}
 
 // addImageWithBuffer adds an image to the list of images with unflushed buffers.
 func addImageWithBuffer(img *Image) {
-	imagesWithBuffers = append(imagesWithBuffers, img)
+	imagesWithBuffers[img] = struct{}{}
+}
+
+// removeImageWithBuffer removes an image from the list of images with unflushed buffers.
+func removeImageWithBuffer(img *Image) {
+	delete(imagesWithBuffers, img)
 }
 
 // flushImageBuffers flushes all the image buffers and send to the command queue.
 // flushImageBuffers should be called before flushing commands.
 func flushImageBuffers() {
-	for i, img := range imagesWithBuffers {
+	for img := range imagesWithBuffers {
 		img.flushBufferedWritePixels()
-		imagesWithBuffers[i] = nil
 	}
-	imagesWithBuffers = imagesWithBuffers[:0]
+
+	if len(imagesWithBuffers) != 0 {
+		panic("graphicscommand: len(imagesWithBuffers) must be empty after flushing")
+	}
 }
 
 // NewImage returns a new image.
@@ -103,6 +110,8 @@ func (i *Image) flushBufferedWritePixels() {
 	theCommandQueueManager.enqueueCommand(c)
 
 	i.bufferedWritePixelsArgs = nil
+
+	removeImageWithBuffer(i)
 }
 
 func (i *Image) Dispose() {
@@ -111,6 +120,8 @@ func (i *Image) Dispose() {
 		target: i,
 	}
 	theCommandQueueManager.enqueueCommand(c)
+
+	removeImageWithBuffer(i)
 }
 
 func (i *Image) InternalSize() (int, int) {