@marimo-team/islands 0.25.3-dev28 → 0.25.3-dev284

package/src/core/cells/add-missing-import.ts

@@ -3,7 +3,8 @@
 import { getRequestClient } from "../network/requests";
 import { store } from "../state/jotai";
 import { variablesAtom } from "../variables/state";
-import { type CellActions, notebookAtom } from "./cells";
+import { notebookAtom } from "./cells";
+import type { CreateCell, CreateCellOptions } from "./actions";
 import { CellId } from "./ids";
 
 /**
@@ -35,7 +36,7 @@ export function maybeAddMissingImport({
     `import[ \t]+${moduleName}[ \t]+as[ \t]+${variableName}`,
     "g",
   );
-  for (const cell of cellIds.inOrderIds) {
+  for (const cell of cellIds) {
     if (regex.test(cellData[cell].code)) {
       return false;
     }
@@ -50,82 +51,85 @@ export function maybeAddMissingImport({
 /**
  * Adds a marimo import to the notebook if not already present.
  * @param autoInstantiate Whether to automatically run the cell.
- * @param createNewCell The function to create a new cell.
- * @param fromCellId The cell to add the import to.
+ * @param createCell The command that creates a cell.
+ * @param fromCellId The cell to add the import next to.
  * @param before Whether to add the import before or after the cell.
  *
  * Returns the ID of the new cell if added, otherwise null.
  */
 export function maybeAddMarimoImport({
   autoInstantiate,
-  createNewCell,
+  createCell,
   fromCellId,
   before,
 }: {
   autoInstantiate: boolean;
-  createNewCell: CellActions["createNewCell"];
+  createCell: CreateCell;
   fromCellId?: CellId | null;
   before?: boolean;
 }): CellId | null {
-  const client = getRequestClient();
-  let newCellId: CellId | null = null;
-  const added = maybeAddMissingImport({
+  return addImport({
     moduleName: "marimo",
     variableName: "mo",
-    onAddImport: (importStatement) => {
-      newCellId = CellId.create();
-      createNewCell({
-        cellId: fromCellId ?? "__end__",
-        before: before ?? false,
-        code: importStatement,
-        lastCodeRun: autoInstantiate ? importStatement : undefined,
-        newCellId: newCellId,
-        skipIfCodeExists: true,
-        autoFocus: false,
-      });
-      if (autoInstantiate) {
-        void client.sendRun({
-          cellIds: [newCellId],
-          codes: [importStatement],
-        });
-      }
-    },
+    autoInstantiate,
+    createCell,
+    placement: before ? { before: fromCellId } : { after: fromCellId },
   });
-  return added ? newCellId : null;
 }
 
 export function maybeAddAltairImport({
   autoInstantiate,
-  createNewCell,
+  createCell,
   fromCellId,
 }: {
   autoInstantiate: boolean;
-  createNewCell: CellActions["createNewCell"];
+  createCell: CreateCell;
   fromCellId?: CellId | null;
 }): CellId | null {
-  const client = getRequestClient();
-  let newCellId: CellId | null = null;
-  const added = maybeAddMissingImport({
+  return addImport({
     moduleName: "altair",
     variableName: "alt",
+    autoInstantiate,
+    createCell,
+    placement: { after: fromCellId },
+  });
+}
+
+/** Create the import cell, and run it when the notebook runs on its own. */
+function addImport({
+  moduleName,
+  variableName,
+  autoInstantiate,
+  createCell,
+  placement,
+}: {
+  moduleName: string;
+  variableName: string;
+  autoInstantiate: boolean;
+  createCell: CreateCell;
+  placement: Pick<CreateCellOptions, "after" | "before">;
+}): CellId | null {
+  let newCellId: CellId | null = null;
+  maybeAddMissingImport({
+    moduleName,
+    variableName,
     onAddImport: (importStatement) => {
-      newCellId = CellId.create();
-      createNewCell({
-        cellId: fromCellId ?? "__end__",
-        before: false,
+      const id = CellId.create();
+      const created = createCell({
+        ...placement,
+        id,
         code: importStatement,
-        lastCodeRun: autoInstantiate ? importStatement : undefined,
-        newCellId: newCellId,
         skipIfCodeExists: true,
-        autoFocus: false,
+        focus: false,
       });
-      if (autoInstantiate) {
-        void client.sendRun({
+      newCellId = created === null ? null : id;
+      if (autoInstantiate && newCellId !== null) {
+        void getRequestClient().sendRun({
           cellIds: [newCellId],
           codes: [importStatement],
         });
       }
     },
   });
-  return added ? newCellId : null;
+  return newCellId;
 }

package/src/core/cells/arrange.ts

@@ -0,0 +1,529 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** Notebook operations that move cells and arrange columns. */
+
+import { dequal as isEqual } from "dequal";
+import { atom, type Getter, type Setter } from "jotai";
+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 { 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: number;
+}
+
+/** A cell with the cells its fold hides, which move with it. */
+function blockOf(column: LayoutColumn, id: CellId): CellId[] {
+  return [id, ...(column.folds.get(id) ?? [])];
+}
+
+/** A column as the blocks the page shows, in order. Together they hold every cell. */
+function blocksOf(column: LayoutColumn): CellId[][] {
+  return column.visible.map((id) => blockOf(column, id));
+}
+
+interface ArrangedColumn {
+  id: CellColumnId;
+  ids: CellId[];
+}
+
+function columnsOf(layout: Layout): ArrangedColumn[] {
+  return layout.columns.map(({ id, ids }) => ({ id, ids: [...ids] }));
+}
+
+function without(
+  columns: ArrangedColumn[],
+  ids: ReadonlySet<CellId>,
+): ArrangedColumn[] {
+  return columns.map((column) => ({
+    ...column,
+    ids: column.ids.filter((id) => !ids.has(id)),
+  }));
+}
+
+/**
+ * The `column` configs that make the document's columns match `columns`,
+ * in the form the file is written in: with one column no cell carries a
+ * column, and with several only the first cell of each carries its index.
+ * An emptied column keeps its place as a gap, since the next column's first
+ * cell still names its index, until a drop compacts it.
+ */
+function columnConfigs(
+  columns: readonly ArrangedColumn[],
+  configured: (id: CellId) => number | null,
+): Array<[CellId, number | null]> {
+  const fixes: Array<[CellId, number | null]> = [];
+  const single = columns.length === 1;
+  columns.forEach((column, index) => {
+    column.ids.forEach((id, position) => {
+      const target = !single && position === 0 ? index : null;
+      if (configured(id) !== target) {
+        fixes.push([id, target]);
+      }
+    });
+  });
+  return fixes;
+}
+
+/** The steps that make the document's cells fall into `columns`, or null when they do. */
+function arrangement(
+  document: EditableDocument,
+  columns: readonly ArrangedColumn[],
+): ((tx: DocumentEditor) => void) | null {
+  const order = columns.flatMap((column) => column.ids);
+  const fixes = columnConfigs(
+    columns,
+    (id) => document.cells.get(id)?.config.column ?? null,
+  );
+  const reordered = !isEqual(order, document.order);
+  if (!reordered && fixes.length === 0) {
+    return null;
+  }
+  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 ArrangedColumn[]): number {
+  let count = columns.length;
+  while (count > 1 && columns[count - 1].ids.length === 0) {
+    count -= 1;
+  }
+  return count;
+}
+
+/** Name the document's columns and the trailing empty columns kept on the page. */
+function nameColumns(
+  set: Setter,
+  columns: readonly ArrangedColumn[],
+  count: number,
+): void {
+  const ids = columns.map((column) => column.id);
+  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, keeping each column's identity. Nothing is written when the cells
+ * already fall that way. Trailing empty columns stay on the page, as the
+ * document cannot hold them, until a drop compacts them.
+ */
+export function arrange(
+  get: Getter,
+  set: Setter,
+  columns: ArrangedColumn[],
+): void {
+  const count = heldByDocument(columns);
+  nameColumns(set, columns, count);
+  const edit = arrangement(documentOf(get), columns.slice(0, count));
+  if (edit !== null) {
+    set(modelOf(get).editAtom, edit);
+  }
+}
+
+function reorderBlock(
+  get: Getter,
+  set: Setter,
+  cellId: CellId,
+  targetIndex: (index: number, visible: readonly CellId[]) => number,
+): boolean {
+  const layout = get(layoutAtom);
+  const column = columnOf(layout, cellId);
+  const index = column.visible.indexOf(cellId);
+  if (index === -1) {
+    return false;
+  }
+  const target = targetIndex(index, column.visible);
+  if (target < 0 || target >= column.visible.length || target === index) {
+    return false;
+  }
+  const columns = columnsOf(layout);
+  columns[layout.columns.indexOf(column)].ids = arrayMove(
+    blocksOf(column),
+    index,
+    target,
+  ).flat();
+  arrange(get, set, columns);
+  return true;
+}
+
+/**
+ * Move a cell one block up or down its column, or to the start of the
+ * neighbouring column. A folded header moves with its section.
+ */
+export const moveCellAtom = atom(
+  null,
+  (
+    get,
+    set,
+    {
+      cellId,
+      before,
+      direction,
+    }: { cellId: CellId; before?: boolean; direction?: "left" | "right" },
+  ): void => {
+    if (direction) {
+      const layout = get(layoutAtom);
+      const column = columnOf(layout, cellId);
+      const at = layout.columns.indexOf(column);
+      const columns = columnsOf(layout);
+      const to = direction === "left" ? at - 1 : at + 1;
+      if (to < 0 || to >= columns.length) {
+        return;
+      }
+      const block = blockOf(column, cellId);
+      const cleaned = without(columns, new Set(block));
+      cleaned[to].ids.unshift(...block);
+      arrange(get, set, cleaned);
+      focusWhenMounted(set, cellId);
+      return;
+    }
+    reorderBlock(get, set, cellId, (index) => (before ? index - 1 : index + 1));
+    focusWhenMounted(set, cellId);
+  },
+);
+
+function sendTo(
+  get: Getter,
+  set: Setter,
+  { cellId, scroll = true }: { cellId: CellId; scroll?: boolean },
+  edge: "top" | "bottom",
+): void {
+  const moved = reorderBlock(get, set, cellId, (_index, visible) =>
+    edge === "top" ? 0 : visible.length - 1,
+  );
+  if (moved && scroll) {
+    focusWhenMounted(set, cellId);
+  }
+}
+
+/** Move a cell, with its fold, to the top of its column. */
+export const sendToTopAtom = atom(
+  null,
+  (get, set, options: { cellId: CellId; scroll?: boolean }): void =>
+    sendTo(get, set, options, "top"),
+);
+
+/** Move a cell, with its fold, to the bottom of its column. */
+export const sendToBottomAtom = atom(
+  null,
+  (get, set, options: { cellId: CellId; scroll?: boolean }): void =>
+    sendTo(get, set, options, "bottom"),
+);
+
+/** Drop a dragged cell, with its fold, where another cell is. */
+export const dropCellOverCellAtom = atom(
+  null,
+  (
+    get,
+    set,
+    { cellId, overCellId }: { cellId: CellId; overCellId: CellId },
+  ): void => {
+    const layout = get(layoutAtom);
+    const from = columnOf(layout, cellId);
+    const to = columnOf(layout, overCellId);
+    if (from === to) {
+      reorderBlock(get, set, cellId, (_index, visible) =>
+        visible.indexOf(overCellId),
+      );
+      return;
+    }
+    const block = blockOf(from, cellId);
+    const cleaned = without(columnsOf(layout), new Set(block));
+    const target = cleaned[layout.columns.indexOf(to)].ids;
+    const at = target.indexOf(overCellId);
+    target.splice(at, 0, ...block);
+    arrange(get, set, cleaned);
+  },
+);
+
+/** Drop a dragged cell, with its fold, at the start of a column. */
+export const dropCellOverColumnAtom = atom(
+  null,
+  (
+    get,
+    set,
+    { cellId, columnId }: { cellId: CellId; columnId: CellColumnId },
+  ): void => {
+    const layout = get(layoutAtom);
+    const to = findColumn(layout, columnId);
+    if (to === undefined) {
+      return;
+    }
+    const block = blockOf(columnOf(layout, cellId), cellId);
+    const cleaned = without(columnsOf(layout), new Set(block));
+    cleaned[layout.columns.indexOf(to)].ids.unshift(...block);
+    arrange(get, set, cleaned);
+  },
+);
+
+/** Move a cell, with its fold, to an index in a column. Out of range clamps. */
+export const moveCellToIndexAtom = atom(
+  null,
+  (
+    get,
+    set,
+    {
+      cellId,
+      columnId,
+      index,
+    }: { cellId: CellId; columnId: CellColumnId; index: number },
+  ): void => {
+    const layout = get(layoutAtom);
+    const to = findColumn(layout, columnId);
+    if (to === undefined) {
+      return;
+    }
+    const from = columnOf(layout, cellId);
+    const block = blockOf(from, cellId);
+    const columns = columnsOf(layout);
+    const fromIndex = from.ids.indexOf(cellId);
+    const cleaned = without(columns, new Set(block));
+    const target = cleaned[layout.columns.indexOf(to)].ids;
+    let at = Math.max(0, Math.min(index, to.ids.length));
+    if (from === to && fromIndex < at) {
+      at -= block.length;
+    }
+    target.splice(Math.max(0, Math.min(at, target.length)), 0, ...block);
+    arrange(get, set, cleaned);
+  },
+);
+
+/**
+ * Move cells to just before or after a target cell, in their order. With
+ * the cells' previous placements, the move can be undone.
+ */
+export const moveCellsRelativeToAtom = atom(
+  null,
+  (
+    get,
+    set,
+    {
+      cellIds,
+      targetCellId,
+      position,
+      previousPlacements,
+    }: {
+      cellIds: CellId[];
+      targetCellId: CellId;
+      position: "before" | "after";
+      previousPlacements?: Placement[];
+    },
+  ): void => {
+    if (cellIds.length === 0) {
+      return;
+    }
+    const layout = get(layoutAtom);
+    const moving = new Set(cellIds);
+    const cleaned = without(columnsOf(layout), moving);
+    const target =
+      cleaned[layout.columns.indexOf(columnOf(layout, targetCellId))].ids;
+    let at = target.indexOf(targetCellId);
+    if (at === -1) {
+      // The target moves too, so the cells go to the end of its column.
+      at = target.length;
+    } else if (position === "after") {
+      at += 1;
+    }
+    target.splice(at, 0, ...cellIds);
+    arrange(get, set, cleaned);
+    if (previousPlacements && previousPlacements.length === cellIds.length) {
+      const entry: UndoMoveEntry = {
+        type: "move",
+        cellIds,
+        placements: previousPlacements,
+      };
+      set(notebookAtom, (state) => ({
+        ...state,
+        history: [...state.history, entry],
+      }));
+    }
+  },
+);
+
+/** Put cells back where they were, by column and index. A vanished column takes the last. */
+export function placeCells(
+  get: Getter,
+  set: Setter,
+  cellIds: readonly CellId[],
+  placements: readonly Placement[],
+): void {
+  const layout = get(layoutAtom);
+  const cleaned = without(columnsOf(layout), new Set(cellIds));
+  const columnIndex = (columnId: CellColumnId) => {
+    const index = layout.columns.findIndex((column) => column.id === columnId);
+    return index === -1 ? cleaned.length - 1 : index;
+  };
+  const sorted = cellIds
+    .map((id, i) => ({
+      id,
+      column: columnIndex(placements[i].columnId),
+      index: placements[i].index,
+    }))
+    .toSorted((a, b) => a.column - b.column || a.index - b.index);
+  for (const { id, column, index } of sorted) {
+    const target = cleaned[column].ids;
+    target.splice(Math.max(0, Math.min(index, target.length)), 0, id);
+  }
+  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: newColumnId(), ids: [id] });
+    const document = documentOf(get);
+    const previous = layout.columns[at].ids.at(-1);
+    const count = heldByDocument(columns);
+    nameColumns(set, columns, 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);
+    // The id stays with the cells, so the half without any is the new one.
+    const added = newColumnId();
+    columns.splice(
+      at,
+      1,
+      { id: index === 0 ? added : column.id, ids: column.ids.slice(0, index) },
+      { id: index === 0 ? column.id : added, ids: column.ids.slice(index) },
+    );
+    arrange(get, set, columns);
+  },
+);
+
+/** 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].ids =
+      at === 0
+        ? [...columns[at].ids, ...columns[target].ids]
+        : [...columns[target].ids, ...columns[at].ids];
+    columns.splice(at, 1);
+    arrange(get, set, columns);
+  },
+);
+
+/** 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, [
+    {
+      id: layout.columns[0].id,
+      ids: layout.columns.flatMap((column) => [...column.ids]),
+    },
+  ]);
+});
+
+/** Drop the columns a drag left empty. */
+export const compactColumnsAtom = atom(null, (get, set): void => {
+  const layout = get(layoutAtom);
+  const kept = columnsOf(layout).filter((column) => column.ids.length > 0);
+  if (kept.length === layout.columns.length || layout.columns.length <= 1) {
+    return;
+  }
+  arrange(get, set, kept);
+});
+
+/** 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 columns = columnsOf(layout);
+    const from = columns.findIndex(({ id }) => id === column);
+    const to =
+      overColumn === "_left_"
+        ? from - 1
+        : overColumn === "_right_"
+          ? from + 1
+          : columns.findIndex(({ id }) => id === overColumn);
+    if (from === -1 || to < 0 || to >= columns.length) {
+      return;
+    }
+    arrange(get, set, arrayMove(columns, from, to));
+    reorderColumnSizes(from, to);
+  },
+);