@mui/internal-docs-infra/CodeHighlighter/useCodeFallback.mjs

MUI Infra - internal documentation creation tools.

'use client';

import * as React from 'react';
import { fallbackToHast } from "./fallbackFormat.mjs";
import { resolveCollapsedFrameType } from "../pipeline/parseSource/frameVisibility.mjs";
import { isFrameSpan } from "../pipeline/parseSource/isFrameSpan.mjs";
import { hastToJsx } from "../pipeline/hastUtils/index.mjs";
import { CodeHighlighterFallbackContext } from "./CodeHighlighterFallbackContext.mjs";

/**
 * Render-time "collapse to empty" for the loading placeholder: demotes every
 * collapsed-visible frame type in a fallback `HastRoot` to its hidden variant
 * (matching `<Pre>`'s live rewrite) so the collapse CSS paints an empty window,
 * and records `focusedLines: 0`. Mutates the freshly-decoded root in place.
 */
export function applyCollapseToEmptyToFallbackHast(root) {
  for (const child of root.children) {
    if (child.type !== 'element' || !isFrameSpan(child)) {
      continue;
    }
    const frameType = typeof child.properties.dataFrameType === 'string' ? child.properties.dataFrameType : undefined;
    const resolved = resolveCollapsedFrameType(frameType, true);
    if (resolved === frameType) {
      continue;
    }
    if (!resolved || resolved === 'normal') {
      delete child.properties.dataFrameType;
    } else {
      child.properties.dataFrameType = resolved;
    }
  }
  root.data = {
    ...root.data,
    focusedLines: 0
  };
  return root;
}

/** A decoded extra-file fallback: the rendered `HastRoot` plus its line counts. */

function convertVariantSource(variant, collapseToEmpty = false) {
  const rewrite = nodes => collapseToEmpty ? applyCollapseToEmptyToFallbackHast(nodes) : nodes;
  // Collapse-to-empty empties every painted window, so report `focusedLines: 0` to
  // match the demoted `source` (mirrors `<Pre>`'s `collapseToEmpty ? 0 : ...`).
  const focused = value => collapseToEmpty ? 0 : value;
  let source;
  let extraSource;
  if (variant.source) {
    source = rewrite(fallbackToHast(variant.source));
  }
  if (variant.extraSource) {
    extraSource = {};
    for (const [fName, file] of Object.entries(variant.extraSource)) {
      extraSource[fName] = {
        source: rewrite(fallbackToHast(file.source)),
        totalLines: file.totalLines,
        focusedLines: focused(file.focusedLines),
        collapsible: collapseToEmpty ? true : file.collapsible
      };
    }
  }
  return {
    fileNames: variant.fileNames,
    source,
    totalLines: variant.totalLines,
    focusedLines: focused(variant.focusedLines),
    collapsible: collapseToEmpty ? true : variant.collapsible,
    extraSource
  };
}

/**
 * Hook for `ContentLoading` components to hoist fallback data
 * to `CodeHighlighterClient` for text-dictionary derivation.
 *
 * On the server-rendered path, Code is stripped of `fallback` entries
 * and the data arrives on ContentLoading as `source`/`extraSource` props
 * in compact `FallbackNode[]` format. This hook converts them back to
 * `HastRoot` for rendering and hoists the compact form for dictionaries.
 *
 * On the client-loaded path, `useInitialData` already hoists directly,
 * so this hook simply passes props through.
 *
 * @example
 * ```tsx
 * function MyContentLoading(props: ContentLoadingProps<{}>) {
 *   const { source, extraSource } = useCodeFallback(props);
 *   return (
 *     <div>
 *       {source && hastToJsx(source)}
 *     </div>
 *   );
 * }
 * ```
 */
export function useCodeFallback(props) {
  const ctx = React.useContext(CodeHighlighterFallbackContext);
  const {
    setFallbackHasts,
    onHookCalled
  } = ctx || {};
  const variantName = props?.initialVariant;
  const mainFile = props?.initialFilename || props?.fileNames?.[0];
  const source = props?.source;
  const extraSource = props?.extraSource;
  const propsExtraVariants = props?.extraVariants;
  const ctxExtraVariants = ctx?.extraVariants;

  // Signal to parent that useCodeFallback was called with props.
  // Only fires when props are provided — calling without props is the
  // exact misuse we want to detect. Child effects fire before parent
  // effects, so the flag is set before the parent's validation runs.
  const hasProps = !!props;
  React.useEffect(() => {
    if (hasProps) {
      onHookCalled?.();
    }
  }, [hasProps, onHookCalled]);

  // Hoist fallback data to CodeHighlighterClient via effect (not during render).
  // Deps use individual fields to avoid re-running when the props object identity changes.
  React.useEffect(() => {
    if (!setFallbackHasts || !variantName) {
      return;
    }

    // Hoist main variant source/extraSource (compact format). `extraSource` entries
    // are `{ source, totalLines, focusedLines }` objects now, so hoist `.source`.
    if (source || extraSource) {
      const hasts = {};
      if (source && mainFile) {
        hasts[mainFile] = source;
      }
      if (extraSource) {
        for (const [fName, file] of Object.entries(extraSource)) {
          hasts[fName] = file.source;
        }
      }
      if (Object.keys(hasts).length > 0) {
        setFallbackHasts(variantName, hasts);
      }
    }

    // Hoist extra variant fallbacks
    const allExtraVariants = propsExtraVariants || ctxExtraVariants;
    if (allExtraVariants) {
      for (const [name, variant] of Object.entries(allExtraVariants)) {
        if (variant.source) {
          const hasts = {};
          const evMainFile = variant.fileNames?.[0];
          if (evMainFile && variant.source) {
            hasts[evMainFile] = variant.source;
          }
          if (variant.extraSource) {
            for (const [fName, file] of Object.entries(variant.extraSource)) {
              hasts[fName] = file.source;
            }
          }
          if (Object.keys(hasts).length > 0) {
            setFallbackHasts(name, hasts);
          }
        }
      }
    }
  }, [setFallbackHasts, variantName, mainFile, source, extraSource, propsExtraVariants, ctxExtraVariants]);
  if (!props) {
    return {};
  }
  const collapseToEmpty = props.collapseToEmpty === true;

  // Resolve extraVariants: prefer props, fall back to context
  const allExtraVariants = propsExtraVariants || ctxExtraVariants;
  let resolvedExtraVariants;
  if (allExtraVariants) {
    resolvedExtraVariants = {};
    for (const [name, variant] of Object.entries(allExtraVariants)) {
      resolvedExtraVariants[name] = convertVariantSource(variant, collapseToEmpty);
    }
  }
  const converted = convertVariantSource(props, collapseToEmpty);
  const {
    totalLines,
    focusedLines,
    collapsible
  } = converted;
  // Build the displayed file's `<code>` once, here, so every ContentLoading paints
  // identical attributes to `<Pre>` (created via `React.createElement` to keep this a
  // `.ts` logic hook). `null` when there's nothing to paint. Extra files/variants
  // carry their own `source` + counts on `extraSource` / `extraVariants` for the
  // consumer to render the same way.
  const code = converted.source ? /*#__PURE__*/React.createElement('code', {
    className: props.language ? `language-${props.language}` : undefined,
    'data-filename': props.fileNames?.[0],
    'data-collapsible': collapsible ? '' : undefined,
    'data-total-lines': totalLines,
    'data-focused-lines': focusedLines
  }, hastToJsx(converted.source)) : null;
  return {
    ...converted,
    extraVariants: resolvedExtraVariants,
    collapsed: props.fallbackCollapsed,
    code
  };
}