From ace31e79342276c11aca5117f13db10166b1eb5f Mon Sep 17 00:00:00 2001
From: Sassoun Derderian <sassound@adobe.com>
Date: Fri, 13 Sep 2024 11:02:08 +0400
Subject: [PATCH] fix: row grouping algorithm

fix: typings for RowGroupingMapperResult
---
 .../core/src/data-editor/row-grouping-api.ts  |  6 +-
 packages/core/src/data-editor/row-grouping.ts | 15 ++--
 packages/core/test/row-grouping-api.test.ts   | 77 +++++++++++++++++++
 3 files changed, 88 insertions(+), 10 deletions(-)
 create mode 100644 packages/core/test/row-grouping-api.test.ts

diff --git a/packages/core/src/data-editor/row-grouping-api.ts b/packages/core/src/data-editor/row-grouping-api.ts
index 23eeb8803..988de9013 100644
--- a/packages/core/src/data-editor/row-grouping-api.ts
+++ b/packages/core/src/data-editor/row-grouping-api.ts
@@ -1,13 +1,13 @@
 import React from "react";
 import type { Item } from "../internal/data-grid/data-grid-types.js";
-import { flattenRowGroups, mapRowIndexToPath, type RowGroup, type RowGroupingOptions } from "./row-grouping.js";
+import { flattenRowGroups, mapRowIndexToPath, type RowGroup, type RowGroupingOptions, type MapResult } from "./row-grouping.js";
 
-export type RowGroupingMapperResult<T> = {
+export interface RowGroupingMapperResult<T> extends Omit<MapResult, 'originalIndex'> {
     path: readonly number[];
     originalIndex: T;
     isGroupHeader: boolean;
     groupRows: number;
-};
+}
 
 export type RowGroupingMapper = {
     (itemOrRow: number): RowGroupingMapperResult<number>;
diff --git a/packages/core/src/data-editor/row-grouping.ts b/packages/core/src/data-editor/row-grouping.ts
index 4f90c86ec..4cda4782b 100644
--- a/packages/core/src/data-editor/row-grouping.ts
+++ b/packages/core/src/data-editor/row-grouping.ts
@@ -167,7 +167,7 @@ export function flattenRowGroups(rowGrouping: RowGroupingOptions, rows: number):
         });
 }
 
-interface MapResult {
+export interface MapResult {
     readonly path: readonly number[];
     readonly isGroupHeader: boolean;
     readonly originalIndex: number;
@@ -199,18 +199,19 @@ export function mapRowIndexToPath(row: number, flattenedRowGroups?: readonly Fla
                 contentIndex: -1,
                 groupRows: group.rows,
             };
-        toGo--;
         if (!group.isCollapsed) {
-            if (toGo < group.rows)
+            if (toGo <= group.rows)
                 return {
-                    path: [...group.path, toGo],
+                    path: [...group.path, toGo - 1],
                     originalIndex: group.headerIndex + toGo,
                     isGroupHeader: false,
-                    groupIndex: toGo,
-                    contentIndex: group.contentIndex + toGo,
+                    groupIndex: toGo - 1,
+                    contentIndex: group.contentIndex + toGo - 1,
                     groupRows: group.rows,
                 };
-            toGo -= group.rows;
+            toGo = toGo - group.rows - 1;
+        } else {
+            toGo--;
         }
     }
     // this shouldn't happen
diff --git a/packages/core/test/row-grouping-api.test.ts b/packages/core/test/row-grouping-api.test.ts
new file mode 100644
index 000000000..4d8d8ff03
--- /dev/null
+++ b/packages/core/test/row-grouping-api.test.ts
@@ -0,0 +1,77 @@
+import { describe, it, expect, afterEach } from "vitest";
+import { renderHook, cleanup } from "@testing-library/react-hooks";
+import { useRowGrouping } from "../src/data-editor/row-grouping-api.js";
+import type { RowGroupingOptions } from "../src/data-editor/row-grouping.js";
+
+describe("useRowGrouping - mapper", () => {
+    afterEach(async () => {
+        await cleanup();
+    });
+
+    it("calculates rows correctly with no row grouping options", () => {
+        const totalRows = 5;
+        const { result } = renderHook(() => useRowGrouping(undefined, totalRows));
+
+        expect(result.current.mapper(0)).toEqual({
+            path: [ 0 ],
+            originalIndex: 0,
+            isGroupHeader: false,
+            groupIndex: 0,
+            contentIndex: 0,
+            groupRows: -1
+        });
+        expect(result.current.mapper(1)).toEqual({
+            path: [ 1 ],
+            originalIndex: 1,
+            isGroupHeader: false,
+            groupIndex: 1,
+            contentIndex: 1,
+            groupRows: -1
+        });
+    });
+
+    it("calculates rows correctly with uncollapsed groups - rowIndex is a group header", () => {
+        const rowGroupingOptions: RowGroupingOptions = {
+            groups: [
+                { headerIndex: 0, isCollapsed: false },
+                { headerIndex: 2, isCollapsed: false },
+            ],
+            height: 30,
+        };
+        const totalRows = 5;
+
+        const { result } = renderHook(() => useRowGrouping(rowGroupingOptions, totalRows));
+
+        expect(result.current.mapper(0)).toEqual({
+            path: [ 0, -1 ],
+            originalIndex: 0,
+            isGroupHeader: true,
+            groupIndex: -1,
+            contentIndex: -1,
+            groupRows: 1
+        });
+    });
+
+    it("calculates rows correctly with uncollapsed groups - rowIndex is a regular row", () => {
+        const rowGroupingOptions: RowGroupingOptions = {
+            groups: [
+                { headerIndex: 0, isCollapsed: false },
+                { headerIndex: 2, isCollapsed: false },
+            ],
+            height: 30,
+        };
+        const totalRows = 5;
+
+        const { result } = renderHook(() => useRowGrouping(rowGroupingOptions, totalRows));
+
+        expect(result.current.mapper(1)).toEqual({
+            path: [ 0, 0 ],
+            originalIndex: 1,
+            isGroupHeader: false,
+            groupIndex: 0,
+            contentIndex: 0,
+            groupRows: 1
+        });
+    });
+
+});