From 951951776a2c7178a52163d0184df8728b50f421 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sat, 2 Sep 2023 12:48:44 +0530
Subject: [PATCH] Ensure group_count is never zero in any ImageRef

Now group_count means the number of refs pointing to the same image from
that ref onwards. This is needed because we can index the list of refs
at any point when drawing not just at the start of a group.

Fixes #6594
---
 kitty/graphics.c        | 12 +++++++-----
 kitty_tests/graphics.py | 19 ++++++++++++++++++-
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/kitty/graphics.c b/kitty/graphics.c
index 30f10ac99..62c78e02f 100644
--- a/kitty/graphics.c
+++ b/kitty/graphics.c
@@ -983,12 +983,14 @@ grman_update_layers(GraphicsManager *self, unsigned int scrolled_by, float scree
     // Calculate the group counts
     i = 0;
     while (i < self->count) {
-        id_type image_id = self->render_data[i].image_id, start = i;
-        if (start == self->count - 1) i = self->count;
-        else {
-            while (i < self->count - 1 && self->render_data[++i].image_id == image_id) {}
+        id_type num_identical = 1, image_id = self->render_data[i].image_id, start = i;
+        while (++i < self->count) {
+            if (self->render_data[i].image_id != image_id) break;
+            num_identical++;
+        }
+        while (num_identical > 0) {
+            self->render_data[start++].group_count = num_identical--;
         }
-        self->render_data[start].group_count = i - start;
     }
     return true;
 }
diff --git a/kitty_tests/graphics.py b/kitty_tests/graphics.py
index a045c963c..730929913 100644
--- a/kitty_tests/graphics.py
+++ b/kitty_tests/graphics.py
@@ -516,13 +516,30 @@ class TestGraphics(BaseTest):
         rect_eq(l2[0]['dest_rect'], left, top, -1 + (1 + s.columns) * dx, top - dy * 5 / ch)
         rect_eq(l2[1]['src_rect'], 0, 0, 1, 1)
         rect_eq(l2[1]['dest_rect'], -1, 1, -1 + dx, 1 - dy)
-        self.ae(l2[0]['group_count'], 1), self.ae(l2[1]['group_count'], 1)
+        self.ae(l2[0]['group_count'], 2)
+        self.ae(l2[1]['group_count'], 1)
         self.ae(s.cursor.x, 0), self.ae(s.cursor.y, 1)
         self.ae(put_image(s, 10, 20, cursor_movement=1)[1], 'OK')
         self.ae(s.cursor.x, 0), self.ae(s.cursor.y, 1)
         s.reset()
         self.assertEqual(s.grman.disk_cache.total_size, 0)
 
+    def test_image_layer_grouping(self):
+        cw, ch = 10, 20
+        s, dx, dy, put_image, put_ref, layers, rect_eq = put_helpers(self, cw, ch)
+
+        def group_counts():
+            return tuple(x['group_count'] for x in layers(s))
+
+        self.ae(put_image(s, 10, 20, id=1)[1], 'OK')
+        self.ae(group_counts(), (1,))
+        put_ref(s, id=1, num_cols=2, num_lines=1, placement_id=2)
+        put_ref(s, id=1, num_cols=2, num_lines=1, placement_id=3, z=-2)
+        put_ref(s, id=1, num_cols=2, num_lines=1, placement_id=4, z=-2)
+        self.ae(group_counts(), (4, 3, 2, 1))
+        self.ae(put_image(s, 8, 16, id=2, z=-1)[1], 'OK')
+        self.ae(group_counts(), (2, 1, 1, 2, 1))
+
     def test_unicode_placeholders(self):
         # This test tests basic image placement using using unicode placeholders
         cw, ch = 10, 20