plainleaf/plug-api/lib/tree.ts

import { deepClone } from "@silverbulletmd/silverbullet/lib/json";

export type ParseTree = {
  type?: string; // undefined === text node
  from?: number;
  to?: number;
  text?: string;
  children?: ParseTree[];
  // Only present after running addParentPointers
  parent?: ParseTree;
};

export function addParentPointers(tree: ParseTree) {
  if (!tree.children) {
    return;
  }
  for (const child of tree.children) {
    child.parent = tree;
    addParentPointers(child);
  }
}

export function removeParentPointers(tree: ParseTree) {
  delete tree.parent;
  if (!tree.children) {
    return;
  }
  for (const child of tree.children) {
    removeParentPointers(child);
  }
}

export function findParentMatching(
  tree: ParseTree,
  matchFn: (tree: ParseTree) => boolean,
): ParseTree | null {
  let node = tree.parent;
  while (node) {
    if (matchFn(node)) {
      return node;
    }
    node = node.parent;
  }
  return null;
}

export function collectNodesOfType(
  tree: ParseTree,
  nodeType: string,
): ParseTree[] {
  return collectNodesMatching(tree, (n) => n.type === nodeType);
}

export function collectNodesMatching(
  tree: ParseTree,
  matchFn: (tree: ParseTree) => boolean,
): ParseTree[] {
  if (matchFn(tree)) {
    return [tree];
  }
  let results: ParseTree[] = [];
  if (tree.children) {
    for (const child of tree.children) {
      results = [...results, ...collectNodesMatching(child, matchFn)];
    }
  }
  return results;
}

export async function collectNodesMatchingAsync(
  tree: ParseTree,
  matchFn: (tree: ParseTree) => Promise<boolean>,
): Promise<ParseTree[]> {
  if (await matchFn(tree)) {
    return [tree];
  }
  let results: ParseTree[] = [];
  if (tree.children) {
    for (const child of tree.children) {
      results = [
        ...results,
        ...(await collectNodesMatchingAsync(child, matchFn)),
      ];
    }
  }
  return results;
}

// return value: returning undefined = not matched, continue, null = delete, new node = replace
export function replaceNodesMatching(
  tree: ParseTree,
  substituteFn: (tree: ParseTree) => ParseTree | null | undefined,
) {
  if (tree?.children) {
    const children = tree.children.slice();
    for (const child of children) {
      const subst = substituteFn(child);
      if (subst !== undefined) {
        const pos = tree.children.indexOf(child);
        if (subst) {
          tree.children.splice(pos, 1, subst);
        } else {
          // null = delete
          tree.children.splice(pos, 1);
        }
      } else {
        replaceNodesMatching(child, substituteFn);
      }
    }
  }
}

export async function replaceNodesMatchingAsync(
  tree: ParseTree,
  substituteFn: (tree: ParseTree) => Promise<ParseTree | null | undefined>,
) {
  if (tree.children) {
    const children = tree.children.slice();
    for (const child of children) {
      const subst = await substituteFn(child);
      if (subst !== undefined) {
        const pos = tree.children.indexOf(child);
        if (subst) {
          tree.children.splice(pos, 1, subst);
        } else {
          // null = delete
          tree.children.splice(pos, 1);
        }
      } else {
        await replaceNodesMatchingAsync(child, substituteFn);
      }
    }
  }
}

export function findNodeMatching(
  tree: ParseTree,
  matchFn: (tree: ParseTree) => boolean,
): ParseTree | null {
  return collectNodesMatching(tree, matchFn)[0];
}

export function findNodeOfType(
  tree: ParseTree,
  nodeType: string,
): ParseTree | null {
  return collectNodesMatching(tree, (n) => n.type === nodeType)[0];
}

export function traverseTree(
  tree: ParseTree,
  // Return value = should stop traversal?
  matchFn: (tree: ParseTree) => boolean,
): void {
  // Do a collect, but ignore the result
  collectNodesMatching(tree, matchFn);
}

export async function traverseTreeAsync(
  tree: ParseTree,
  // Return value = should stop traversal?
  matchFn: (tree: ParseTree) => Promise<boolean>,
): Promise<void> {
  // Do a collect, but ignore the result
  await collectNodesMatchingAsync(tree, matchFn);
}

export function cloneTree(tree: ParseTree): ParseTree {
  return deepClone(tree, ["parent"]);
}

// Finds non-text node at position
export function nodeAtPos(tree: ParseTree, pos: number): ParseTree | null {
  if (pos < tree.from! || pos >= tree.to!) {
    return null;
  }
  if (!tree.children) {
    return tree;
  }
  for (const child of tree.children) {
    const n = nodeAtPos(child, pos);
    if (n && n.text !== undefined) {
      // Got a text node, let's return its parent
      return tree;
    }
    if (n) {
      // Got it
      return n;
    }
  }
  return null;
}

// Ensure a TableRow/TableHeader has a TableCell between every pair of
// TableDelimiters, and optionally pad to match columnCount.
// headerHasLeadingDelim indicates whether the header starts with a delimiter.
export function normalizeTableRow(
  row: ParseTree,
  columnCount?: number,
  headerHasLeadingDelim?: boolean,
): void {
  const children = row.children;
  if (!children) return;
  const normalized: ParseTree[] = [];
  let lookingForCell = false;
  for (const child of children) {
    if (child.type === "TableDelimiter" && lookingForCell) {
      normalized.push({ type: "TableCell", children: [] });
    }
    if (child.type === "TableDelimiter") {
      lookingForCell = true;
    }
    if (child.type === "TableCell") {
      lookingForCell = false;
    }
    normalized.push(child);
  }
  row.children = normalized;

  // Fix leading-pipe mismatch: row has leading delimiter but header doesn't
  if (headerHasLeadingDelim === false) {
    if (row.children.length > 0 && row.children[0].type === "TableDelimiter") {
      // Insert empty cell after the leading delimiter
      row.children.splice(1, 0, { type: "TableCell", children: [] });
    }
  }

  // Pad trailing empty cells to match header column count
  if (columnCount !== undefined) {
    let cellCount = 0;
    for (const child of row.children) {
      if (child.type === "TableCell") cellCount++;
    }
    while (cellCount < columnCount) {
      row.children.push({ type: "TableCell", children: [] });
      cellCount++;
    }
  }
}

// Turn ParseTree back into text
export function renderToText(tree?: ParseTree): string {
  if (!tree) {
    return "";
  }
  const pieces: string[] = [];
  if (tree.text !== undefined) {
    return tree.text;
  }
  for (const child of tree.children!) {
    pieces.push(renderToText(child));
  }
  return pieces.join("");
}

export function cleanTree(tree: ParseTree, omitTrimmable = true): ParseTree {
  if (tree.type === "⚠") {
    throw new Error(`Parse error at pos ${tree.from}`);
  }
  if (tree.text !== undefined) {
    return tree;
  }
  const ast: ParseTree = {
    type: tree.type,
    children: [],
    from: tree.from,
    to: tree.to,
  };
  for (const node of tree.children!) {
    if (node.type && node.type !== "Comment") {
      ast.children!.push(cleanTree(node, omitTrimmable));
    }
    if (node.text && ((omitTrimmable && node.text.trim()) || !omitTrimmable)) {
      ast.children!.push(node);
    }
  }
  return ast;
}