improvement: add move cell left/right shortcuts of multi-column, vari…

…ous perf improvements (#3373)

Add two new hotkeys for moving left/right when in multi-column.
`Mod+Shift+7`, `Mod+Shift+8`.

Also performance improvements in 3 places that would re-render all cells
when moving only one cell

docs/guides/editor_features/hotkeys.md

@@ -20,8 +20,10 @@ You can find a list of available hotkeys below:
 | `cell.format`               |
 | `cell.goToDefinition`       |
 | `cell.hideCode`             |
-| `cell.moveDown`             |
 | `cell.moveUp`               |
+| `cell.moveDown`             |
+| `cell.moveLeft`             |
+| `cell.moveRight`            |
 | `cell.redo`                 |
 | `cell.run`                  |
 | `cell.runAndNewAbove`       |

frontend/src/components/editor/Cell.tsx

@@ -439,6 +439,10 @@ const CellComponent = (
     "cell.createBelow": createBelow,
     "cell.moveUp": () => moveCell({ cellId, before: true }),
     "cell.moveDown": () => moveCell({ cellId, before: false }),
+    "cell.moveLeft": () =>
+      canMoveX ? moveCell({ cellId, direction: "left" }) : undefined,
+    "cell.moveRight": () =>
+      canMoveX ? moveCell({ cellId, direction: "right" }) : undefined,
     "cell.hideCode": () => {
       const nextHideCode = !cellConfig.hide_code;
       // Fire-and-forget

frontend/src/components/editor/SortableCell.tsx

@@ -1,14 +1,17 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import React, { memo, useContext } from "react";
 import { useSortable } from "@dnd-kit/sortable";
-import { CSS } from "@dnd-kit/utilities";
+import { CSS, type Transform } from "@dnd-kit/utilities";
 import { GripVerticalIcon } from "lucide-react";
 import type { CellId } from "@/core/cells/ids";
 import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 import { mergeRefs } from "@/utils/mergeRefs";
+import type { SyntheticListenerMap } from "@dnd-kit/core/dist/hooks/utilities";
+import type { DraggableAttributes } from "@dnd-kit/core";
 
 interface Props extends React.HTMLAttributes<HTMLDivElement> {
+  children: React.ReactNode;
   cellId: CellId;
   canMoveX?: boolean;
 }
@@ -24,12 +27,25 @@ export const CellDragHandle: React.FC = memo(() => {
 });
 CellDragHandle.displayName = "DragHandle";
 
-const SortableCellInternal = React.forwardRef(
+function isTransformNoop(transform: Transform | null) {
+  if (!transform) {
+    return true;
+  }
+  return (
+    transform.x === 0 &&
+    transform.y === 0 &&
+    transform.scaleX === 1 &&
+    transform.scaleY === 1
+  );
+}
+
+export const SortableCell = React.forwardRef(
   (
     { cellId, canMoveX, ...props }: Props,
     ref: React.ForwardedRef<HTMLDivElement>,
   ) => {
-    // Sort
+    // This hook re-renders every time _any_ cell is dragged,
+    // so we should avoid any expensive operations in this component
     const {
       attributes,
       listeners,
@@ -39,6 +55,58 @@ const SortableCellInternal = React.forwardRef(
       isDragging,
     } = useSortable({ id: cellId.toString() });
 
+    // Perf:
+    // If the transform is a noop, keep it as null
+    const transformOrNull = isTransformNoop(transform) ? null : transform;
+    // If there is no transform, we don't need a transition
+    const transitionOrUndefined =
+      transformOrNull == null ? undefined : transition;
+
+    // Use a new component to avoid re-rendering when the cell is dragged
+    return (
+      <SortableCellInternal
+        ref={ref}
+        cellId={cellId}
+        canMoveX={canMoveX}
+        {...props}
+        attributes={attributes}
+        listeners={listeners}
+        setNodeRef={setNodeRef}
+        transform={transformOrNull}
+        transition={transitionOrUndefined}
+        isDragging={isDragging}
+      />
+    );
+  },
+);
+SortableCell.displayName = "SortableCell";
+
+interface SortableCellInternalProps extends Props {
+  children: React.ReactNode;
+  attributes: DraggableAttributes;
+  listeners: SyntheticListenerMap | undefined;
+  setNodeRef: (node: HTMLElement | null) => void;
+  transform: Transform | null;
+  transition: string | undefined;
+  isDragging: boolean;
+}
+
+const SortableCellInternal = React.forwardRef(
+  (
+    {
+      cellId,
+      canMoveX,
+      children,
+      attributes,
+      listeners,
+      setNodeRef,
+      transform,
+      transition,
+      isDragging,
+      ...props
+    }: SortableCellInternalProps,
+    ref: React.ForwardedRef<HTMLDivElement>,
+  ) => {
     const style: React.CSSProperties = {
       transform: transform
         ? CSS.Transform.toString({
@@ -83,12 +151,10 @@ const SortableCellInternal = React.forwardRef(
         style={style}
       >
         <DragHandleSlot.Provider value={dragHandle}>
-          {props.children}
+          {children}
         </DragHandleSlot.Provider>
       </div>
     );
   },
 );
-SortableCellInternal.displayName = "SortableCell";
-
-export const SortableCell = memo(SortableCellInternal);
+SortableCellInternal.displayName = "SortableCellInternal";

frontend/src/components/editor/actions/useCellActionButton.tsx

@@ -3,11 +3,7 @@ import { downloadCellOutput } from "@/components/export/export-output-button";
 import { Switch } from "@/components/ui/switch";
 import { formatEditorViews } from "@/core/codemirror/format";
 import { toggleToLanguage } from "@/core/codemirror/language/commands";
-import {
-  hasOnlyOneCellAtom,
-  useCellActions,
-  useCellIds,
-} from "@/core/cells/cells";
+import { hasOnlyOneCellAtom, useCellActions } from "@/core/cells/cells";
 import {
   ImageIcon,
   Code2Icon,
@@ -27,6 +23,8 @@ import {
   DatabaseIcon,
   Columns2Icon,
   XCircleIcon,
+  ChevronLeftIcon,
+  ChevronRightIcon,
 } from "lucide-react";
 import type { ActionButton } from "./types";
 import { MultiIcon } from "@/components/icons/multi-icon";
@@ -89,7 +87,6 @@ export function useCellActionButtons({ cell }: Props) {
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
   const aiEnabled = useAtomValue(aiEnabledAtom);
   const autoInstantiate = useAtomValue(autoInstantiateAtom);
-  const cellIds = useCellIds();
   const kioskMode = useAtomValue(kioskModeAtom);
   const appWidth = useAtomValue(appWidthAtom);
 
@@ -98,7 +95,6 @@ export function useCellActionButtons({ cell }: Props) {
   }
 
   const { cellId, config, getEditorView, name, hasOutput, status } = cell;
-  const cellIdx = cellIds.inOrderIds.indexOf(cellId);
 
   const toggleDisabled = async () => {
     const newConfig = { disabled: !config.disabled };
@@ -142,7 +138,7 @@ export function useCellActionButtons({ cell }: Props) {
               <div className="flex items-center justify-between">
                 <Label htmlFor="cell-name">Cell name</Label>
                 <NameCellInput
-                  placeholder={`cell_${cellIdx}`}
+                  placeholder={"cell name"}
                   value={name}
                   onKeyDown={(e) => {
                     if (e.key === "Enter") {
@@ -161,7 +157,7 @@ export function useCellActionButtons({ cell }: Props) {
         },
         rightElement: (
           <NameCellInput
-            placeholder={`cell_${cellIdx}`}
+            placeholder={"cell name"}
             value={name}
             onChange={(newName) => updateCellName({ cellId, name: newName })}
           />
@@ -321,6 +317,20 @@ export function useCellActionButtons({ cell }: Props) {
         hotkey: "cell.moveDown",
         handle: () => moveCell({ cellId, before: false }),
       },
+      {
+        icon: <ChevronLeftIcon size={13} strokeWidth={1.5} />,
+        label: "Move cell left",
+        hotkey: "cell.moveLeft",
+        handle: () => moveCell({ cellId, direction: "left" }),
+        hidden: appWidth !== "columns",
+      },
+      {
+        icon: <ChevronRightIcon size={13} strokeWidth={1.5} />,
+        label: "Move cell right",
+        hotkey: "cell.moveRight",
+        handle: () => moveCell({ cellId, direction: "right" }),
+        hidden: appWidth !== "columns",
+      },
       {
         icon: <ChevronsUpIcon size={13} strokeWidth={1.5} />,
         label: "Send to top",

frontend/src/components/editor/cell/code/cell-editor.tsx

@@ -7,7 +7,11 @@ import React, { memo, useCallback, useEffect, useRef, useMemo } from "react";
 import { setupCodeMirror } from "@/core/codemirror/cm";
 import { getFeatureFlag } from "@/core/config/feature-flag";
 import useEvent from "react-use-event-hook";
-import { type CellActions, useCellActions } from "@/core/cells/cells";
+import {
+  type CellActions,
+  notebookAtom,
+  useCellActions,
+} from "@/core/cells/cells";
 import type { CellRuntimeState, CellData } from "@/core/cells/types";
 import type { UserConfig } from "@/core/config/config-schema";
 import type { Theme } from "@/theme/useTheme";
@@ -35,6 +39,7 @@ import { invariant } from "@/utils/invariant";
 import { connectionAtom } from "@/core/network/connection";
 import { WebSocketState } from "@/core/websocket/types";
 import { realTimeCollaboration } from "@/core/codemirror/rtc/extension";
+import { store } from "@/core/state/jotai";
 
 export interface CellEditorProps
   extends Pick<CellRuntimeState, "status">,
@@ -361,7 +366,12 @@ const CellEditorInternal = ({
   const shouldFocus =
     editorViewRef.current === null || serializedEditorState !== null;
   useEffect(() => {
-    if (shouldFocus && allowFocus) {
+    // Perf:
+    // We don't pass this in from the props since it causes lots of re-renders for unrelated cells
+    const hasNotebookKey = store.get(notebookAtom).scrollKey !== null;
+
+    // Only focus if the notebook does not currently have a scrollKey (which means we are focusing on another cell)
+    if (shouldFocus && allowFocus && !hasNotebookKey) {
       // Focus and scroll into view; request an animation frame to
       // avoid a race condition when new editors are created
       // very rapidly by holding a hotkey

frontend/src/components/editor/cell/code/language-toggle.tsx

@@ -107,6 +107,14 @@ export const LanguageToggle: React.FC<Props> = ({
   displayName,
   onAfterToggle,
 }) => {
+  const handleClick = () => {
+    if (!editorView) {
+      return;
+    }
+    switchLanguage(editorView, toType);
+    onAfterToggle();
+  };
+
   if (!canSwitchToLanguage) {
     return null;
   }
@@ -122,13 +130,7 @@ export const LanguageToggle: React.FC<Props> = ({
         variant="text"
         size="xs"
         className="opacity-80 px-1"
-        onClick={() => {
-          if (!editorView) {
-            return;
-          }
-          switchLanguage(editorView, toType);
-          onAfterToggle();
-        }}
+        onClick={handleClick}
       >
         {icon}
       </Button>

frontend/src/components/editor/renderers/CellArray.tsx

@@ -214,7 +214,7 @@ const SortableColumn: React.FC<{
               key={cellData.id.toString()}
               theme={theme}
               showPlaceholder={hasOnlyOneCell}
-              allowFocus={!invisible && !notebook.scrollKey}
+              allowFocus={!invisible}
               id={cellData.id}
               code={cellData.code}
               outline={cellRuntime.outline}

frontend/src/core/cells/__tests__/cells.test.ts

@@ -321,6 +321,80 @@ describe("cell reducer", () => {
       [1] ''
       "
     `);
+
+    // Add a column breakpoint to test left/right movement
+    actions.addColumnBreakpoint({ cellId: "1" as CellId });
+    expect(formatCells(state)).toMatchInlineSnapshot(`
+      "
+      > col 0
+      [0] ''
+
+      > col 1
+      [1] ''
+      "
+    `);
+
+    // Move cell right
+    actions.moveCell({
+      cellId: firstCellId,
+      direction: "right",
+    });
+    expect(formatCells(state)).toMatchInlineSnapshot(`
+      "
+      > col 0
+
+
+      > col 1
+      [0] ''
+
+      [1] ''
+      "
+    `);
+
+    // Move cell left
+    actions.moveCell({
+      cellId: firstCellId,
+      direction: "left",
+    });
+    expect(formatCells(state)).toMatchInlineSnapshot(`
+      "
+      > col 0
+      [0] ''
+
+      > col 1
+      [1] ''
+      "
+    `);
+
+    // Try to move cell left when it's already in leftmost column (should noop)
+    actions.moveCell({
+      cellId: firstCellId,
+      direction: "left",
+    });
+    expect(formatCells(state)).toMatchInlineSnapshot(`
+      "
+      > col 0
+      [0] ''
+
+      > col 1
+      [1] ''
+      "
+    `);
+
+    // Try to move cell right when it's already in rightmost column (should noop)
+    actions.moveCell({
+      cellId: "1" as CellId,
+      direction: "right",
+    });
+    expect(formatCells(state)).toMatchInlineSnapshot(`
+      "
+      > col 0
+      [0] ''
+
+      > col 1
+      [1] ''
+      "
+    `);
   });
 
   it("can drag and drop a cell", () => {