Enter in table cell adds a row after the current one (#4186)

Co-authored-by: Tom Moor <tom.moor@gmail.com>

shared/editor/nodes/Table.ts

@@ -12,15 +12,18 @@ import {
   toggleHeaderCell,
   toggleHeaderColumn,
   toggleHeaderRow,
+  CellSelection,
 } from "prosemirror-tables";
 import {
   addRowAt,
   createTable,
   getCellsInColumn,
   moveRow,
+  setTextSelection,
 } from "prosemirror-utils";
 import { Decoration, DecorationSet } from "prosemirror-view";
 import { MarkdownSerializerState } from "../lib/markdown/serializer";
+import { getRowIndexFromText } from "../queries/getRowIndex";
 import tablesRule from "../rules/tables";
 import { Dispatch } from "../types";
 import Node from "./Node";
@@ -123,12 +126,23 @@ export default class Table extends Node {
         if (!isInTable(state)) {
           return false;
         }
+        const index = getRowIndexFromText(
+          (state.selection as unknown) as CellSelection
+        );
 
-        // TODO: Adding row at the end for now, can we find the current cell
-        // row index and add the row below that?
-        const cells = getCellsInColumn(0)(state.selection) || [];
+        if (index === 0) {
+          const cells = getCellsInColumn(0)(state.selection);
+          if (!cells) {
+            return false;
+          }
 
-        dispatch(addRowAt(cells.length, true)(state.tr));
+          const tr = addRowAt(index + 2, true)(state.tr);
+          dispatch(
+            setTextSelection(cells[1].pos)(moveRow(index + 2, index + 1)(tr))
+          );
+        } else {
+          dispatch(addRowAt(index + 1, true)(state.tr));
+        }
         return true;
       },
     };

shared/editor/queries/getRowIndex.ts

@@ -9,3 +9,13 @@ export default function getRowIndex(selection: CellSelection) {
   const path = (selection.$from as any).path;
   return path[path.length - 8];
 }
+
+export function getRowIndexFromText(selection: CellSelection) {
+  const isRowSelection = selection.isRowSelection && selection.isRowSelection();
+  const path = (selection.$from as any).path;
+  if (isRowSelection) {
+    return path[path.length - 8];
+  } else {
+    return path[path.length - 11];
+  }
+}