From fa4f909d337d67aee66d8d822ab00dad724bbcd1 Mon Sep 17 00:00:00 2001
From: lijinke666 <a1231236677287@163.com>
Date: Fri, 22 Mar 2024 15:07:37 +0800
Subject: [PATCH 1/3] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E8=A7=92=E5=A4=B4?=
 =?UTF-8?q?=E5=92=8C=E8=A1=8C=E5=A4=B4=E6=8A=98=E5=8F=A0=E5=B1=95=E5=BC=80?=
 =?UTF-8?q?=20icon=20=E7=9A=84=E7=8A=B6=E6=80=81=E6=9C=AA=E5=90=8C?=
 =?UTF-8?q?=E6=AD=A5=E4=BB=A5=E5=8F=8A=E5=B1=95=E5=BC=80=E5=BC=82=E5=B8=B8?=
 =?UTF-8?q?=E7=9A=84=E9=97=AE=E9=A2=98=20close=20#2607?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../spread-sheet-collapse-spec.ts.snap        |  46 ++++++
 .../spreadsheet/spread-sheet-collapse-spec.ts | 138 ++++++++++++++++--
 packages/s2-core/src/cell/corner-cell.ts      |  14 +-
 .../layout/build-row-custom-tree-hierarchy.ts |   2 +-
 .../s2-core/src/sheet-type/pivot-sheet.ts     |   1 -
 .../playground/components/CustomTree.tsx      |   5 +
 packages/s2-react/playground/config.tsx       |  12 +-
 s2-site/docs/manual/migration-v2.zh.md        |   6 +-
 8 files changed, 202 insertions(+), 22 deletions(-)
 create mode 100644 packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap

diff --git a/packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap b/packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap
new file mode 100644
index 0000000000..84e96f835d
--- /dev/null
+++ b/packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap
@@ -0,0 +1,46 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`SpreadSheet Collapse/Expand Tests Tree Mode should only expand current group row nodes 1`] = `
+Array [
+  "root[&]浙江省",
+  "root[&]浙江省[&]杭州市",
+  "root[&]浙江省[&]绍兴市",
+  "root[&]浙江省[&]宁波市",
+  "root[&]浙江省[&]舟山市",
+  "root[&]四川省",
+]
+`;
+
+exports[`SpreadSheet Collapse/Expand Tests Tree Mode should only expand current group row nodes for custom tree 1`] = `
+Array [
+  "root[&]a-1",
+  "root[&]a-1[&]a-1-1",
+  "root[&]a-1[&]a-1-2",
+  "root[&]a-2",
+  "root[&]a-3",
+]
+`;
+
+exports[`SpreadSheet Collapse/Expand Tests Tree Mode should sync corner cell collapse all icon status 1`] = `
+Array [
+  "root[&]浙江省",
+  "root[&]四川省",
+]
+`;
+
+exports[`SpreadSheet Collapse/Expand Tests should only expand current group row nodes for custom tree 1`] = `
+Array [
+  "root[&]a-1",
+  "root[&]a-1[&]a-1-1",
+  "root[&]a-1[&]a-1-2",
+  "root[&]a-2",
+  "root[&]a-3",
+]
+`;
+
+exports[`SpreadSheet Collapse/Expand Tests should sync corner cell collapse all icon status 1`] = `
+Array [
+  "root[&]浙江省",
+  "root[&]四川省",
+]
+`;
diff --git a/packages/s2-core/__tests__/spreadsheet/spread-sheet-collapse-spec.ts b/packages/s2-core/__tests__/spreadsheet/spread-sheet-collapse-spec.ts
index 5f93873aa9..a791828dc0 100644
--- a/packages/s2-core/__tests__/spreadsheet/spread-sheet-collapse-spec.ts
+++ b/packages/s2-core/__tests__/spreadsheet/spread-sheet-collapse-spec.ts
@@ -1,6 +1,13 @@
 import * as mockDataConfig from 'tests/data/simple-data.json';
-import { getContainer } from 'tests/util/helpers';
-import type { RowCellCollapsedParams } from '../../src/common/interface';
+import { createPivotSheet, getContainer, sleep } from 'tests/util/helpers';
+import {
+  CornerNodeType,
+  type RowCellCollapsedParams,
+  type S2DataConfig,
+  type S2Options,
+} from '../../src/common/interface';
+import { customRowGridFields } from '../data/custom-grid-fields';
+import { CustomGridData } from '../data/data-custom-grid';
 import { S2Event } from './../../src/common/constant/events/basic';
 import { PivotSheet, SpreadSheet } from '@/sheet-type';
 import type { Node } from '@/facet/layout/node';
@@ -9,9 +16,30 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
   let container: HTMLElement;
   let s2: SpreadSheet;
 
+  const s2Options: S2Options = {
+    width: 600,
+    height: 400,
+    hierarchyType: 'tree',
+    style: {
+      rowCell: {
+        expandDepth: undefined,
+      },
+    },
+  };
+
   const mapNodes = (spreadsheet: SpreadSheet) =>
     spreadsheet.facet.getRowLeafNodes().map((node) => node.id);
 
+  const expectCornerIconName = (spreadsheet: SpreadSheet, iconName: string) => {
+    const cornerCell = spreadsheet.facet
+      .getCornerCells()
+      .find((cell) => cell.getMeta().cornerType === CornerNodeType.Row)!;
+
+    const cornerIcon = cornerCell.getTreeIcon();
+
+    expect(cornerIcon?.name).toEqual(iconName);
+  };
+
   beforeEach(async () => {
     container = getContainer();
     s2 = new PivotSheet(
@@ -25,22 +53,13 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
           valueInCols: true,
         },
       },
-      {
-        width: 600,
-        height: 200,
-        hierarchyType: 'tree',
-        style: {
-          rowCell: {
-            expandDepth: undefined,
-          },
-        },
-      },
+      s2Options,
     );
     await s2.render();
   });
 
   afterEach(() => {
-    s2.destroy();
+    // s2.destroy();
   });
 
   describe('Tree Mode', () => {
@@ -101,6 +120,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
                 "root[&]浙江",
               ]
           `);
+      expectCornerIconName(s2, 'Plus');
     });
 
     test('should collapse by field', async () => {
@@ -121,6 +141,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
                 "root[&]浙江",
               ]
           `);
+      expectCornerIconName(s2, 'Plus');
     });
 
     test('should collapse by field id', async () => {
@@ -144,6 +165,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
                 "root[&]浙江[&]杭州[&]笔",
               ]
           `);
+      expectCornerIconName(s2, 'Plus');
     });
 
     test('should collapse use collapseFields first', async () => {
@@ -167,6 +189,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
                 "root[&]浙江[&]杭州",
               ]
           `);
+      expectCornerIconName(s2, 'Plus');
     });
 
     test('should collapse use expandDepth first', async () => {
@@ -189,6 +212,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
                 "root[&]浙江[&]杭州[&]笔",
               ]
           `);
+      expectCornerIconName(s2, 'Minus');
     });
 
     test('should collapse use collapseFields first when contain collapseAll and expandDepth config', async () => {
@@ -213,6 +237,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
                 "root[&]浙江[&]杭州",
               ]
           `);
+      expectCornerIconName(s2, 'Plus');
     });
 
     test('should collapse use collapseFields by node id first', async () => {
@@ -233,6 +258,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
           "root[&]浙江",
         ]
       `);
+      expectCornerIconName(s2, 'Plus');
     });
 
     test('should collapse all nodes if collapseAll is true and collapseFields is undefined', async () => {
@@ -251,6 +277,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
           "root[&]浙江",
         ]
       `);
+      expectCornerIconName(s2, 'Plus');
     });
 
     test('should emit collapse event', () => {
@@ -275,6 +302,7 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
       s2.emit(S2Event.ROW_CELL_COLLAPSED__PRIVATE, treeRowType);
 
       expect(onCollapsed).toHaveBeenCalledWith(params);
+      expectCornerIconName(s2, 'Minus');
     });
 
     test('should emit collapse all event', () => {
@@ -285,6 +313,90 @@ describe('SpreadSheet Collapse/Expand Tests', () => {
       s2.emit(S2Event.ROW_CELL_ALL_COLLAPSED__PRIVATE, false);
 
       expect(onCollapsed).toHaveBeenCalledWith(true);
+      expectCornerIconName(s2, 'Minus');
+    });
+
+    // https://github.com/antvis/S2/issues/2607
+    test('should only expand current group row nodes', async () => {
+      s2 = createPivotSheet(
+        {
+          ...s2Options,
+          style: {
+            rowCell: {
+              collapseAll: true,
+            },
+          },
+        },
+        { useSimpleData: false },
+      );
+
+      await s2.render();
+
+      const rowNodes = s2.facet.getRowNodes(0);
+
+      s2.emit(S2Event.ROW_CELL_COLLAPSED__PRIVATE, {
+        isCollapsed: false,
+        node: rowNodes[0],
+      });
+
+      await sleep(500);
+
+      expectCornerIconName(s2, 'Minus');
+      expect(s2.facet.getRowNodes().map(({ id }) => id)).toMatchSnapshot();
+    });
+
+    test('should sync corner cell collapse all icon status', async () => {
+      s2 = createPivotSheet(
+        {
+          ...s2Options,
+          style: {
+            rowCell: {
+              collapseAll: false,
+              collapseFields: {
+                'root[&]浙江省': true,
+                'root[&]四川省': true,
+              },
+            },
+          },
+        },
+        { useSimpleData: false },
+      );
+
+      await s2.render();
+
+      expectCornerIconName(s2, 'Plus');
+      expect(s2.facet.getRowNodes().map(({ id }) => id)).toMatchSnapshot();
+    });
+
+    test('should only expand current group row nodes for custom tree', async () => {
+      const customRowDataCfg: S2DataConfig = {
+        data: CustomGridData,
+        fields: customRowGridFields,
+      };
+
+      s2 = createPivotSheet({
+        ...s2Options,
+        style: {
+          rowCell: {
+            collapseAll: true,
+          },
+        },
+      });
+
+      s2.setDataCfg(customRowDataCfg);
+      await s2.render();
+
+      const rowNodes = s2.facet.getRowNodes(0);
+
+      s2.emit(S2Event.ROW_CELL_COLLAPSED__PRIVATE, {
+        isCollapsed: false,
+        node: rowNodes[0],
+      });
+
+      await sleep(500);
+
+      expectCornerIconName(s2, 'Minus');
+      expect(s2.facet.getRowNodes().map(({ id }) => id)).toMatchSnapshot();
     });
   });
 });
diff --git a/packages/s2-core/src/cell/corner-cell.ts b/packages/s2-core/src/cell/corner-cell.ts
index 3b45b1eacd..28b9b9a12d 100644
--- a/packages/s2-core/src/cell/corner-cell.ts
+++ b/packages/s2-core/src/cell/corner-cell.ts
@@ -1,5 +1,5 @@
 import type { PointLike } from '@antv/g';
-import { last } from 'lodash';
+import { last, values } from 'lodash';
 import {
   CellType,
   KEY_GROUP_CORNER_RESIZE_AREA,
@@ -61,10 +61,16 @@ export class CornerCell extends HeaderCell<CornerHeaderConfig> {
       return;
     }
 
-    const { collapseAll } = this.spreadsheet.options.style?.rowCell!;
+    const { collapseFields = {} } = this.spreadsheet.options.style?.rowCell!;
     const { size = 0 } = this.getStyle()!.icon!;
     const { fill } = this.getTextStyle();
     const area = this.getBBoxByType(CellClipBox.CONTENT_BOX);
+    const rootRowNodes = this.spreadsheet.facet.getRowNodes(0);
+    // 任意一级节点的 icon 展开/收起, 同步更新角头 icon 的状态
+    const isAllCollapsed =
+      values(collapseFields!).filter(Boolean).length === rootRowNodes.length ||
+      rootRowNodes.every((node) => node.isCollapsed);
+    const isCollapsed = isAllCollapsed;
 
     this.treeIcon = renderTreeIcon({
       group: this,
@@ -75,12 +81,12 @@ export class CornerCell extends HeaderCell<CornerHeaderConfig> {
         height: size,
         fill,
       },
-      isCollapsed: collapseAll,
+      isCollapsed,
       onClick: () => {
         this.spreadsheet.facet.resetScrollY();
         this.spreadsheet.emit(
           S2Event.ROW_CELL_ALL_COLLAPSED__PRIVATE,
-          collapseAll!,
+          isCollapsed,
         );
       },
     });
diff --git a/packages/s2-core/src/facet/layout/build-row-custom-tree-hierarchy.ts b/packages/s2-core/src/facet/layout/build-row-custom-tree-hierarchy.ts
index 7264f8d542..2343ef096d 100644
--- a/packages/s2-core/src/facet/layout/build-row-custom-tree-hierarchy.ts
+++ b/packages/s2-core/src/facet/layout/build-row-custom-tree-hierarchy.ts
@@ -40,7 +40,7 @@ export const buildCustomTreeHierarchy = (params: CustomTreeHeaderParams) => {
 
     const defaultCollapsed = collapsed ?? false;
     const isDefaultCollapsed =
-      collapseFields?.[nodeId] || collapseFields?.[field];
+      collapseFields?.[nodeId] ?? collapseFields?.[field];
     const isCollapsed = isDefaultCollapsed ?? (collapseAll || defaultCollapsed);
 
     // TODO: 平铺模式支持 折叠/展开
diff --git a/packages/s2-core/src/sheet-type/pivot-sheet.ts b/packages/s2-core/src/sheet-type/pivot-sheet.ts
index 7ae5725d23..211278cc30 100644
--- a/packages/s2-core/src/sheet-type/pivot-sheet.ts
+++ b/packages/s2-core/src/sheet-type/pivot-sheet.ts
@@ -114,7 +114,6 @@ export class PivotSheet extends SpreadSheet {
     this.setOptions({
       style: {
         rowCell: {
-          collapseAll: false,
           collapseFields,
         },
       },
diff --git a/packages/s2-react/playground/components/CustomTree.tsx b/packages/s2-react/playground/components/CustomTree.tsx
index 66bf8732c0..be3dcf7ef9 100644
--- a/packages/s2-react/playground/components/CustomTree.tsx
+++ b/packages/s2-react/playground/components/CustomTree.tsx
@@ -41,6 +41,11 @@ export const CustomTreeOptions: SheetComponentOptions = {
       withFormat: true,
     },
   },
+  style: {
+    rowCell: {
+      collapseAll: true,
+    },
+  },
   // cornerText: '指标',
 };
 
diff --git a/packages/s2-react/playground/config.tsx b/packages/s2-react/playground/config.tsx
index 7bae02f71e..8e8ec5b39f 100644
--- a/packages/s2-react/playground/config.tsx
+++ b/packages/s2-react/playground/config.tsx
@@ -315,7 +315,7 @@ export const s2Options: SheetComponentOptions = {
   debug: true,
   width: 800,
   height: 600,
-  hierarchyType: 'grid',
+  hierarchyType: 'tree',
   seriesNumber: {
     enable: true,
   },
@@ -384,7 +384,15 @@ export const s2Options: SheetComponentOptions = {
   //   ],
   // ],
   tooltip: S2TooltipOptions,
-  style: {},
+  style: {
+    rowCell: {
+      collapseAll: true,
+      // collapseFields: {
+      //   'root[&]浙江省': true,
+      //   'root[&]四川省': true,
+      // },
+    },
+  },
 };
 
 export const sliderOptions: SliderSingleProps = {
diff --git a/s2-site/docs/manual/migration-v2.zh.md b/s2-site/docs/manual/migration-v2.zh.md
index 227b26586a..c5a849faee 100644
--- a/s2-site/docs/manual/migration-v2.zh.md
+++ b/s2-site/docs/manual/migration-v2.zh.md
@@ -357,7 +357,11 @@ const s2DataConfig = {
 
 具体请查看 [自定义行列头分组](/manual/advanced/custom/custom-header) 相关文档。
 
-#### 行列冻结配置调整
+#### 树状结构 icon 折叠展开状态同步
+
+现在行头节点的 icon 展开/收起,会同步更新角头 icon(全部展开/收起)的状态。
+
+#### 行列冻结配置
 
 透视表和明细表的行列冻结配置统一收拢到 `frozen`.
 

From 9eab6745c2187ec0e687f09f0b1ddf1f11202fcb Mon Sep 17 00:00:00 2001
From: lijinke666 <a1231236677287@163.com>
Date: Fri, 22 Mar 2024 15:22:15 +0800
Subject: [PATCH 2/3] chore: fix

---
 packages/s2-react/playground/config.tsx | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/packages/s2-react/playground/config.tsx b/packages/s2-react/playground/config.tsx
index 8e8ec5b39f..f9f4c3401d 100644
--- a/packages/s2-react/playground/config.tsx
+++ b/packages/s2-react/playground/config.tsx
@@ -384,15 +384,7 @@ export const s2Options: SheetComponentOptions = {
   //   ],
   // ],
   tooltip: S2TooltipOptions,
-  style: {
-    rowCell: {
-      collapseAll: true,
-      // collapseFields: {
-      //   'root[&]浙江省': true,
-      //   'root[&]四川省': true,
-      // },
-    },
-  },
+  style: {},
 };
 
 export const sliderOptions: SliderSingleProps = {

From af0b48fef58e957fcf77cfcee90ff8297be4c308 Mon Sep 17 00:00:00 2001
From: lijinke666 <a1231236677287@163.com>
Date: Fri, 22 Mar 2024 16:09:34 +0800
Subject: [PATCH 3/3] =?UTF-8?q?test:=20=E5=BF=AB=E7=85=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../spread-sheet-collapse-spec.ts.snap          | 17 -----------------
 1 file changed, 17 deletions(-)

diff --git a/packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap b/packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap
index 84e96f835d..f9c60166c0 100644
--- a/packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap
+++ b/packages/s2-core/__tests__/spreadsheet/__snapshots__/spread-sheet-collapse-spec.ts.snap
@@ -27,20 +27,3 @@ Array [
   "root[&]四川省",
 ]
 `;
-
-exports[`SpreadSheet Collapse/Expand Tests should only expand current group row nodes for custom tree 1`] = `
-Array [
-  "root[&]a-1",
-  "root[&]a-1[&]a-1-1",
-  "root[&]a-1[&]a-1-2",
-  "root[&]a-2",
-  "root[&]a-3",
-]
-`;
-
-exports[`SpreadSheet Collapse/Expand Tests should sync corner cell collapse all icon status 1`] = `
-Array [
-  "root[&]浙江省",
-  "root[&]四川省",
-]
-`;