From 05c7f10756b31a91106a7067131bc30c723c911b Mon Sep 17 00:00:00 2001
From: Chandler Prall <chandler.prall@gmail.com>
Date: Tue, 19 Nov 2019 15:12:35 -0700
Subject: [PATCH] Refactored cellsUpdateFocus usage

---
 src/components/datagrid/data_grid.tsx      | 45 +++++++++-------------
 src/components/datagrid/data_grid_cell.tsx |  4 +-
 2 files changed, 20 insertions(+), 29 deletions(-)

diff --git a/src/components/datagrid/data_grid.tsx b/src/components/datagrid/data_grid.tsx
index 02d64bdaa72..ac1de404e27 100644
--- a/src/components/datagrid/data_grid.tsx
+++ b/src/components/datagrid/data_grid.tsx
@@ -340,7 +340,7 @@ function createKeyDownHandler(
               ? focusedCell[1]
               : newPageRowCount - 1;
           setFocusedCell([focusedCell[0], rowIndex]);
-          requestAnimationFrame(() => updateFocus([focusedCell[0], rowIndex]));
+          updateFocus([focusedCell[0], rowIndex]);
         }
       }
     } else if (keyCode === keyCodes.PAGE_UP) {
@@ -349,19 +349,7 @@ function createKeyDownHandler(
         const pageIndex = props.pagination.pageIndex;
         if (pageIndex > 0) {
           props.pagination.onChangePage(pageIndex - 1);
-          const newPageRowCount = computeVisibleRows({
-            rowCount,
-            pagination: {
-              ...props.pagination,
-              pageIndex: pageIndex - 1,
-            },
-          });
-          const rowIndex =
-            focusedCell[1] < newPageRowCount
-              ? focusedCell[1]
-              : newPageRowCount - 1;
-          setFocusedCell([focusedCell[0], focusedCell[1]]);
-          requestAnimationFrame(() => updateFocus([focusedCell[0], rowIndex]));
+          updateFocus(focusedCell);
         }
       }
     } else if (keyCode === (ctrlKey && keyCodes.END)) {
@@ -611,30 +599,33 @@ export const EuiDataGrid: FunctionComponent<EuiDataGridProps> = props => {
   );
 
   const [cellsUpdateFocus, setCellsUpdateFocus] = useState<
-    Array<Function[] | null[]>
-  >([]);
+    Map<string, Function>
+  >(new Map());
 
   const updateFocus = (focusedCell: [number, number]) => {
-    const updateFocus = cellsUpdateFocus[focusedCell[0]][focusedCell[1]];
-
-    if (updateFocus) {
-      updateFocus();
+    const key = `${focusedCell[0]}-${focusedCell[1]}`;
+    if (cellsUpdateFocus.has(key)) {
+      requestAnimationFrame(() => {
+        cellsUpdateFocus.get(key)!();
+      });
     }
   };
 
   const datagridContext = {
     onFocusUpdate: (cell: [number, number], updateFocus: Function) => {
       if (pagination) {
-        if (!cellsUpdateFocus[cell[0]]) {
-          cellsUpdateFocus[cell[0]] = [];
-        }
-
-        cellsUpdateFocus[cell[0]][cell[1]] = updateFocus;
+        const key = `${cell[0]}-${cell[1]}`;
 
-        setCellsUpdateFocus(cellsUpdateFocus);
+        setCellsUpdateFocus(cellsUpdateFocus => {
+          cellsUpdateFocus.set(key, updateFocus);
+          return cellsUpdateFocus;
+        });
 
         return () => {
-          cellsUpdateFocus[cell[0]][cell[1]] = null;
+          setCellsUpdateFocus(cellsUpdateFocus => {
+            cellsUpdateFocus.delete(key);
+            return cellsUpdateFocus;
+          });
         };
       }
     },
diff --git a/src/components/datagrid/data_grid_cell.tsx b/src/components/datagrid/data_grid_cell.tsx
index cbe10210316..c5d6e31fcce 100644
--- a/src/components/datagrid/data_grid_cell.tsx
+++ b/src/components/datagrid/data_grid_cell.tsx
@@ -107,7 +107,7 @@ export class EuiDataGridCell extends Component<
     cellProps: {},
     popoverIsOpen: false,
   };
-  unsubscribeCell: Function = () => {};
+  unsubscribeCell?: Function = () => {};
 
   static contextType = DataGridContext;
 
@@ -128,7 +128,7 @@ export class EuiDataGridCell extends Component<
   }
 
   componentWillUnmount() {
-    if (this.unsubscribeCell !== undefined) {
+    if (this.unsubscribeCell) {
       this.unsubscribeCell();
     }
   }