@lexical/table 0.1.12 → 0.1.15

package/LexicalTable.dev.js

@@ -6,58 +6,89 @@
  */
 'use strict';
 
+var utils = require('@lexical/utils');
 var lexical = require('lexical');
 
-/**
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- *
- * 
- */
-function addClassNamesToElement(element, ...classNames) {
-  classNames.forEach(className => {
-    if (className != null && typeof className === 'string') {
-      element.classList.add(...className.split(' '));
-    }
-  });
-}
-
-/**
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- *
- * 
- */
+/* eslint-disable sort-keys-fix/sort-keys-fix */
+const TableCellHeaderStates = {
+  NO_STATUS: 0,
+  ROW: 1,
+  COLUMN: 2,
+  BOTH: 3
+};
 class TableCellNode extends lexical.GridCellNode {
   static getType() {
     return 'tablecell';
   }
 
   static clone(node) {
-    return new TableCellNode(false, node.__colSpan, node.__key);
+    return new TableCellNode(node.__headerState, node.__colSpan, node.__width, node.__key);
   }
 
-  constructor(isHeader = false, colSpan = 1, key) {
+  constructor(headerState = TableCellHeaderStates.NO_STATUS, colSpan = 1, width, key) {
     super(colSpan, key);
-    this.__isHeader = isHeader;
-  }
-
-  getTag() {
-    return this.__isHeader ? 'th' : 'td';
+    this.__headerState = headerState;
+    this.__width = width;
   }
 
   createDOM(config) {
     const element = document.createElement(this.getTag());
-    addClassNamesToElement(element, config.theme.tableCell, this.__isHeader === true && config.theme.tableCellHeader);
+
+    if (this.__width) {
+      element.style.width = `${this.__width}px`;
+    }
+
+    utils.addClassNamesToElement(element, config.theme.tableCell, this.hasHeader() && config.theme.tableCellHeader);
     return element;
   }
 
-  updateDOM() {
-    return false;
+  getTag() {
+    return this.hasHeader() ? 'th' : 'td';
+  }
+
+  setHeaderStyles(headerState) {
+    const self = this.getWritable();
+    self.__headerState = headerState;
+    return this.__headerState;
+  }
+
+  getHeaderStyles() {
+    return this.getLatest().__headerState;
+  }
+
+  setWidth(width) {
+    const self = this.getWritable();
+    self.__width = width;
+    return this.__width;
+  }
+
+  getWidth() {
+    return this.getLatest().__width;
+  }
+
+  toggleHeaderStyle(headerStateToToggle) {
+    const self = this.getWritable();
+
+    if ((self.__headerState & headerStateToToggle) === headerStateToToggle) {
+      self.__headerState -= headerStateToToggle;
+    } else {
+      self.__headerState += headerStateToToggle;
+    }
+
+    self.__headerState = self.__headerState;
+    return self;
+  }
+
+  hasHeaderState(headerState) {
+    return (this.getHeaderStyles() & headerState) === headerState;
+  }
+
+  hasHeader() {
+    return this.getLatest().__headerState !== TableCellHeaderStates.NO_STATUS;
+  }
+
+  updateDOM(prevNode) {
+    return prevNode.__headerState !== this.__headerState || prevNode.__width !== this.__width;
   }
 
   collapseAtStart() {
@@ -69,8 +100,8 @@ class TableCellNode extends lexical.GridCellNode {
   }
 
 }
-function $createTableCellNode(isHeader) {
-  return new TableCellNode(isHeader);
+function $createTableCellNode(headerState, colSpan = 1, width) {
+  return new TableCellNode(headerState, colSpan, width);
 }
 function $isTableCellNode(node) {
   return node instanceof TableCellNode;
@@ -84,19 +115,9 @@ function $isTableCellNode(node) {
  *
  * 
  */
-function $findMatchingParent(startingNode, findFn) {
-  let curr = startingNode;
+const getSelection = () => window.getSelection();
 
-  while (curr !== lexical.$getRoot() && curr != null) {
-    if (findFn(curr)) {
-      return curr;
-    }
-
-    curr = curr.getParent();
-  }
-
-  return null;
-}
+var getDOMSelection = getSelection;
 
 /**
  * Copyright (c) Meta Platforms, Inc. and affiliates.
@@ -106,189 +127,307 @@ function $findMatchingParent(startingNode, findFn) {
  *
  * 
  */
-const LowPriority = 1;
-const CriticalPriority = 4;
 const removeHighlightStyle = document.createElement('style');
 removeHighlightStyle.appendChild(document.createTextNode('::selection{background-color: transparent}'));
-function getCellFromTarget(node) {
-  let currentNode = node;
-
-  while (currentNode != null) {
-    const nodeName = currentNode.nodeName;
-
-    if (nodeName === 'TD' || nodeName === 'TH') {
-      // $FlowFixMe: internal field
-      const cell = currentNode._cell;
-
-      if (cell === undefined) {
-        return null;
-      }
+class TableSelection {
+  constructor(editor, tableNodeKey) {
+    this.isMouseDown = false;
+    this.isHighlightingCells = false;
+    this.startX = -1;
+    this.startY = -1;
+    this.currentX = -1;
+    this.currentY = -1;
+    this.listenersToRemove = new Set();
+    this.tableNodeKey = tableNodeKey;
+    this.editor = editor;
+    this.grid = {
+      cells: [],
+      columns: 0,
+      rows: 0
+    };
+    this.gridSelection = null;
+    this.anchorCellNodeKey = null;
+    this.focusCellNodeKey = null;
+    this.trackTableGrid();
+  }
 
-      return cell;
-    }
+  getGrid() {
+    return this.grid;
+  }
 
-    currentNode = currentNode.parentNode;
+  removeListeners() {
+    Array.from(this.listenersToRemove).forEach(removeListener => removeListener());
   }
 
-  return null;
-}
-function trackTableGrid(tableNode, tableElement, editor) {
-  const cells = [];
-  const grid = {
-    cells,
-    columns: 0,
-    rows: 0
-  };
+  trackTableGrid() {
+    const observer = new MutationObserver(records => {
+      this.editor.update(() => {
+        let gridNeedsRedraw = false;
 
-  const calcSize = () => {
-    let currentNode = tableElement.firstChild;
-    let x = 0;
-    let y = 0;
-    cells.length = 0;
-
-    while (currentNode != null) {
-      const nodeMame = currentNode.nodeName;
-
-      if (nodeMame === 'TD' || nodeMame === 'TH') {
-        // $FlowFixMe: TD is always an HTMLElement
-        const elem = currentNode;
-        const cell = {
-          elem,
-          highlighted: false,
-          x,
-          y
-        }; // $FlowFixMe: internal field
+        for (let i = 0; i < records.length; i++) {
+          const record = records[i];
+          const target = record.target;
+          const nodeName = target.nodeName;
 
-        currentNode._cell = cell;
+          if (nodeName === 'TABLE' || nodeName === 'TR') {
+            gridNeedsRedraw = true;
+            break;
+          }
+        }
 
-        if (cells[y] === undefined) {
-          cells[y] = [];
+        if (!gridNeedsRedraw) {
+          return;
         }
 
-        cells[y][x] = cell;
-      } else {
-        const child = currentNode.firstChild;
+        const tableElement = this.editor.getElementByKey(this.tableNodeKey);
 
-        if (child != null) {
-          currentNode = child;
-          continue;
+        if (!tableElement) {
+          throw new Error('Expected to find TableElement in DOM');
         }
+
+        this.grid = getTableGrid(tableElement);
+      });
+    });
+    this.editor.update(() => {
+      const tableElement = this.editor.getElementByKey(this.tableNodeKey);
+
+      if (!tableElement) {
+        throw new Error('Expected to find TableElement in DOM');
       }
 
-      const sibling = currentNode.nextSibling;
+      this.grid = getTableGrid(tableElement);
+      observer.observe(tableElement, {
+        childList: true,
+        subtree: true
+      });
+    });
+  }
 
-      if (sibling != null) {
-        x++;
-        currentNode = sibling;
-        continue;
+  clearHighlight() {
+    this.editor.update(() => {
+      const tableNode = lexical.$getNodeByKey(this.tableNodeKey);
+
+      if (!$isTableNode(tableNode)) {
+        throw new Error('Expected TableNode.');
       }
 
-      const parent = currentNode.parentNode;
+      const tableElement = this.editor.getElementByKey(this.tableNodeKey);
+
+      if (!tableElement) {
+        throw new Error('Expected to find TableElement in DOM');
+      }
+
+      const grid = getTableGrid(tableElement);
+      this.isHighlightingCells = false;
+      this.isMouseDown = false;
+      this.startX = -1;
+      this.startY = -1;
+      this.currentX = -1;
+      this.currentY = -1;
+      $updateDOMForSelection(grid, null);
+      this.gridSelection = null;
+      lexical.$setSelection(null);
+      const parent = removeHighlightStyle.parentNode;
 
       if (parent != null) {
-        const parentSibling = parent.nextSibling;
+        parent.removeChild(removeHighlightStyle);
+      }
+    });
+  }
 
-        if (parentSibling == null) {
-          break;
-        }
+  adjustFocusCellForSelection(cell, ignoreStart = false) {
+    this.editor.update(() => {
+      this.isMouseDown = true;
+      const tableNode = lexical.$getNodeByKey(this.tableNodeKey);
 
-        y++;
-        x = 0;
-        currentNode = parentSibling;
+      if (!$isTableNode(tableNode)) {
+        throw new Error('Expected TableNode.');
       }
-    }
 
-    grid.columns = x + 1;
-    grid.rows = y + 1;
-    tableNode.setGrid(grid);
-  };
+      const tableElement = this.editor.getElementByKey(this.tableNodeKey);
 
-  const observer = new MutationObserver(records => {
-    editor.update(() => {
-      let gridNeedsRedraw = false;
+      if (!tableElement) {
+        throw new Error('Expected to find TableElement in DOM');
+      }
 
-      for (let i = 0; i < records.length; i++) {
-        const record = records[i];
-        const target = record.target;
-        const nodeName = target.nodeName;
+      const cellX = cell.x;
+      const cellY = cell.y;
 
-        if (nodeName === 'TABLE' || nodeName === 'TR') {
-          gridNeedsRedraw = true;
-          break;
+      if (!this.isHighlightingCells && (this.startX !== cellX || this.startY !== cellY || ignoreStart)) {
+        const domSelection = getDOMSelection();
+        const anchorNode = domSelection.anchorNode;
+
+        if (anchorNode !== null) {
+          // Collapse the selection
+          domSelection.setBaseAndExtent(anchorNode, 0, anchorNode, 0);
         }
-      }
 
-      if (!gridNeedsRedraw) {
+        this.isHighlightingCells = true;
+
+        if (document.body) {
+          document.body.appendChild(removeHighlightStyle);
+        }
+      } else if (cellX === this.currentX && cellY === this.currentY) {
         return;
       }
 
-      calcSize();
+      this.currentX = cellX;
+      this.currentY = cellY;
+
+      if (this.isHighlightingCells) {
+        const focusTableCellNode = lexical.$getNearestNodeFromDOMNode(cell.elem);
+
+        if (this.gridSelection != null && this.anchorCellNodeKey != null && $isTableCellNode(focusTableCellNode)) {
+          const focusNodeKey = focusTableCellNode.getKey();
+          this.gridSelection = lexical.$createGridSelection();
+          this.focusCellNodeKey = focusNodeKey;
+          this.gridSelection.set(this.tableNodeKey, // $FlowFixMe This is not null, as you can see in the statement above.
+          this.anchorCellNodeKey, this.focusCellNodeKey);
+          lexical.$setSelection(this.gridSelection);
+          $updateDOMForSelection(this.grid, this.gridSelection);
+        }
+      }
     });
-  });
-  observer.observe(tableElement, {
-    childList: true,
-    subtree: true
-  });
-  calcSize();
-  return grid;
-}
-function updateCells(fromX, toX, fromY, toY, cells) {
-  const highlighted = [];
+  }
 
-  for (let y = 0; y < cells.length; y++) {
-    const row = cells[y];
+  setAnchorCellForSelection(cell) {
+    this.editor.update(() => {
+      this.startX = cell.x;
+      this.startY = cell.y;
+      this.isMouseDown = true;
+      const anchorTableCellNode = lexical.$getNearestNodeFromDOMNode(cell.elem);
+
+      if ($isTableCellNode(anchorTableCellNode)) {
+        const anchorNodeKey = anchorTableCellNode.getKey();
+        this.gridSelection = lexical.$createGridSelection();
+        this.anchorCellNodeKey = anchorNodeKey;
+      }
+    });
+    document.addEventListener('mouseup', () => {
+      this.isMouseDown = false;
+    }, {
+      capture: true,
+      once: true
+    });
+  }
 
-    for (let x = 0; x < row.length; x++) {
-      const cell = row[x];
-      const elemStyle = cell.elem.style;
+  formatCells(type) {
+    this.editor.update(() => {
+      const selection = lexical.$getSelection();
 
-      if (x >= fromX && x <= toX && y >= fromY && y <= toY) {
-        if (!cell.highlighted) {
-          cell.highlighted = true;
-          elemStyle.setProperty('background-color', 'rgb(163, 187, 255)');
-          elemStyle.setProperty('caret-color', 'transparent');
+      if (!lexical.$isGridSelection(selection)) {
+        {
+          throw Error(`Expected grid selection`);
         }
+      } // This is to make Flow play ball.
 
-        highlighted.push(cell);
-      } else if (cell.highlighted) {
-        cell.highlighted = false;
-        elemStyle.removeProperty('background-color');
-        elemStyle.removeProperty('caret-color');
 
-        if (!cell.elem.getAttribute('style')) {
-          cell.elem.removeAttribute('style');
+      const formatSelection = lexical.$createRangeSelection();
+      const anchor = formatSelection.anchor;
+      const focus = formatSelection.focus;
+      selection.getNodes().forEach(cellNode => {
+        if (lexical.$isElementNode(cellNode) && cellNode.getTextContentSize() !== 0) {
+          anchor.set(cellNode.getKey(), 0, 'element');
+          focus.set(cellNode.getKey(), cellNode.getChildrenSize(), 'element');
+          formatSelection.formatText(type);
         }
+      });
+      lexical.$setSelection(selection);
+    });
+  }
+
+  clearText() {
+    this.editor.update(() => {
+      const tableNode = lexical.$getNodeByKey(this.tableNodeKey);
+
+      if (!$isTableNode(tableNode)) {
+        throw new Error('Expected TableNode.');
       }
-    }
+
+      const selection = lexical.$getSelection();
+
+      if (!lexical.$isGridSelection(selection)) {
+        {
+          throw Error(`Expected grid selection`);
+        }
+      }
+
+      const selectedNodes = selection.getNodes();
+
+      if (selectedNodes.length === this.grid.columns * this.grid.rows) {
+        tableNode.selectPrevious(); // Delete entire table
+
+        tableNode.remove();
+        this.clearHighlight();
+        return;
+      }
+
+      selectedNodes.forEach(cellNode => {
+        if (lexical.$isElementNode(cellNode)) {
+          const paragraphNode = lexical.$createParagraphNode();
+          const textNode = lexical.$createTextNode();
+          paragraphNode.append(textNode);
+          cellNode.append(paragraphNode);
+          cellNode.getChildren().forEach(child => {
+            if (child !== paragraphNode) {
+              child.remove();
+            }
+          });
+        }
+      });
+      $updateDOMForSelection(this.grid, null);
+      lexical.$setSelection(null);
+    });
   }
 
-  return highlighted;
 }
-function $applyCustomTableHandlers(tableNode, tableElement, editor) {
+
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * 
+ */
+const CriticalPriority = 4;
+const LEXICAL_ELEMENT_KEY = '__lexicalTableSelection';
+function applyTableHandlers(tableNode, tableElement, editor) {
   const rootElement = editor.getRootElement();
 
   if (rootElement === null) {
     throw new Error('No root element.');
   }
 
-  trackTableGrid(tableNode, tableElement, editor);
-  const grid = tableNode.getGrid();
-  let isSelected = false;
-  let isHighlightingCells = false;
-  let startX = -1;
-  let startY = -1;
-  let currentX = -1;
-  let currentY = -1;
-  let highlightedCells = [];
-  const editorListeners = new Set();
-  let deleteCharacterListener = null;
+  const tableSelection = new TableSelection(editor, tableNode.getKey());
+  attachTableSelectionToTableElement(tableElement, tableSelection);
+  tableElement.addEventListener('dblclick', event => {
+    // $FlowFixMe: event.target is always a Node on the DOM
+    const cell = getCellFromTarget(event.target);
+
+    if (cell !== null) {
+      event.preventDefault();
+      event.stopImmediatePropagation();
+      event.stopPropagation();
+      tableSelection.setAnchorCellForSelection(cell);
+      tableSelection.adjustFocusCellForSelection(cell, true);
+      tableSelection.isMouseDown = false;
+    }
+  }); // This is the anchor of the selection.
 
-  if (grid == null) {
-    throw new Error('Table grid not found.');
-  }
+  tableElement.addEventListener('mousedown', event => {
+    setTimeout(() => {
+      // $FlowFixMe: event.target is always a Node on the DOM
+      const cell = getCellFromTarget(event.target);
+
+      if (cell !== null) {
+        tableSelection.setAnchorCellForSelection(cell);
+      }
+    }, 0);
+  }); // This is adjusting the focus of the selection.
 
   tableElement.addEventListener('mousemove', event => {
-    if (isSelected) {
+    if (tableSelection.isMouseDown) {
       // $FlowFixMe: event.target is always a Node on the DOM
       const cell = getCellFromTarget(event.target);
 
@@ -296,280 +435,306 @@ function $applyCustomTableHandlers(tableNode, tableElement, editor) {
         const cellX = cell.x;
         const cellY = cell.y;
 
-        if (!isHighlightingCells && (startX !== cellX || startY !== cellY)) {
+        if (tableSelection.isMouseDown && (tableSelection.startX !== cellX || tableSelection.startY !== cellY || tableSelection.isHighlightingCells)) {
           event.preventDefault();
-          const windowSelection = window.getSelection(); // Collapse the selection
-
-          windowSelection.setBaseAndExtent(windowSelection.anchorNode, 0, windowSelection.anchorNode, 0);
-          isHighlightingCells = true;
-
-          if (document.body) {
-            document.body.appendChild(removeHighlightStyle);
-          }
-
-          if (deleteCharacterListener === null) {
-            deleteCharacterListener = editor.addListener('command', (type, payload) => {
-              if (type === 'deleteCharacter') {
-                if (highlightedCells.length === grid.columns * grid.rows) {
-                  tableNode.selectPrevious(); // Delete entire table
-
-                  tableNode.remove();
-                  clearHighlight();
-                  return true;
-                }
-
-                highlightedCells.forEach(({
-                  elem
-                }) => {
-                  const cellNode = lexical.$getNearestNodeFromDOMNode(elem);
-
-                  if (lexical.$isElementNode(cellNode)) {
-                    const paragraphNode = lexical.$createParagraphNode();
-                    const textNode = lexical.$createTextNode();
-                    paragraphNode.append(textNode);
-                    cellNode.append(paragraphNode);
-                    cellNode.getChildren().forEach(child => {
-                      if (child !== paragraphNode) {
-                        child.remove();
-                      }
-                    });
-                  }
-                });
-                tableNode.setSelectionState(null);
-                lexical.$setSelection(null);
-                return true;
-              } else if (type === 'formatText') {
-                formatCells(payload);
-                return true;
-              } else if (type === 'insertText') {
-                clearHighlight();
-                return false;
-              }
-
-              return false;
-            }, LowPriority);
-            editorListeners.add(deleteCharacterListener);
-          }
-        } else if (cellX === currentX && cellY === currentY) {
-          return;
-        }
-
-        currentX = cellX;
-        currentY = cellY;
-
-        if (isHighlightingCells) {
-          const fromX = Math.min(startX, currentX);
-          const toX = Math.max(startX, currentX);
-          const fromY = Math.min(startY, currentY);
-          const toY = Math.max(startY, currentY);
-          editor.update(() => {
-            highlightedCells = tableNode.setSelectionState({
-              fromX,
-              fromY,
-              toX,
-              toY
-            });
-          });
+          tableSelection.adjustFocusCellForSelection(cell);
         }
       }
     }
   });
+  tableElement.addEventListener('mouseup', event => {
+    if (tableSelection.isMouseDown) {
+      tableSelection.isMouseDown = false;
+    }
+  }); // Select entire table at this point, when grid selection is ready.
 
-  const clearHighlight = () => {
-    editor.update(() => {
-      isHighlightingCells = false;
-      isSelected = false;
-      startX = -1;
-      startY = -1;
-      currentX = -1;
-      currentY = -1;
-      editor.update(() => {
-        tableNode.setSelectionState(null);
-      });
-      highlightedCells = [];
-
-      if (deleteCharacterListener !== null) {
-        deleteCharacterListener();
-        deleteCharacterListener = null;
-        editorListeners.delete(deleteCharacterListener);
-      }
+  tableElement.addEventListener('mouseleave', event => {
+    if (tableSelection.isMouseDown) {
+      return;
+    }
+  }); // Clear selection when clicking outside of dom.
 
-      const parent = removeHighlightStyle.parentNode;
+  const mouseDownCallback = e => {
+    editor.update(() => {
+      const selection = lexical.$getSelection();
 
-      if (parent != null) {
-        parent.removeChild(removeHighlightStyle);
+      if (lexical.$isGridSelection(selection) && selection.gridKey === tableSelection.tableNodeKey && rootElement.contains(e.target)) {
+        return tableSelection.clearHighlight();
       }
     });
   };
 
-  tableElement.addEventListener('mouseleave', event => {
-    if (isSelected) {
-      return;
-    }
-  });
+  window.addEventListener('mousedown', mouseDownCallback);
+  tableSelection.listenersToRemove.add(() => window.removeEventListener(mouseDownCallback));
+  tableSelection.listenersToRemove.add(editor.addListener('command', (type, payload) => {
+    const selection = lexical.$getSelection();
 
-  const formatCells = type => {
-    let selection = lexical.$getSelection();
+    if (lexical.$isGridSelection(selection)) {
+      if (type === 'deleteCharacter' || type === 'keyBackspace') {
+        const event = payload;
+        event.preventDefault();
+        event.stopPropagation();
+        tableSelection.clearText();
+        return true;
+      } else if (type === 'formatText') {
+        tableSelection.formatCells(payload);
+        return true;
+      } else if (type === 'insertText') {
+        tableSelection.clearHighlight();
+        return false;
+      }
+    } else if (lexical.$isRangeSelection(selection)) {
+      const tableCellNode = utils.$findMatchingParent(selection.anchor.getNode(), n => $isTableCellNode(n));
 
-    if (!lexical.$isRangeSelection(selection)) {
-      selection = lexical.$createRangeSelection();
-    } // This is to make Flow play ball.
+      if (!$isTableCellNode(tableCellNode)) {
+        return false;
+      }
 
+      if (type === 'deleteCharacter') {
+        if (selection.isCollapsed() && selection.anchor.offset === 0 && selection.anchor.getNode().getPreviousSiblings().length === 0) {
+          return true;
+        }
+      }
 
-    const formatSelection = selection;
-    const anchor = formatSelection.anchor;
-    const focus = formatSelection.focus;
-    highlightedCells.forEach(highlightedCell => {
-      const cellNode = lexical.$getNearestNodeFromDOMNode(highlightedCell.elem);
+      if (type === 'keyTab') {
+        const event = payload;
 
-      if (lexical.$isElementNode(cellNode) && cellNode.getTextContentSize() !== 0) {
-        anchor.set(cellNode.getKey(), 0, 'element');
-        focus.set(cellNode.getKey(), cellNode.getChildrenSize(), 'element');
-        formatSelection.formatText(type);
+        if (selection.isCollapsed()) {
+          const currentCords = tableNode.getCordsFromCellNode(tableCellNode, tableSelection.grid);
+          event.preventDefault();
+          selectGridNodeInDirection(tableSelection, tableNode, currentCords.x, currentCords.y, !event.shiftKey && type === 'keyTab' ? 'forward' : 'backward');
+          return true;
+        }
       }
-    }); // Collapse selection
 
-    selection.anchor.set(selection.anchor.key, selection.anchor.offset, selection.anchor.type);
-    selection.focus.set(selection.anchor.key, selection.anchor.offset, selection.anchor.type);
-    lexical.$setSelection(selection);
-  };
+      if (type === 'keyArrowDown' || type === 'keyArrowUp' || type === 'keyArrowLeft' || type === 'keyArrowRight') {
+        const event = payload;
 
-  tableElement.addEventListener('mousedown', event => {
-    if (isSelected) {
-      if (isHighlightingCells) {
-        clearHighlight();
-      }
+        if (selection.isCollapsed()) {
+          const currentCords = tableNode.getCordsFromCellNode(tableCellNode, tableSelection.grid);
+          const elementParentNode = utils.$findMatchingParent(selection.anchor.getNode(), n => lexical.$isElementNode(n));
 
-      return;
+          if (elementParentNode == null) {
+            throw new Error('Expected BlockNode Parent');
+          }
+
+          const firstChild = tableCellNode.getFirstChild();
+          const lastChild = tableCellNode.getLastChild();
+          const isSelectionInFirstBlock = firstChild && elementParentNode.isParentOf(firstChild) || elementParentNode === firstChild;
+          const isSelectionInLastBlock = lastChild && elementParentNode.isParentOf(lastChild) || elementParentNode === lastChild;
+
+          if (type === 'keyArrowUp' && isSelectionInFirstBlock || type === 'keyArrowDown' && isSelectionInLastBlock) {
+            event.preventDefault();
+            event.stopImmediatePropagation();
+            event.stopPropagation();
+            selectGridNodeInDirection(tableSelection, tableNode, currentCords.x, currentCords.y, type === 'keyArrowUp' ? 'up' : 'down');
+            return true;
+          }
+
+          if (type === 'keyArrowLeft' && selection.anchor.offset === 0 || type === 'keyArrowRight' && selection.anchor.offset === selection.anchor.getNode().getTextContentSize()) {
+            event.preventDefault();
+            event.stopImmediatePropagation();
+            event.stopPropagation();
+            selectGridNodeInDirection(tableSelection, tableNode, currentCords.x, currentCords.y, type === 'keyArrowLeft' ? 'backward' : 'forward');
+            return true;
+          }
+        }
+      }
     }
 
-    setTimeout(() => {
-      if (isHighlightingCells) {
-        clearHighlight();
-      } // $FlowFixMe: event.target is always a Node on the DOM
+    return false;
+  }, CriticalPriority));
+  return tableSelection;
+}
+function attachTableSelectionToTableElement(tableElement, tableSelection) {
+  // $FlowFixMe
+  tableElement[LEXICAL_ELEMENT_KEY] = tableSelection;
+}
+function getTableSelectionFromTableElement(tableElement) {
+  // $FlowFixMe
+  return tableElement[LEXICAL_ELEMENT_KEY];
+}
+function getCellFromTarget(node) {
+  let currentNode = node;
 
+  while (currentNode != null) {
+    const nodeName = currentNode.nodeName;
 
-      const cell = getCellFromTarget(event.target);
+    if (nodeName === 'TD' || nodeName === 'TH') {
+      // $FlowFixMe: internal field
+      const cell = currentNode._cell;
 
-      if (cell !== null) {
-        isSelected = true;
-        startX = cell.x;
-        startY = cell.y;
-        document.addEventListener('mouseup', () => {
-          isSelected = false;
-        }, {
-          capture: true,
-          once: true
-        });
+      if (cell === undefined) {
+        return null;
       }
-    }, 0);
-  });
-  window.addEventListener('click', e => {
-    if (highlightedCells.length > 0 && !tableElement.contains(e.target) && rootElement.contains(e.target)) {
-      editor.update(() => {
-        tableNode.setSelectionState(null);
-      });
+
+      return cell;
     }
-  });
 
-  const selectGridNodeInDirection = (x, y, direction) => {
-    switch (direction) {
-      case 'backward':
-      case 'forward':
-        {
-          const isForward = direction === 'forward';
+    currentNode = currentNode.parentNode;
+  }
 
-          if (x !== (isForward ? grid.columns - 1 : 0)) {
-            tableNode.getCellNodeFromCords(x + (isForward ? 1 : -1), y).select();
-          } else {
-            if (y !== (isForward ? grid.rows - 1 : 0)) {
-              tableNode.getCellNodeFromCords(isForward ? 0 : grid.columns - 1, y + (isForward ? 1 : -1)).select();
-            } else if (!isForward) {
-              tableNode.selectPrevious();
-            } else {
-              tableNode.selectNext();
-            }
-          }
+  return null;
+}
+function getTableGrid(tableElement) {
+  const cells = [];
+  const grid = {
+    cells,
+    columns: 0,
+    rows: 0
+  };
+  let currentNode = tableElement.firstChild;
+  let x = 0;
+  let y = 0;
+  cells.length = 0;
 
-          return true;
-        }
+  while (currentNode != null) {
+    const nodeMame = currentNode.nodeName;
+
+    if (nodeMame === 'TD' || nodeMame === 'TH') {
+      // $FlowFixMe: TD is always an HTMLElement
+      const elem = currentNode;
+      const cell = {
+        elem,
+        highlighted: false,
+        x,
+        y
+      }; // $FlowFixMe: internal field
+
+      currentNode._cell = cell;
+
+      if (cells[y] === undefined) {
+        cells[y] = [];
+      }
 
-      case 'up':
-        {
-          if (y !== 0) {
-            tableNode.getCellNodeFromCords(x, y - 1).select();
-          } else {
-            tableNode.selectPrevious();
-          }
+      cells[y][x] = cell;
+    } else {
+      const child = currentNode.firstChild;
 
-          return true;
-        }
+      if (child != null) {
+        currentNode = child;
+        continue;
+      }
+    }
 
-      case 'down':
-        {
-          if (y !== grid.rows - 1) {
-            tableNode.getCellNodeFromCords(x, y + 1).select();
-          } else {
-            tableNode.selectNext();
-          }
+    const sibling = currentNode.nextSibling;
 
-          return true;
-        }
+    if (sibling != null) {
+      x++;
+      currentNode = sibling;
+      continue;
     }
 
-    return false;
-  };
+    const parent = currentNode.parentNode;
 
-  const genericCommandListener = editor.addListener('command', (type, payload) => {
-    const selection = lexical.$getSelection();
+    if (parent != null) {
+      const parentSibling = parent.nextSibling;
+
+      if (parentSibling == null) {
+        break;
+      }
 
-    if (!lexical.$isRangeSelection(selection)) {
-      return false;
+      y++;
+      x = 0;
+      currentNode = parentSibling;
     }
+  }
 
-    const tableCellNode = $findMatchingParent(selection.anchor.getNode(), n => $isTableCellNode(n));
+  grid.columns = x + 1;
+  grid.rows = y + 1;
+  return grid;
+}
+function $updateDOMForSelection(grid, gridSelection) {
+  const highlightedCells = [];
+  const {
+    cells
+  } = grid;
+  const selectedCellNodes = new Set(gridSelection ? gridSelection.getNodes() : []);
 
-    if (!$isTableCellNode(tableCellNode)) {
-      return false;
-    }
+  for (let y = 0; y < cells.length; y++) {
+    const row = cells[y];
 
-    if (type === 'deleteCharacter') {
-      if (highlightedCells.length === 0 && selection.isCollapsed() && selection.anchor.offset === 0 && selection.anchor.getNode().getPreviousSiblings().length === 0) {
-        return true;
+    for (let x = 0; x < row.length; x++) {
+      const cell = row[x];
+      const elemStyle = cell.elem.style;
+      const lexicalNode = lexical.$getNearestNodeFromDOMNode(cell.elem);
+
+      if (lexicalNode && selectedCellNodes.has(lexicalNode)) {
+        cell.highlighted = true;
+        elemStyle.setProperty('background-color', 'rgb(163, 187, 255)');
+        elemStyle.setProperty('caret-color', 'transparent');
+        highlightedCells.push(cell);
+      } else {
+        cell.highlighted = false;
+        elemStyle.removeProperty('background-color');
+        elemStyle.removeProperty('caret-color');
+
+        if (!cell.elem.getAttribute('style')) {
+          cell.elem.removeAttribute('style');
+        }
       }
     }
+  }
 
-    if (type === 'keyTab') {
-      const event = payload;
+  return highlightedCells;
+}
+
+const selectGridNodeInDirection = (tableSelection, tableNode, x, y, direction) => {
+  switch (direction) {
+    case 'backward':
+    case 'forward':
+      {
+        const isForward = direction === 'forward';
+
+        if (x !== (isForward ? tableSelection.grid.columns - 1 : 0)) {
+          selectTableCellNode(tableNode.getCellNodeFromCordsOrThrow(x + (isForward ? 1 : -1), y, tableSelection.grid));
+        } else {
+          if (y !== (isForward ? tableSelection.grid.rows - 1 : 0)) {
+            selectTableCellNode(tableNode.getCellNodeFromCordsOrThrow(isForward ? 0 : tableSelection.grid.columns - 1, y + (isForward ? 1 : -1), tableSelection.grid));
+          } else if (!isForward) {
+            tableNode.selectPrevious();
+          } else {
+            tableNode.selectNext();
+          }
+        }
 
-      if (selection.isCollapsed() && highlightedCells.length === 0) {
-        const currentCords = tableNode.getCordsFromCellNode(tableCellNode);
-        event.preventDefault();
-        selectGridNodeInDirection(currentCords.x, currentCords.y, !event.shiftKey && type === 'keyTab' ? 'forward' : 'backward');
         return true;
       }
-    }
 
-    if (type === 'keyArrowDown' || type === 'keyArrowUp') {
-      const event = payload;
+    case 'up':
+      {
+        if (y !== 0) {
+          selectTableCellNode(tableNode.getCellNodeFromCordsOrThrow(x, y - 1, tableSelection.grid));
+        } else {
+          tableNode.selectPrevious();
+        }
 
-      if (selection.isCollapsed() && highlightedCells.length === 0) {
-        const currentCords = tableNode.getCordsFromCellNode(tableCellNode);
-        const elementParentNode = $findMatchingParent(selection.anchor.getNode(), n => lexical.$isElementNode(n));
+        return true;
+      }
 
-        if (type === 'keyArrowUp' && elementParentNode === tableCellNode.getFirstChild() || type === 'keyArrowDown' && elementParentNode === tableCellNode.getLastChild()) {
-          event.preventDefault();
-          event.stopImmediatePropagation();
-          selectGridNodeInDirection(currentCords.x, currentCords.y, type === 'keyArrowUp' ? 'up' : 'down');
-          return true;
+    case 'down':
+      {
+        if (y !== tableSelection.grid.rows - 1) {
+          selectTableCellNode(tableNode.getCellNodeFromCordsOrThrow(x, y + 1, tableSelection.grid));
+        } else {
+          tableNode.selectNext();
         }
+
+        return true;
       }
-    }
+  }
 
-    return false;
-  }, CriticalPriority);
-  editorListeners.add(genericCommandListener);
-  return () => Array.from(editorListeners).forEach(removeListener => removeListener ? removeListener() : null);
+  return false;
+};
+
+function selectTableCellNode(tableCell) {
+  const possibleParagraph = tableCell.getChildren().find(n => lexical.$isParagraphNode(n));
+
+  if (lexical.$isParagraphNode(possibleParagraph)) {
+    possibleParagraph.selectEnd();
+  } else {
+    tableCell.selectEnd();
+  }
 }
 
 /**
@@ -585,21 +750,17 @@ class TableNode extends lexical.GridNode {
     return 'table';
   }
 
-  static clone(node, selectionShape, grid) {
-    // TODO: selectionShape and grid aren't being deeply cloned?
-    // They shouldn't really be on the table node IMO.
-    return new TableNode(node.__key, node.__selectionShape, node.__grid);
+  static clone(node) {
+    return new TableNode(node.__key);
   }
 
-  constructor(key, selectionShape, grid) {
+  constructor(key) {
     super(key);
-    this.__selectionShape = selectionShape;
-    this.__grid = grid;
   }
 
   createDOM(config, editor) {
     const element = document.createElement('table');
-    addClassNamesToElement(element, config.theme.table);
+    utils.addClassNamesToElement(element, config.theme.table);
     return element;
   }
 
@@ -615,26 +776,7 @@ class TableNode extends lexical.GridNode {
     return false;
   }
 
-  setSelectionState(selectionShape) {
-    const self = this.getWritable();
-    self.__selectionShape = selectionShape;
-    const grid = this.getGrid();
-    if (grid == null) return [];
-
-    if (!selectionShape) {
-      return updateCells(-1, -1, -1, -1, grid.cells);
-    }
-
-    return updateCells(selectionShape.fromX, selectionShape.toX, selectionShape.fromY, selectionShape.toY, grid.cells);
-  }
-
-  getSelectionState() {
-    return this.getLatest().__selectionShape;
-  }
-
-  getCordsFromCellNode(tableCellNode) {
-    const grid = this.getGrid();
-
+  getCordsFromCellNode(tableCellNode, grid) {
     if (!grid) {
       throw Error(`Grid not found.`);
     }
@@ -669,9 +811,7 @@ class TableNode extends lexical.GridNode {
     throw new Error('Cell not found in table.');
   }
 
-  getCellNodeFromCords(x, y) {
-    const grid = this.getGrid();
-
+  getCellNodeFromCords(x, y, grid) {
     if (!grid) {
       throw Error(`Grid not found.`);
     }
@@ -682,13 +822,13 @@ class TableNode extends lexical.GridNode {
     const row = cells[y];
 
     if (row == null) {
-      throw new Error(`Table row x:"${y}" not found.`);
+      return null;
     }
 
     const cell = row[x];
 
     if (cell == null) {
-      throw new Error(`Table cell y:"${x}" in row x:"${y}" not found.`);
+      return null;
     }
 
     const node = lexical.$getNearestNodeFromDOMNode(cell.elem);
@@ -697,17 +837,17 @@ class TableNode extends lexical.GridNode {
       return node;
     }
 
-    throw new Error('Node at cords not TableCellNode.');
+    return null;
   }
 
-  setGrid(grid) {
-    const self = this.getWritable();
-    self.__grid = grid;
-    return self;
-  }
+  getCellNodeFromCordsOrThrow(x, y, grid) {
+    const node = this.getCellNodeFromCords(x, y, grid);
 
-  getGrid() {
-    return this.getLatest().__grid;
+    if (!node) {
+      throw new Error('Node at cords not TableCellNode.');
+    }
+
+    return node;
   }
 
   canSelectBefore() {
@@ -715,6 +855,15 @@ class TableNode extends lexical.GridNode {
   }
 
 }
+function $getElementGridForTableNode(editor, tableNode) {
+  const tableElement = editor.getElementByKey(tableNode.getKey());
+
+  if (tableElement == null) {
+    throw new Error('Table Element Not Found');
+  }
+
+  return getTableGrid(tableElement);
+}
 function $createTableNode() {
   return new TableNode();
 }
@@ -736,21 +885,37 @@ class TableRowNode extends lexical.GridRowNode {
   }
 
   static clone(node) {
-    return new TableRowNode(node.__key);
+    return new TableRowNode(node.__height, node.__key);
   }
 
-  constructor(key) {
+  constructor(height, key) {
     super(key);
+    this.__height = height;
   }
 
   createDOM(config) {
     const element = document.createElement('tr');
-    addClassNamesToElement(element, config.theme.tableRow);
+
+    if (this.__height) {
+      element.style.height = `${this.__height}px`;
+    }
+
+    utils.addClassNamesToElement(element, config.theme.tableRow);
     return element;
   }
 
-  updateDOM() {
-    return false;
+  setHeight(height) {
+    const self = this.getWritable();
+    self.__height = height;
+    return this.__height;
+  }
+
+  getHeight() {
+    return this.getLatest().__height;
+  }
+
+  updateDOM(prevNode) {
+    return prevNode.__height !== this.__height;
   }
 
   canBeEmpty() {
@@ -758,8 +923,8 @@ class TableRowNode extends lexical.GridRowNode {
   }
 
 }
-function $createTableRowNode() {
-  return new TableRowNode();
+function $createTableRowNode(height) {
+  return new TableRowNode(height);
 }
 function $isTableRowNode(node) {
   return node instanceof TableRowNode;
@@ -773,14 +938,21 @@ function $isTableRowNode(node) {
  *
  * 
  */
-function $createTableNodeWithDimensions(rowCount, columnCount, includeHeader = true) {
+function $createTableNodeWithDimensions(rowCount, columnCount, includeHeaders = true) {
   const tableNode = $createTableNode();
 
   for (let iRow = 0; iRow < rowCount; iRow++) {
     const tableRowNode = $createTableRowNode();
 
     for (let iColumn = 0; iColumn < columnCount; iColumn++) {
-      const tableCellNode = $createTableCellNode(iRow === 0 && includeHeader);
+      let headerState = TableCellHeaderStates.NO_STATUS;
+
+      if (includeHeaders) {
+        if (iRow === 0) headerState |= TableCellHeaderStates.ROW;
+        if (iColumn === 0) headerState |= TableCellHeaderStates.COLUMN;
+      }
+
+      const tableCellNode = $createTableCellNode(headerState);
       const paragraphNode = lexical.$createParagraphNode();
       paragraphNode.append(lexical.$createTextNode());
       tableCellNode.append(paragraphNode);
@@ -793,7 +965,7 @@ function $createTableNodeWithDimensions(rowCount, columnCount, includeHeader = t
   return tableNode;
 }
 function $getTableCellNodeFromLexicalNode(startingNode) {
-  const node = $findMatchingParent(startingNode, n => $isTableCellNode(n));
+  const node = utils.$findMatchingParent(startingNode, n => $isTableCellNode(n));
 
   if ($isTableCellNode(node)) {
     return node;
@@ -802,7 +974,7 @@ function $getTableCellNodeFromLexicalNode(startingNode) {
   return null;
 }
 function $getTableRowNodeFromTableCellNodeOrThrow(startingNode) {
-  const node = $findMatchingParent(startingNode, n => $isTableRowNode(n));
+  const node = utils.$findMatchingParent(startingNode, n => $isTableRowNode(n));
 
   if ($isTableRowNode(node)) {
     return node;
@@ -811,7 +983,7 @@ function $getTableRowNodeFromTableCellNodeOrThrow(startingNode) {
   throw new Error('Expected table cell to be inside of table row.');
 }
 function $getTableNodeFromLexicalNodeOrThrow(startingNode) {
-  const node = $findMatchingParent(startingNode, n => $isTableNode(n));
+  const node = utils.$findMatchingParent(startingNode, n => $isTableNode(n));
 
   if ($isTableNode(node)) {
     return node;
@@ -828,6 +1000,19 @@ function $getTableColumnIndexFromTableCellNode(tableCellNode) {
   const tableRowNode = $getTableRowNodeFromTableCellNodeOrThrow(tableCellNode);
   return tableRowNode.getChildren().findIndex(n => n.is(tableCellNode));
 }
+function $getTableCellSiblingsFromTableCellNode(tableCellNode, grid) {
+  const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);
+  const {
+    x,
+    y
+  } = tableNode.getCordsFromCellNode(tableCellNode, grid);
+  return {
+    above: tableNode.getCellNodeFromCords(x, y - 1, grid),
+    below: tableNode.getCellNodeFromCords(x, y + 1, grid),
+    left: tableNode.getCellNodeFromCords(x - 1, y, grid),
+    right: tableNode.getCellNodeFromCords(x + 1, y, grid)
+  };
+}
 function $removeTableRowAtIndex(tableNode, indexToDelete) {
   const tableRows = tableNode.getChildren();
 
@@ -839,7 +1024,7 @@ function $removeTableRowAtIndex(tableNode, indexToDelete) {
   targetRowNode.remove();
   return tableNode;
 }
-function $insertTableRow(tableNode, targetIndex, shouldInsertAfter = true, rowCount) {
+function $insertTableRow(tableNode, targetIndex, shouldInsertAfter = true, rowCount, grid) {
   const tableRows = tableNode.getChildren();
 
   if (targetIndex >= tableRows.length || targetIndex < 0) {
@@ -849,12 +1034,30 @@ function $insertTableRow(tableNode, targetIndex, shouldInsertAfter = true, rowCo
   const targetRowNode = tableRows[targetIndex];
 
   if ($isTableRowNode(targetRowNode)) {
-    for (let i = 0; i < rowCount; i++) {
-      const tableColumnCount = targetRowNode.getChildren().length;
+    for (let r = 0; r < rowCount; r++) {
+      const tableRowCells = targetRowNode.getChildren();
+      const tableColumnCount = tableRowCells.length;
       const newTableRowNode = $createTableRowNode();
 
-      for (let j = 0; j < tableColumnCount; j++) {
-        const tableCellNode = $createTableCellNode(false);
+      for (let c = 0; c < tableColumnCount; c++) {
+        const tableCellFromTargetRow = tableRowCells[c];
+
+        if (!$isTableCellNode(tableCellFromTargetRow)) {
+          throw Error(`Expected table cell`);
+        }
+
+        const {
+          above,
+          below
+        } = $getTableCellSiblingsFromTableCellNode(tableCellFromTargetRow, grid);
+        let headerState = TableCellHeaderStates.NO_STATUS;
+        const width = above && above.getWidth() || below && below.getWidth() || null;
+
+        if (above && above.hasHeaderState(TableCellHeaderStates.COLUMN) || below && below.hasHeaderState(TableCellHeaderStates.COLUMN)) {
+          headerState |= TableCellHeaderStates.COLUMN;
+        }
+
+        const tableCellNode = $createTableCellNode(headerState, 1, width);
         tableCellNode.append(lexical.$createParagraphNode());
         newTableRowNode.append(tableCellNode);
       }
@@ -874,12 +1077,18 @@ function $insertTableRow(tableNode, targetIndex, shouldInsertAfter = true, rowCo
 function $insertTableColumn(tableNode, targetIndex, shouldInsertAfter = true, columnCount) {
   const tableRows = tableNode.getChildren();
 
-  for (let i = 0; i < tableRows.length; i++) {
-    const currentTableRowNode = tableRows[i];
+  for (let r = 0; r < tableRows.length; r++) {
+    const currentTableRowNode = tableRows[r];
 
     if ($isTableRowNode(currentTableRowNode)) {
-      for (let j = 0; j < columnCount; j++) {
-        const newTableCell = $createTableCellNode(i === 0);
+      for (let c = 0; c < columnCount; c++) {
+        let headerState = TableCellHeaderStates.NO_STATUS;
+
+        if (r === 0) {
+          headerState |= TableCellHeaderStates.ROW;
+        }
+
+        const newTableCell = $createTableCellNode(headerState);
         newTableCell.append(lexical.$createParagraphNode());
         const tableRowChildren = currentTableRowNode.getChildren();
 
@@ -920,12 +1129,12 @@ function $deleteTableColumn(tableNode, targetIndex) {
   return tableNode;
 }
 
-exports.$applyCustomTableHandlers = $applyCustomTableHandlers;
 exports.$createTableCellNode = $createTableCellNode;
 exports.$createTableNode = $createTableNode;
 exports.$createTableNodeWithDimensions = $createTableNodeWithDimensions;
 exports.$createTableRowNode = $createTableRowNode;
 exports.$deleteTableColumn = $deleteTableColumn;
+exports.$getElementGridForTableNode = $getElementGridForTableNode;
 exports.$getTableCellNodeFromLexicalNode = $getTableCellNodeFromLexicalNode;
 exports.$getTableColumnIndexFromTableCellNode = $getTableColumnIndexFromTableCellNode;
 exports.$getTableNodeFromLexicalNodeOrThrow = $getTableNodeFromLexicalNodeOrThrow;
@@ -937,6 +1146,11 @@ exports.$isTableCellNode = $isTableCellNode;
 exports.$isTableNode = $isTableNode;
 exports.$isTableRowNode = $isTableRowNode;
 exports.$removeTableRowAtIndex = $removeTableRowAtIndex;
+exports.TableCellHeaderStates = TableCellHeaderStates;
 exports.TableCellNode = TableCellNode;
 exports.TableNode = TableNode;
 exports.TableRowNode = TableRowNode;
+exports.TableSelection = TableSelection;
+exports.applyTableHandlers = applyTableHandlers;
+exports.getCellFromTarget = getCellFromTarget;
+exports.getTableSelectionFromTableElement = getTableSelectionFromTableElement;