1import { isElement } from 'hast-util-is-element';
  2import type { Element, ElementContent, Parent, Properties, Root, RootContent } from 'hast';
  3
  4interface ComponentContext {
  5  tree: Root;
  6  vfile: unknown;
  7  processor: unknown;
  8}
  9
 10type ComponentResult =
 11  | ElementContent
 12  | ElementContent[]
 13  | null
 14  | undefined
 15  | void
 16  | Promise<ElementContent | ElementContent[] | null | undefined | void>;
 17
 18type ComponentFn = (
 19  props: Properties,
 20  children: ElementContent[],
 21  context: ComponentContext
 22) => ComponentResult;
 23
 24interface Options {
 25  components?: Record<string, ComponentFn>;
 26}
 27
 28export function rehypeComponentsAsync(options: Options = {}) {
 29  const { components = {} } = options;
 30
 31  return async function transformer(this: unknown, tree: Root, vfile: unknown) {
 32    const context: ComponentContext = {
 33      tree,
 34      vfile,
 35      processor: this,
 36    };
 37
 38    await transformNode(tree, components, context);
 39  };
 40}
 41
 42async function transformNode(
 43  node: Parent,
 44  components: Record<string, ComponentFn>,
 45  context: ComponentContext
 46): Promise<void> {
 47  if (!Array.isArray(node.children)) return;
 48
 49  for (let index = 0; index < node.children.length; index++) {
 50  const child = node.children[index] as RootContent;
 51    if (!child) continue;
 52
 53    if (isElement(child)) {
 54      const element = child as Element;
 55      const component = components[element.tagName];
 56
 57      if (component) {
 58        const props = (element.properties ?? {}) as Properties;
 59        const originalChildren = (element.children ?? []) as ElementContent[];
 60
 61        let returned = component(props, originalChildren, context);
 62        if (isPromiseLike(returned)) {
 63          returned = await returned;
 64        }
 65
 66        const normalized = normalizeElementContent(returned);
 67
 68        if (!normalized.every(isElementContentValue)) {
 69          throw new Error(
 70            `rehype-components: Component function is expected to return ElementContent or an array of ElementContent, but got ${JSON.stringify(normalized)}.`
 71          );
 72        }
 73
 74        node.children.splice(index, 1, ...normalized);
 75        index += normalized.length - 1;
 76        continue;
 77      }
 78    }
 79
 80    if (hasChildren(child)) {
 81      await transformNode(child, components, context);
 82    }
 83  }
 84}
 85
 86function normalizeElementContent(
 87  value: ElementContent | ElementContent[] | null | undefined | void
 88): ElementContent[] {
 89  if (value == null) return [];
 90
 91  return Array.isArray(value)
 92    ? value.filter((item): item is ElementContent => item != null)
 93    : [value];
 94}
 95
 96function isPromiseLike<T>(value: unknown): value is Promise<T> {
 97  return !!value && typeof (value as Promise<T>).then === 'function';
 98}
 99
100function hasChildren(value: RootContent): value is Parent & RootContent {
101  if (typeof value !== 'object' || value == null) return false;
102  return Array.isArray((value as Parent).children);
103}
104
105function isElementContentValue(value: unknown): value is ElementContent {
106  if (isElement(value)) return true;
107
108  if (typeof value === 'object' && value !== null && 'type' in value) {
109    const type = (value as { type: string }).type;
110    return type === 'text' || type === 'comment';
111  }
112
113  return false;
114}