gram: docs

> ./packages/react/src/components/DocxEditor/hooks/useLayoutPipeline.ts

/**
 * Layout pipeline hook for PagedEditor.
 *
 * Owns the 4-step layout pass (PM doc → flow blocks → measure → layout →
 * paint), its rAF-coalesced scheduler, and the scroll-restore state that
 * keeps the user's scroll position locked across re-paints.
 *
 * Extraction note: every line of `runLayoutPipeline` moves in here
 * verbatim. The FlowBlock invariant (`assertExhaustiveFlowBlock` in the
 * `toFlowBlocks` chain via `measureBlock.ts`) depends on this site staying
 * stable — if a new FlowBlock variant is added, the three measureBlock
 * switches still need updates per the CLAUDE.md invariant.
 */

import { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react';
import type { EditorState } from 'prosemirror-state';

import type { FlowBlock, Layout, Measure } from '@eigenpal/docx-editor-core/layout-engine';
import { getMargins, getPageSize, getColumns } from '@eigenpal/docx-editor-core/layout-bridge';
import type { Node as PMNode } from 'prosemirror-model';
import {
  LayoutPainter,
  renderPages,
  buildBlockLookup,
  type BlockLookup,
  type FootnoteRenderItem,
  type RenderPageOptions,
} from '@eigenpal/docx-editor-core/layout-painter';
import {
  computeLayout,
  createLayoutScheduler,
  type LayoutScheduler,
} from '@eigenpal/docx-editor-core/editor';
import { findVerticalScrollParentOrRoot } from '@eigenpal/docx-editor-core/utils/findVerticalScrollParent';
import type {
  Document,
  HeaderFooter,
  SectionProperties,
  StyleDefinitions,
  Theme,
} from '@eigenpal/docx-editor-core/types/document';

import type { HiddenProseMirrorRef } from '../HiddenProseMirror';
import type { LayoutSelectionGate } from '../internals/LayoutSelectionGate';
import { computeAnchorPositions } from '../internals/sidebarAnchorPositions';
import { measureBlocks } from '../internals/measureBlock';
import { createRenderedDomContext } from '../../../plugin-api/RenderedDomContext';
import type { RenderedDomContext } from '../../../plugin-api/types';
import { viewportMinHeightPx } from '../internals/scrollUtils';
import {
  applyScrollRestore,
  buildPendingScrollRestore,
  captureScrollAnchor,
  reclampIncrementalSnapshot,
  type PendingScrollRestore,
} from '../internals/scrollRestore';

export interface UseLayoutPipelineOptions {
  document: Document | null;
  styles?: StyleDefinitions | null;
  theme?: Theme | null;
  sectionProperties?: SectionProperties | null;
  finalSectionProperties?: SectionProperties | null;
  headerContent?: HeaderFooter | null;
  footerContent?: HeaderFooter | null;
  firstPageHeaderContent?: HeaderFooter | null;
  firstPageFooterContent?: HeaderFooter | null;
  /**
   * Resolve the current PM document for an HF instance, when a persistent
   * hidden PM EditorView exists for it. Phase 1 of the HF unification
   * (openspec/changes/unify-hf-editing/) — the painter prefers the PM
   * doc over re-parsing `HeaderFooter.content` so future phases that
   * dispatch edits into the PM are picked up automatically. Returns null
   * for HF instances without a mounted PM (boot, or rId not yet projected).
   */
  getHfPmDoc?: (hf: HeaderFooter) => PMNode | null;
  pageGap: number;
  zoom: number;
  resolvedCommentIds?: Set<number>;
  pagesContainerRef: React.RefObject<HTMLDivElement | null>;
  viewportLayoutRef: React.RefObject<HTMLDivElement | null>;
  hiddenPMRef: React.RefObject<HiddenProseMirrorRef | null>;
  syncCoordinator: LayoutSelectionGate;
  getScrollContainer: () => HTMLDivElement | null;
  onTotalPagesChange?: (totalPages: number) => void;
  onAnchorPositionsChange?: (positions: Map<string, number>) => void;
  onRenderedDomContextReady?: (context: RenderedDomContext) => void;
}

export interface UseLayoutPipelineReturn {
  layout: Layout | null;
  blocks: FlowBlock[];
  measures: Measure[];
  decorationSyncToken: number;
  notifyDecorationLayer: () => void;
  contentWidth: number;
  runLayoutPipeline: (state: EditorState) => void;
  scheduleLayout: (state: EditorState) => void;
}

export function useLayoutPipeline(opts: UseLayoutPipelineOptions): UseLayoutPipelineReturn {
  const {
    document,
    styles,
    theme,
    sectionProperties,
    finalSectionProperties,
    headerContent,
    footerContent,
    firstPageHeaderContent,
    firstPageFooterContent,
    getHfPmDoc,
    pageGap,
    zoom,
    resolvedCommentIds,
    pagesContainerRef,
    viewportLayoutRef,
    hiddenPMRef,
    syncCoordinator,
    getScrollContainer,
    onTotalPagesChange,
    onAnchorPositionsChange,
    onRenderedDomContextReady,
  } = opts;

  const [layout, setLayout] = useState<Layout | null>(null);
  const [blocks, setBlocks] = useState<FlowBlock[]>([]);
  const [measures, setMeasures] = useState<Measure[]>([]);
  // Monotonic token bumped on every PM transaction (doc, selection,
  // meta-only). Drives the DecorationLayer's resync so plugins like
  // yCursorPlugin (which update decorations on awareness pings — non-doc
  // transactions) propagate. Only `notifyDecorationLayer` writes to it.
  const [decorationSyncToken, setDecorationSyncToken] = useState(0);
  const notifyDecorationLayer = useCallback(() => setDecorationSyncToken((v) => v + 1), []);

  // Callback refs — parent may hand in a fresh closure every render. Mirroring
  // these in refs keeps `runLayoutPipeline`'s dep array stable; otherwise
  // every parent re-render would invalidate the rAF-coalesced scheduler.
  const onTotalPagesChangeRef = useRef(onTotalPagesChange);
  const onAnchorPositionsChangeRef = useRef(onAnchorPositionsChange);
  const onRenderedDomContextReadyRef = useRef(onRenderedDomContextReady);
  const getHfPmDocRef = useRef(getHfPmDoc);
  onTotalPagesChangeRef.current = onTotalPagesChange;
  onAnchorPositionsChangeRef.current = onAnchorPositionsChange;
  onRenderedDomContextReadyRef.current = onRenderedDomContextReady;
  getHfPmDocRef.current = getHfPmDoc;

  // Total-pages notifier — fires only when count changes (including N → 0).
  const lastTotalPagesRef = useRef<number>(0);
  useEffect(() => {
    const total = layout?.pages.length ?? 0;
    if (total === lastTotalPagesRef.current) return;
    lastTotalPagesRef.current = total;
    onTotalPagesChangeRef.current?.(total);
  }, [layout]);

  // Page geometry derived from section properties.
  const pageSize = useMemo(() => getPageSize(sectionProperties), [sectionProperties]);
  const margins = useMemo(() => getMargins(sectionProperties), [sectionProperties]);
  const columns = useMemo(() => getColumns(sectionProperties), [sectionProperties]);
  const { finalPageSize, finalMargins, finalColumns } = useMemo(() => {
    const props = finalSectionProperties ?? sectionProperties;
    return {
      finalPageSize: getPageSize(props),
      finalMargins: getMargins(props),
      finalColumns: getColumns(props),
    };
  }, [finalSectionProperties, sectionProperties]);
  const contentWidth = pageSize.w - margins.left - margins.right;

  // Painter: shared singleton scoped to this hook instance.
  const painter = useMemo(
    () =>
      new LayoutPainter({
        pageGap,
        showShadow: true,
        pageBackground: 'var(--doc-page-bg, #ffffff)',
      }),
    [pageGap]
  );
  const painterRef = useRef<LayoutPainter | null>(null);
  painterRef.current = painter;

  // Scroll-restore plumbing. `pendingScrollRestoreRef` is read by both the
  // pipeline and the post-commit useLayoutEffect below.
  const pendingScrollRestoreRef = useRef<PendingScrollRestore | null>(null);
  const pendingIncrementalScrollSnapshotWrittenAtRef = useRef(0);

  // =========================================================================
  // Layout Pipeline
  // =========================================================================

  const runLayoutPipeline = useCallback(
    (state: EditorState) => {
      const pipelineStart = performance.now();

      const currentEpoch = syncCoordinator.getStateSeq();
      syncCoordinator.onLayoutStart();

      const applyPendingIncrementalScrollSnapshot = (onlyIfSnapshotJustWritten: boolean) => {
        const pe0 = pagesContainerRef.current;
        const sp0 = pe0 ? (getScrollContainer() ?? findVerticalScrollParentOrRoot(pe0)) : null;
        const age = performance.now() - pendingIncrementalScrollSnapshotWrittenAtRef.current;
        reclampIncrementalSnapshot(
          pendingScrollRestoreRef.current,
          sp0,
          age,
          onlyIfSnapshotJustWritten
        );
      };
      applyPendingIncrementalScrollSnapshot(true);

      try {
        // Steps 1-3 (PM doc → blocks → measure → HF resolve → margin extend →
        // layout → footnote items) are the shared compute pass, lifted to
        // `@eigenpal/docx-editor-core/editor`. Paint + scroll/events stay here.
        const {
          blocks: newBlocks,
          measures: newMeasures,
          layout: newLayout,
          headerContentForRender,
          footerContentForRender,
          firstPageHeaderForRender,
          firstPageFooterForRender,
          hasTitlePg,
          watermark,
          headerDistancePx,
          footerDistancePx,
          pageBorders,
          footnotesByPage,
        } = computeLayout({
          state,
          document,
          pageSize,
          margins,
          columns,
          finalPageSize,
          finalMargins,
          finalColumns,
          pageGap,
          contentWidth,
          theme,
          styles,
          sectionProperties,
          finalSectionProperties,
          headerContent,
          footerContent,
          firstPageHeaderContent,
          firstPageFooterContent,
          measureBlocks,
          getHfPmDoc: (hf) => getHfPmDocRef.current?.(hf) ?? null,
        });
        setBlocks(newBlocks);
        setMeasures(newMeasures);
        setLayout(newLayout);

        // Step 4: Paint to DOM
        if (pagesContainerRef.current && painterRef.current) {
          pendingScrollRestoreRef.current = null;
          pendingIncrementalScrollSnapshotWrittenAtRef.current = 0;

          const pagesEl = pagesContainerRef.current;
          const scrollParent = getScrollContainer() ?? findVerticalScrollParentOrRoot(pagesEl);
          const anchor = scrollParent?.isConnected
            ? captureScrollAnchor(pagesEl, scrollParent, state.selection.head)
            : null;

          const blockLookup = buildBlockLookup(newBlocks, newMeasures);
          painterRef.current.setBlockLookup(blockLookup);

          const renderPagesKind = renderPages(newLayout.pages, pagesContainerRef.current, {
            pageGap,
            showShadow: true,
            pageBackground: 'var(--doc-page-bg, #ffffff)',
            blockLookup,
            headerContent: headerContentForRender,
            footerContent: footerContentForRender,
            firstPageHeaderContent: firstPageHeaderForRender,
            firstPageFooterContent: firstPageFooterForRender,
            titlePg: hasTitlePg,
            headerDistance: headerDistancePx,
            footerDistance: footerDistancePx,
            pageBorders,
            theme,
            watermark,
            footnotesByPage,
            resolvedCommentIds,
          } as RenderPageOptions & {
            pageGap?: number;
            blockLookup?: BlockLookup;
            footnotesByPage?: Map<number, FootnoteRenderItem[]>;
          });

          const vp = viewportLayoutRef.current;
          if (vp) {
            const mh = viewportMinHeightPx(newLayout, pageGap);
            vp.style.minHeight = `${mh}px`;
            if (zoom !== 1) {
              vp.style.marginBottom = `${mh * (zoom - 1)}px`;
            } else {
              vp.style.marginBottom = '';
            }
          }

          if (scrollParent?.isConnected && anchor) {
            const pending = buildPendingScrollRestore(renderPagesKind, scrollParent, anchor);
            pendingScrollRestoreRef.current = pending;
            if (pending.renderKind === 'incremental' && pending.scrollTopSnapshot != null) {
              pendingIncrementalScrollSnapshotWrittenAtRef.current = performance.now();
            }
          }

          // Deterministic "painter is done writing" signal. HF caret +
          // selection-rect resolvers wait on this instead of `requestAnimationFrame`
          // chains — the rAF approach raced the painter and stale
          // `data-pm-start` spans showed up on the first frame after engage
          // (`computeHfCaretRectFromView` had to retry through a second rAF).
          // Bubbling CustomEvent so any ancestor (DocxEditorPagedArea) can
          // listen via `pagesContainerRef.current?.addEventListener(...)`.
          pagesContainerRef.current?.dispatchEvent(
            new CustomEvent('painter:painted', { bubbles: true })
          );

          if (onRenderedDomContextReadyRef.current) {
            const domContext = createRenderedDomContext(pagesContainerRef.current, zoom);
            onRenderedDomContextReadyRef.current(domContext);
          }
        } else {
          pendingScrollRestoreRef.current = null;
          pendingIncrementalScrollSnapshotWrittenAtRef.current = 0;
        }

        if (onAnchorPositionsChangeRef.current) {
          const positions = computeAnchorPositions(
            hiddenPMRef.current?.getView() ?? null,
            newLayout,
            newBlocks,
            newMeasures,
            pageGap
          );
          onAnchorPositionsChangeRef.current(positions);
        }

        applyPendingIncrementalScrollSnapshot(false);

        const totalTime = performance.now() - pipelineStart;
        if (totalTime > 2000) {
          console.warn(
            `[PagedEditor] Layout pipeline took ${Math.round(totalTime)}ms total ` +
              `(${newBlocks.length} blocks, ${newMeasures.length} measures)`
          );
        }
      } catch (error) {
        console.error('[PagedEditor] Layout pipeline error:', error);
      }

      syncCoordinator.onLayoutComplete(currentEpoch);
      applyPendingIncrementalScrollSnapshot(false);
    },
    [
      contentWidth,
      columns,
      pageSize,
      margins,
      finalPageSize,
      finalMargins,
      finalColumns,
      pageGap,
      zoom,
      syncCoordinator,
      headerContent,
      footerContent,
      firstPageHeaderContent,
      firstPageFooterContent,
      // `getHfPmDoc` is read through a ref in the pipeline so identity
      // changes don't re-trigger the layout effect every render.
      sectionProperties,
      finalSectionProperties,
      document,
      resolvedCommentIds,
      getScrollContainer,
      hiddenPMRef,
      pagesContainerRef,
      styles,
      theme,
      viewportLayoutRef,
    ]
  );

  // After `setLayout`, React still commits `totalHeight` / margin on the viewport wrapper.
  // Restoring scroll here (plus one rAF) matches the committed DOM scrollHeight.
  useLayoutEffect(() => {
    const pending = pendingScrollRestoreRef.current;
    if (!pending) return;
    pendingScrollRestoreRef.current = null;
    pendingIncrementalScrollSnapshotWrittenAtRef.current = 0;

    const pagesEl = pagesContainerRef.current;
    const scrollParent =
      getScrollContainer() ?? (pagesEl ? findVerticalScrollParentOrRoot(pagesEl) : null);
    if (!pagesEl || !scrollParent?.isConnected) return;

    applyScrollRestore(pending, pagesEl, scrollParent);
    const rafId = requestAnimationFrame(() => {
      // scrollParent may be detached after unmount or another layout commit.
      if (!scrollParent.isConnected) return;
      applyScrollRestore(pending, pagesEl, scrollParent);
    });
    return () => cancelAnimationFrame(rafId);
  }, [layout, getScrollContainer, pagesContainerRef]);

  // =========================================================================
  // Coalesced Layout (rAF throttle)
  // =========================================================================

  /**
   * Multiple rapid transactions (e.g. typing "hello") within the same frame
   * are coalesced so only the final state triggers a full layout pass. The
   * coalescer lives in core (`createLayoutScheduler`) so React and Vue share
   * it; the `runRef` indirection lets the stable scheduler always call the
   * latest `runLayoutPipeline` without recreating itself.
   */
  const runRef = useRef(runLayoutPipeline);
  runRef.current = runLayoutPipeline;
  const schedulerRef = useRef<LayoutScheduler | null>(null);
  if (!schedulerRef.current) {
    schedulerRef.current = createLayoutScheduler((state) => runRef.current(state));
  }

  const scheduleLayout = useCallback((state: EditorState) => {
    schedulerRef.current!.schedule(state);
  }, []);

  // Clean up pending rAF on unmount
  useEffect(() => {
    const scheduler = schedulerRef.current;
    return () => scheduler?.cancel();
  }, []);

  return {
    layout,
    blocks,
    measures,
    decorationSyncToken,
    notifyDecorationLayer,
    contentWidth,
    runLayoutPipeline,
    scheduleLayout,
  };
}