@marimo-team/islands 0.25.3-dev197 → 0.25.3-dev198

package/src/core/cells/arrange.ts

@@ -1,26 +1,34 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-/** Notebook operations that move cells within and across columns. */
+/** Notebook operations that move cells and arrange columns. */
 
 import { dequal as isEqual } from "dequal";
 import { atom, type Getter, type Setter } from "jotai";
-import type { CellColumnId, CellIndex } from "@/utils/id-tree";
+import { reorderColumnSizes } from "@/components/editor/columns/storage";
+import type {
+  DocumentEditor,
+  EditableDocument,
+} from "@/core/notebook/document";
 import { arrayMove } from "@/utils/arrays";
 import { notebookAtom, type UndoMoveEntry } from "./cells";
 import { documentOf, focusWhenMounted, modelOf } from "./edit";
-import type { CellId } from "./ids";
+import { CellId } from "./ids";
 import {
+  type CellColumnId,
+  columnIdsAtom,
   columnOf,
+  emptyColumnIdsAtom,
   findColumn,
   type Layout,
   type LayoutColumn,
   layoutAtom,
+  newColumnId,
 } from "./layout";
 
 /** A placement to restore a cell to: its column and its index there. */
 export interface Placement {
   columnId: CellColumnId;
-  index: CellIndex;
+  index: number;
 }
 
 /** A cell with the cells its fold hides, which move with it. */
@@ -66,12 +74,11 @@ function columnConfigs(
   return fixes;
 }
 
-/**
- * Make the document's cells fall into `columns`, in that order, as one edit.
- * Nothing is written when they already do.
- */
-export function arrange(get: Getter, set: Setter, columns: CellId[][]): void {
-  const document = documentOf(get);
+/** The steps that make the document's cells fall into `columns`, or null when they do. */
+function arrangement(
+  document: EditableDocument,
+  columns: readonly (readonly CellId[])[],
+): ((tx: DocumentEditor) => void) | null {
   const order = columns.flat();
   const fixes = columnConfigs(
     columns,
@@ -79,16 +86,61 @@ export function arrange(get: Getter, set: Setter, columns: CellId[][]): void {
   );
   const reordered = !isEqual(order, document.order);
   if (!reordered && fixes.length === 0) {
-    return;
+    return null;
   }
-  set(modelOf(get).editAtom, (tx) => {
+  return (tx) => {
     if (reordered) {
       tx.setOrder(order);
     }
     for (const [id, column] of fixes) {
       tx.updateConfig(id, { column });
     }
-  });
+  };
+}
+
+/**
+ * How many of `columns` the document holds: all but the trailing empty
+ * ones, which the page alone keeps.
+ */
+function heldByDocument(columns: readonly (readonly CellId[])[]): number {
+  let count = columns.length;
+  while (count > 1 && columns[count - 1].length === 0) {
+    count -= 1;
+  }
+  return count;
+}
+
+/** Name the page's columns: `count` of `ids` for the document's, the rest kept empty. */
+function nameColumns(
+  set: Setter,
+  ids: readonly CellColumnId[],
+  count: number,
+): void {
+  set(columnIdsAtom, ids.slice(0, count));
+  set(emptyColumnIdsAtom, ids.slice(count));
+}
+
+/**
+ * Make the document's cells fall into `columns`, in that order, as one
+ * edit, and name the columns `ids`, by default as they stand. Nothing is
+ * written when the cells already fall that way. Trailing columns left
+ * empty stay on the page, as the document cannot hold them, until a drop
+ * compacts them.
+ */
+export function arrange(
+  get: Getter,
+  set: Setter,
+  columns: CellId[][],
+  ids: readonly CellColumnId[] = get(layoutAtom).columns.map(
+    (column) => column.id,
+  ),
+): void {
+  const count = heldByDocument(columns);
+  nameColumns(set, ids, count);
+  const edit = arrangement(documentOf(get), columns.slice(0, count));
+  if (edit !== null) {
+    set(modelOf(get).editAtom, edit);
+  }
 }
 
 /**
@@ -339,3 +391,144 @@ export function placeCells(
   }
   arrange(get, set, cleaned);
 }
+
+// Columns --------------------------------------------------------------------
+
+/** Add a column after another, with one new cell to start it, and focus the cell. */
+export const addColumnAtom = atom(
+  null,
+  (get, set, { columnId }: { columnId: CellColumnId }): void => {
+    const layout = get(layoutAtom);
+    const at = layout.columns.findIndex((column) => column.id === columnId);
+    if (at === -1) {
+      return;
+    }
+    const id = CellId.create();
+    const columns = columnsOf(layout);
+    columns.splice(at + 1, 0, [id]);
+    const ids = layout.columns.map((column) => column.id);
+    ids.splice(at + 1, 0, newColumnId());
+    const document = documentOf(get);
+    const previous = layout.columns[at].ids.at(-1);
+    const count = heldByDocument(columns);
+    nameColumns(set, ids, count);
+    set(modelOf(get).editAtom, (tx) => {
+      tx.createCell(
+        { id, code: "", config: { column: at + 1 } },
+        previous === undefined ? {} : { after: previous },
+      );
+      arrangement(document, columns.slice(0, count))?.(tx);
+    });
+    focusWhenMounted(set, id);
+  },
+);
+
+/** Start a new column at a cell, unless it already starts the first column. */
+export const addColumnBreakpointAtom = atom(
+  null,
+  (get, set, { cellId }: { cellId: CellId }): void => {
+    const layout = get(layoutAtom);
+    const column = columnOf(layout, cellId);
+    const at = layout.columns.indexOf(column);
+    const index = column.ids.indexOf(cellId);
+    if (at === 0 && index === 0) {
+      return;
+    }
+    const columns = columnsOf(layout);
+    columns.splice(at, 1, column.ids.slice(0, index), column.ids.slice(index));
+    const ids = layout.columns.map((column) => column.id);
+    // The id stays with the cells, so the half without any is the new one.
+    ids.splice(index === 0 ? at : at + 1, 0, newColumnId());
+    arrange(get, set, columns, ids);
+  },
+);
+
+/** Remove a column, moving its cells to the column before it, or after it for the first. */
+export const deleteColumnAtom = atom(
+  null,
+  (get, set, { columnId }: { columnId: CellColumnId }): void => {
+    const layout = get(layoutAtom);
+    const at = layout.columns.findIndex((column) => column.id === columnId);
+    if (at === -1 || layout.columns.length <= 1) {
+      return;
+    }
+    const columns = columnsOf(layout);
+    const target = at === 0 ? 1 : at - 1;
+    columns[target] =
+      at === 0
+        ? [...columns[at], ...columns[target]]
+        : [...columns[target], ...columns[at]];
+    columns.splice(at, 1);
+    const ids = layout.columns.map((column) => column.id);
+    ids.splice(at, 1);
+    arrange(get, set, columns, ids);
+  },
+);
+
+/** Bring every cell into the first column. */
+export const mergeAllColumnsAtom = atom(null, (get, set): void => {
+  const layout = get(layoutAtom);
+  if (layout.columns.length <= 1) {
+    return;
+  }
+  arrange(
+    get,
+    set,
+    [layout.columns.flatMap((column) => [...column.ids])],
+    [layout.columns[0].id],
+  );
+});
+
+/** Drop the columns a drag left empty. */
+export const compactColumnsAtom = atom(null, (get, set): void => {
+  const layout = get(layoutAtom);
+  const kept = layout.columns.filter((column) => column.ids.length > 0);
+  if (kept.length === layout.columns.length || layout.columns.length <= 1) {
+    return;
+  }
+  arrange(
+    get,
+    set,
+    kept.map((column) => [...column.ids]),
+    kept.map((column) => column.id),
+  );
+});
+
+/** Move a column where another is, or one step left or right, with its saved width. */
+export const moveColumnAtom = atom(
+  null,
+  (
+    get,
+    set,
+    {
+      column,
+      overColumn,
+    }: {
+      column: CellColumnId;
+      overColumn: CellColumnId | "_left_" | "_right_";
+    },
+  ): void => {
+    if (column === overColumn) {
+      return;
+    }
+    const layout = get(layoutAtom);
+    const ids = layout.columns.map((column) => column.id);
+    const from = ids.indexOf(column);
+    const to =
+      overColumn === "_left_"
+        ? from - 1
+        : overColumn === "_right_"
+          ? from + 1
+          : ids.indexOf(overColumn);
+    if (from === -1 || to < 0 || to >= ids.length) {
+      return;
+    }
+    arrange(
+      get,
+      set,
+      arrayMove(columnsOf(layout), from, to),
+      arrayMove(ids, from, to),
+    );
+    reorderColumnSizes(from, to);
+  },
+);

package/src/core/cells/cells.ts

@@ -6,11 +6,7 @@ import { atomFamily, selectAtom, splitAtom } from "jotai/utils";
 import { createRef, type ReducerWithoutAction } from "react";
 import type { CellHandle } from "@/components/editor/notebook-cell";
 import type { EditableDocument, Placement } from "@/core/notebook/document";
-import {
-  type CellColumnId,
-  type CellIndex,
-  MultiColumn,
-} from "@/utils/id-tree";
+import type { CellColumnId } from "./layout";
 import { invariant } from "@/utils/invariant";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
@@ -18,8 +14,6 @@ import { extractAllTracebackInfo, type TracebackInfo } from "@/utils/traceback";
 import { createReducerAndAtoms } from "../../utils/createReducer";
 import { foldAllBulk, unfoldAllBulk } from "../codemirror/editing/commands";
 import { editorMountScheduler } from "../codemirror/editor-mount-scheduler";
-import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
-import { splitEditor } from "../codemirror/language/utils";
 import type { SerializedEditorState } from "../codemirror/types";
 import { mergeOutlines } from "../dom/outline";
 import type { CellMessage } from "../kernel/messages";
@@ -30,7 +24,7 @@ import { isWasm } from "../wasm/utils";
 import { prepareCellForExecution, transitionCell } from "./cell";
 import { applyDocument, documentEditMiddleware } from "./document";
 import { isCodeHidden, shownHiddenCodeAtom } from "./editor";
-import { CellId, SCRATCH_CELL_ID } from "./ids";
+import { type CellId, SCRATCH_CELL_ID } from "./ids";
 import { type CellLog, getCellLogsForMessage } from "./logs";
 import type { Outline } from "./outline";
 import { focusAndScrollCellIntoView } from "./scrollCellIntoView";
@@ -74,7 +68,7 @@ export interface UndoDeleteEntry {
 export interface UndoMoveEntry {
   type: "move";
   cellIds: CellId[];
-  placements: Array<{ columnId: CellColumnId; index: CellIndex }>;
+  placements: Array<{ columnId: CellColumnId; index: number }>;
 }
 
 export type HistoryEntry = UndoDeleteEntry | UndoMoveEntry;
@@ -86,7 +80,8 @@ export interface NotebookState {
   /**
    * Order of cells on the page.
    */
-  cellIds: MultiColumn<CellId>;
+  /** The cells in document order. */
+  cellIds: CellId[];
   /**
    * Map of cells to their view state
    */
@@ -144,7 +139,7 @@ function withScratchCell(notebookState: NotebookState): NotebookState {
  */
 export function initialNotebookState(): NotebookState {
   return withScratchCell({
-    cellIds: MultiColumn.from([]),
+    cellIds: [],
     cellData: {},
     cellRuntime: {},
     cellHandles: {},
@@ -164,74 +159,6 @@ const {
   useActions,
   valueAtom: notebookAtom,
 } = createReducerAndAtoms(initialNotebookState, {
-  moveColumn: (
-    state,
-    action: {
-      column: CellColumnId;
-      overColumn: CellColumnId | "_left_" | "_right_";
-    },
-  ) => {
-    if (action.column === action.overColumn) {
-      return state;
-    }
-    return {
-      ...state,
-      cellIds: state.cellIds.moveColumn(action.column, action.overColumn),
-    };
-  },
-  addColumn: (state, action: { columnId: CellColumnId }) => {
-    // Add column and new cell
-    const newCellId = CellId.create();
-    return {
-      ...state,
-      cellIds: state.cellIds.addColumn(action.columnId, [newCellId]),
-      cellData: {
-        ...state.cellData,
-        [newCellId]: createCell({
-          id: newCellId,
-        }),
-      },
-      cellRuntime: {
-        ...state.cellRuntime,
-        [newCellId]: createCellRuntimeState(),
-      },
-      cellHandles: {
-        ...state.cellHandles,
-        [newCellId]: createRef(),
-      },
-      scrollKey: newCellId,
-    };
-  },
-  addColumnBreakpoint: (state, action: { cellId: CellId }) => {
-    const { cellId } = action;
-    if (state.cellIds.getColumns()[0].inOrderIds[0] === cellId) {
-      return state;
-    }
-    return {
-      ...state,
-      cellIds: state.cellIds.insertBreakpoint(cellId),
-    };
-  },
-  deleteColumn: (state, action: { columnId: CellColumnId }) => {
-    // Move all cells in the column to the previous column
-    const { columnId } = action;
-    return {
-      ...state,
-      cellIds: state.cellIds.delete(columnId),
-    };
-  },
-  mergeAllColumns: (state) => {
-    return {
-      ...state,
-      cellIds: state.cellIds.mergeAllColumns(),
-    };
-  },
-  compactColumns: (state) => {
-    return {
-      ...state,
-      cellIds: state.cellIds.compact(),
-    };
-  },
   /**
    * Put moved cells back where they were.
    *
@@ -332,7 +259,7 @@ const {
   /** Follow the host document's cells and order. */
   reconcileDocument: (state, document: EditableDocument) => {
     const nextState = applyDocument(state, document);
-    for (const cellId of state.cellIds.inOrderIds) {
+    for (const cellId of state.cellIds) {
       if (!(cellId in nextState.cellData)) {
         releaseCellAtoms(cellId);
       }
@@ -382,7 +309,7 @@ const {
     },
   ) => {
     const cellData = { ...state.cellData };
-    for (const cellId of state.cellIds.inOrderIds) {
+    for (const cellId of state.cellIds) {
       const cell = state.cellData[cellId];
       const lastCodeRun = action.lastCodeRun[cellId] ?? null;
       cellData[cellId] = {
@@ -402,7 +329,7 @@ const {
   resetCellRuntimes: (state) => {
     const preserveSnapshot = isWasm();
     const cellRuntime = { ...state.cellRuntime };
-    for (const cellId of [...state.cellIds.inOrderIds, SCRATCH_CELL_ID]) {
+    for (const cellId of [...state.cellIds, SCRATCH_CELL_ID]) {
       const previous = state.cellRuntime[cellId];
       const hasSnapshot =
         preserveSnapshot &&
@@ -419,12 +346,7 @@ const {
       return state;
     }
 
-    const column = state.cellIds.findWithId(scrollKey);
-    const index = column.indexOfOrThrow(scrollKey);
-
-    const cellId =
-      index === column.length - 1 ? column.last() : column.atOrThrow(index);
-
+    const cellId = scrollKey;
     focusAndScrollCellIntoView({
       cellId: cellId,
       cell: state.cellHandles[cellId],
@@ -462,110 +384,6 @@ const {
       cellLogs: [],
     };
   },
-  splitCell: (state, action: { cellId: CellId }) => {
-    const { cellId } = action;
-    const column = state.cellIds.findWithId(cellId);
-    const index = column.indexOfOrThrow(cellId);
-    const cell = state.cellData[cellId];
-    const cellHandle = state.cellHandles[cellId].current;
-
-    if (cellHandle?.editorView == null) {
-      return state;
-    }
-
-    const { beforeCursorCode, afterCursorCode } = splitEditor(
-      cellHandle.editorView,
-    );
-
-    const update = prepareCellCodeUpdate(
-      cellHandle.editorView.state,
-      beforeCursorCode,
-    );
-
-    if (update) {
-      cellHandle.editorView.dispatch(update);
-    }
-
-    const newCellId = CellId.create();
-
-    return {
-      ...state,
-      cellIds: state.cellIds.insertId(newCellId, column.id, index + 1),
-      cellData: {
-        ...state.cellData,
-        [cellId]: {
-          ...cell,
-          code: beforeCursorCode,
-          edited:
-            Boolean(beforeCursorCode) &&
-            beforeCursorCode.trim() !== cell.lastCodeRun?.trim(),
-        },
-        [newCellId]: createCell({
-          id: newCellId,
-          code: afterCursorCode,
-          edited: Boolean(afterCursorCode),
-        }),
-      },
-      cellRuntime: {
-        ...state.cellRuntime,
-        [cellId]: {
-          ...state.cellRuntime[cellId],
-          output: null,
-          consoleOutputs: [],
-        },
-        [newCellId]: createCellRuntimeState(),
-      },
-      cellHandles: {
-        ...state.cellHandles,
-        [newCellId]: createRef(),
-      },
-      scrollKey: newCellId,
-    };
-  },
-  undoSplitCell: (state, action: { cellId: CellId; snapshot: string }) => {
-    const { cellId, snapshot } = action;
-
-    const cell = state.cellData[cellId];
-    const cellHandle = state.cellHandles[cellId].current;
-
-    if (cellHandle?.editorView == null) {
-      return state;
-    }
-
-    const update = prepareCellCodeUpdate(cellHandle.editorView.state, snapshot);
-
-    if (update) {
-      cellHandle.editorView.dispatch(update);
-    }
-
-    return {
-      ...state,
-      cellIds: state.cellIds.transformWithCellId(cellId, (column) => {
-        const newCellIndex = column.indexOfOrThrow(cellId) + 1;
-        return column.deleteAtIndex(newCellIndex);
-      }),
-      cellData: {
-        ...state.cellData,
-        [cellId]: {
-          ...cell,
-          code: snapshot,
-          edited:
-            Boolean(snapshot) && snapshot?.trim() !== cell.lastCodeRun?.trim(),
-        },
-      },
-      cellRuntime: {
-        ...state.cellRuntime,
-        [cellId]: {
-          ...state.cellRuntime[cellId],
-          output: null,
-          consoleOutputs: [],
-        },
-      },
-      cellHandles: {
-        ...state.cellHandles,
-      },
-    };
-  },
   clearCellOutput: (state, action: { cellId: CellId }) => {
     const { cellId } = action;
     return updateCellRuntimeState({
@@ -594,7 +412,7 @@ const {
   },
   clearAllCellOutputs: (state) => {
     const newCellRuntime = { ...state.cellRuntime };
-    for (const cellId of state.cellIds.inOrderIds) {
+    for (const cellId of state.cellIds) {
       newCellRuntime[cellId] = {
         ...newCellRuntime[cellId],
         output: null,
@@ -660,43 +478,12 @@ function updateCellData({
   };
 }
 
+/** Each cell's config in order. The document's `column` configs are canonical. */
 export function getCellConfigs(state: NotebookState): CellConfig[] {
-  const cells = state.cellData;
-
-  // We set to null by default to prevent inconsistencies between undefined & null
-  const defaultCellConfig: Partial<CellConfig> = { column: null };
-
-  // Handle the case where there's only one column
-  // We don't want to set the column config
-  const hasMultipleColumns = state.cellIds.getColumns().length > 1;
-  if (!hasMultipleColumns) {
-    return state.cellIds.getColumns().flatMap((column) => {
-      return column.inOrderIds.map((cellId) => {
-        return {
-          ...cells[cellId].config,
-          ...defaultCellConfig,
-        };
-      });
-    });
-  }
-
-  return state.cellIds.getColumns().flatMap((column, columnIndex) => {
-    return column.inOrderIds.map((cellId, cellIndex) => {
-      const config: Partial<CellConfig> = { ...defaultCellConfig };
-
-      // Only set the column index for the first cell in the column
-      if (cellIndex === 0) {
-        config.column = columnIndex;
-      }
-
-      const newConfig = {
-        ...cells[cellId].config,
-        ...config,
-      };
-
-      return newConfig;
-    });
-  });
+  return state.cellIds.map((cellId) => ({
+    ...state.cellData[cellId].config,
+    column: state.cellData[cellId].config.column ?? null,
+  }));
 }
 
 export {
@@ -716,9 +503,7 @@ export const addLogs = imperativeActions.addLogs;
 
 export const cellIdsAtom = atom((get) => get(notebookAtom).cellIds);
 
-export const hasOnlyOneCellAtom = atom((get) =>
-  get(cellIdsAtom).hasOnlyOneId(),
-);
+export const hasOnlyOneCellAtom = atom((get) => get(cellIdsAtom).length === 1);
 
 export const hasDisabledCellsAtom = atom(
   (get) => disabledCellIds(get(notebookAtom)).length > 0,
@@ -736,7 +521,7 @@ export const needsRunAtom = atom((get) => notebookNeedsRun(get(notebookAtom)));
 
 export const cellErrorsAtom = atom((get) => {
   const { cellIds, cellRuntime, cellData } = get(notebookAtom);
-  const errors = cellIds.inOrderIds
+  const errors = cellIds
     .map((cellId) => {
       const cell = cellRuntime[cellId];
       const { name } = cellData[cellId];
@@ -785,9 +570,7 @@ export const cellOutlinesAtom = createDeepEqualAtom(
 export const notebookOutline = atom((get) => {
   const cellIds = get(cellIdsAtom);
   const cellRuntime = get(cellsRuntimeAtom);
-  const outlines = cellIds.inOrderIds.map(
-    (cellId) => cellRuntime[cellId].outline,
-  );
+  const outlines = cellIds.map((cellId) => cellRuntime[cellId].outline);
   return mergeOutlines(outlines);
 });
 
@@ -797,7 +580,7 @@ export const cellIdToNamesMap = createDeepEqualAtom(
   atom((get) => {
     const { cellIds, cellData } = get(notebookAtom);
     const names: Record<CellId, string | undefined> = Objects.fromEntries(
-      cellIds.inOrderIds.map((cellId) => [cellId, cellData[cellId]?.name]),
+      cellIds.map((cellId) => [cellId, cellData[cellId]?.name]),
     );
     return names;
   }),
@@ -854,12 +637,12 @@ export const getCells = () => store.get(notebookAtom).cellIds;
  */
 export const getCellNames = () => {
   const { cellIds, cellData } = store.get(notebookAtom);
-  return cellIds.inOrderIds.map((id) => cellData[id]?.name).filter(Boolean);
+  return cellIds.map((id) => cellData[id]?.name).filter(Boolean);
 };
 
 const cellDataAtoms = splitAtom(
   selectAtom(notebookAtom, (cells) =>
-    cells.cellIds.inOrderIds.map((id) => cells.cellData[id]),
+    cells.cellIds.map((id) => cells.cellData[id]),
   ),
 );
 export const useCellDataAtoms = () => useAtom(cellDataAtoms);
@@ -879,16 +662,7 @@ export const notebookQueuedOrRunningCountAtom = atom((get) =>
   notebookQueueOrRunningCount(get(notebookAtom)),
 );
 
-export const numColumnsAtom = atom(
-  (get) => get(notebookAtom).cellIds.colLength,
-);
-export const hasCellsAtom = atom(
-  (get) => get(notebookAtom).cellIds.idLength > 0,
-);
-export const columnIdsAtom = atom((get) =>
-  get(notebookAtom).cellIds.getColumnIds(),
-);
-
+export const hasCellsAtom = atom((get) => get(notebookAtom).cellIds.length > 0);
 export const cellDataAtom = atomFamily((cellId: CellId) =>
   atom((get) => get(notebookAtom).cellData[cellId]),
 );
@@ -929,7 +703,7 @@ export const useCellHandle = (cellId: CellId) =>
  */
 export const getAllEditorViews = () => {
   const { cellIds, cellHandles } = store.get(notebookAtom);
-  return cellIds.inOrderIds
+  return cellIds
     .map((cellId) => cellHandles[cellId]?.current?.editorViewOrNull)
     .filter(Boolean);
 };
@@ -953,12 +727,10 @@ export function flattenTopLevelNotebookCells(
   state: NotebookState,
 ): (CellData & CellRuntimeState)[] {
   const { cellIds, cellData, cellRuntime } = state;
-  return cellIds.getColumns().flatMap((column) =>
-    column.topLevelIds.map((cellId) => ({
-      ...cellData[cellId],
-      ...cellRuntime[cellId],
-    })),
-  );
+  return cellIds.map((cellId) => ({
+    ...cellData[cellId],
+    ...cellRuntime[cellId],
+  }));
 }
 
 export function createTracebackInfoAtom(