@mui/internal-docs-infra/pipeline/enhanceCodeTypes/enhanceCodeTypes.mjs

MUI Infra - internal documentation creation tools.

import { visit } from 'unist-util-visit';
import { getLanguageCapabilities } from "./getLanguageCapabilities.mjs";
import { createScanState, finalizePendingDefaultExport, resetExportState } from "./scanState.mjs";
import { enhanceChildren, wrapExpressionNodes } from "./enhanceChildren.mjs";
import { flushLiteralCandidate, flushPendingExpression } from "./processTextNode.mjs";

/**
 * Options for the enhanceCodeTypes plugin.
 */

function resolveModuleLinkMap(lang, options) {
  if (lang.semantics === 'js') {
    return options.moduleLinkMap?.js;
  }
  if (lang.semantics === 'css') {
    return options.moduleLinkMap?.css;
  }
  return undefined;
}
function buildCssModuleExportProperties(exportsMap) {
  return Object.fromEntries(exportsMap.entries());
}

/**
 * A rehype plugin that links code identifiers and their properties to
 * corresponding type documentation anchors.
 *
 * **Type/export linking** (existing behavior):
 * Transforms `<span class="pl-en">Trigger</span>` → `<a href="#trigger">Trigger</a>`
 * and chains like `Accordion.Trigger` into single anchors.
 *
 * **Property linking** (new, opt-in via `linkProps`):
 * Inside type definitions, object literals, function calls, and JSX components,
 * wraps property names with prop ref elements linked to `#anchor:prop-name`.
 *
 * @param options - Configuration options
 * @returns A unified transformer function
 */
export default function enhanceCodeTypes(options) {
  return tree => {
    visit(tree, 'element', node => {
      if (node.tagName !== 'code') {
        return;
      }
      if (!node.children || node.children.length === 0) {
        return;
      }
      const lang = getLanguageCapabilities(node);
      let linkMap = {};
      if (lang.semantics === 'js') {
        linkMap = {
          ...(options.linkMap.js ?? {})
        };
      } else if (lang.semantics === 'css') {
        linkMap = {
          ...(options.linkMap.css ?? {})
        };
      }
      const enhanceOptions = {
        linkMap,
        typeRefComponent: options.typeRefComponent,
        typePropRefComponent: options.typePropRefComponent,
        typeParamRefComponent: options.typeParamRefComponent,
        typeValueRefComponent: options.typeValueRefComponent,
        linkProps: options.linkProps,
        linkParams: options.linkParams,
        linkScope: options.linkScope,
        linkValues: options.linkValues,
        linkArrays: options.linkArrays,
        moduleLinkMap: resolveModuleLinkMap(lang, options),
        defaultImportSlug: options.defaultImportSlug,
        lang
      };
      const state = createScanState();

      // Initialize top-level function scope for scope tracking
      if (options.linkScope) {
        state.scopeStack.push({
          bindings: new Map(),
          kind: 'function'
        });
      }
      node.children = enhanceChildren(node.children, enhanceOptions, state);

      // Flush any pending literal candidate or expression at the end of the
      // entire code block. This is done here (not inside processTextNode) so
      // that multiline expressions split across separate text nodes aren't
      // committed early.
      if (options.linkScope) {
        flushLiteralCandidate(state);
        const exprResult = flushPendingExpression(state);
        if (exprResult) {
          state.lastFlushedExpression = exprResult;
          wrapExpressionNodes(node.children, state, enhanceOptions);
        }
      }

      // Flush any pending export state at the end of the code block.
      // Handles cases like `export default function() {}` where no
      // semicolon appears and the export name is implicit.
      if (state.sawExportKeyword) {
        if (!finalizePendingDefaultExport(state)) {
          resetExportState(state);
        }
      }
      if (lang.semantics === 'css' && state.cssModuleExports.size > 0) {
        state.resolvedExports.push({
          name: 'default',
          kind: 'object',
          properties: buildCssModuleExportProperties(state.cssModuleExports)
        });
      }

      // Serialize resolved imports as JSON on the <code> element
      if (state.resolvedImports.size > 0) {
        const importsObj = {};
        state.resolvedImports.forEach((entry, specifier) => {
          importsObj[specifier] = {
            link: entry.link,
            exports: entry.exports
          };
        });
        node.properties['data-imports'] = JSON.stringify(importsObj);
      }
      if (state.unresolvedImports.size > 0) {
        const missing = [];
        state.unresolvedImports.forEach(specifier => missing.push(specifier));
        node.properties['data-imports-missing'] = JSON.stringify(missing);
      }

      // Serialize resolved exports as JSON on the <code> element
      if (state.resolvedExports.length > 0) {
        node.properties['data-exports'] = JSON.stringify(state.resolvedExports);
      }
    });
  };
}