From 2b00a31d4d0f2aa183dbe7e9d8558ea619252a2d Mon Sep 17 00:00:00 2001
From: Ben Wittenberg <benjaminwittenberg@gmail.com>
Date: Tue, 4 Jun 2024 19:13:26 +0000
Subject: [PATCH] chore: spike - improve performance with early return

---
 packages/react/src/DataTable/DataTable.tsx    |  6 +-
 .../DataTable/__tests__/DataTable.test.tsx    | 63 ++++++++++++++++++-
 packages/react/src/DataTable/useTable.ts      |  9 ++-
 3 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/packages/react/src/DataTable/DataTable.tsx b/packages/react/src/DataTable/DataTable.tsx
index e2d59e171a96..7f89c545fb91 100644
--- a/packages/react/src/DataTable/DataTable.tsx
+++ b/packages/react/src/DataTable/DataTable.tsx
@@ -61,13 +61,17 @@ function DataTable<Data extends UniqueRow>({
   initialSortColumn,
   initialSortDirection,
 }: DataTableProps<Data>) {
-  const {headers, rows, actions, gridTemplateColumns} = useTable({
+  const {headers, rows, actions, gridTemplateColumns, willReactThrowOutRenderResult} = useTable({
     data,
     columns,
     initialSortColumn,
     initialSortDirection,
   })
 
+  if (willReactThrowOutRenderResult) {
+    return null
+  }
+
   return (
     <Table
       aria-labelledby={labelledby}
diff --git a/packages/react/src/DataTable/__tests__/DataTable.test.tsx b/packages/react/src/DataTable/__tests__/DataTable.test.tsx
index d95f0f09456b..e85f7e4f60f2 100644
--- a/packages/react/src/DataTable/__tests__/DataTable.test.tsx
+++ b/packages/react/src/DataTable/__tests__/DataTable.test.tsx
@@ -1,6 +1,6 @@
 import userEvent from '@testing-library/user-event'
 import {render, screen, queryAllByRole} from '@testing-library/react'
-import React from 'react'
+import React, {Profiler, type ComponentProps} from 'react'
 import {DataTable, Table} from '../../DataTable'
 import type {Column} from '../column'
 import {createColumnHelper} from '../column'
@@ -857,6 +857,41 @@ describe('DataTable', () => {
       expect(customSortFn).toHaveBeenCalled()
       expect(getCellContentByColumnHeaderName('Value')).toEqual(['3', '2', '1'])
     })
+
+    it('should sort again when data changes', async () => {
+      type Data = {id: string; value: string}
+      const columns: ComponentProps<typeof DataTable<Data>>['columns'] = [
+        {
+          header: 'Value',
+          field: 'value',
+          sortBy: true,
+        },
+      ]
+      const data: ComponentProps<typeof DataTable<Data>>['data'] = [
+        {
+          id: '3',
+          value: '3',
+        },
+        {
+          id: '2',
+          value: '2',
+        },
+      ]
+      const {rerender} = render(
+        <DataTable data={data} columns={columns} initialSortColumn="value" initialSortDirection="ASC" />,
+      )
+
+      rerender(
+        <DataTable
+          data={data.concat({id: '1', value: '1'})}
+          columns={columns}
+          initialSortColumn="value"
+          initialSortDirection="ASC"
+        />,
+      )
+
+      expect(getCellContentByColumnHeaderName('Value')).toEqual(['1', '2', '3'])
+    })
   })
 
   describe('column widths', () => {
@@ -1036,4 +1071,30 @@ describe('DataTable', () => {
       })
     })
   })
+
+  describe('performance tests', () => {
+    it('should not render twice on initial render', () => {
+      type Data = {id: string; value: string}
+      const columns: ComponentProps<typeof DataTable<Data>>['columns'] = [
+        {
+          header: 'Value',
+          field: 'value',
+        },
+      ]
+      const data: ComponentProps<typeof DataTable<Data>>['data'] = [
+        {
+          id: '1',
+          value: 'one',
+        },
+      ]
+      let renderCount = 0
+      render(
+        <Profiler id="DataTable" onRender={() => renderCount++}>
+          <DataTable data={data} columns={columns} />
+        </Profiler>,
+      )
+
+      expect(renderCount).toBe(1)
+    })
+  })
 })
diff --git a/packages/react/src/DataTable/useTable.ts b/packages/react/src/DataTable/useTable.ts
index d046f736e88e..a85d73897d26 100644
--- a/packages/react/src/DataTable/useTable.ts
+++ b/packages/react/src/DataTable/useTable.ts
@@ -47,7 +47,7 @@ export function useTable<Data extends UniqueRow>({
   data,
   initialSortColumn,
   initialSortDirection,
-}: TableConfig<Data>): Table<Data> {
+}: TableConfig<Data>): Table<Data> & {willReactThrowOutRenderResult: boolean} {
   const [rowOrder, setRowOrder] = useState(data)
   const [prevData, setPrevData] = useState<typeof data | null>(null)
   const [prevColumns, setPrevColumns] = useState(columns)
@@ -58,7 +58,8 @@ export function useTable<Data extends UniqueRow>({
 
   // Reset the `sortByColumn` state if the columns change and that column is no
   // longer provided
-  if (columns !== prevColumns) {
+  const haveColumnsChanged = columns !== prevColumns
+  if (haveColumnsChanged) {
     setPrevColumns(columns)
     if (sortByColumn) {
       const column = columns.find(column => {
@@ -94,7 +95,8 @@ export function useTable<Data extends UniqueRow>({
   })
 
   // Update the row order and apply the current sort column to the incoming data
-  if (data !== prevData) {
+  const hasDataChanged = data !== prevData
+  if (hasDataChanged) {
     setPrevData(data)
     setRowOrder(data)
     if (sortByColumn) {
@@ -207,6 +209,7 @@ export function useTable<Data extends UniqueRow>({
       sortBy,
     },
     gridTemplateColumns,
+    willReactThrowOutRenderResult: hasDataChanged || haveColumnsChanged,
   }
 }