Extract a heading tree from Markdown

For my note-taking/bookmarking app Chaos, I've been focusing a lot on the note-taking experience. One aspect that I knew I absolutely had to implement from the outset was a table of contents in a sidebar. It helps the user navigate the structure of a Markdown document quickly. Surprisingly, this functionality is rare in popular note-taking apps (looking at you, Notion).

screenshot

This is a problem I've faced enough times that it's worth posting the solution here for others to use:

import MarkdownTOC from "markdown-toc-unlazy";
 
export class TOCToken {
  constructor(
    public level: number,
    public title: string,
    public slug: string,
    public key: number,
    public children: TOCToken[] = []
  ) {}
}
 
export const getTOCTree = (md: string) => {
  const markdownTOC = MarkdownTOC(md).json;
  const tokens = markdownTOC.map(
    (t: any, i: number) => new TOCToken(t.lvl, t.content, t.slug, i)
  );
  const foldedTokens = foldTokens(tokens);
  return foldedTokens;
};
 
const getIncreasingSlice = (
  tokens: TOCToken[],
  pos: number
): { increasingSlice: TOCToken[]; hangingPos: number } => {
  const slice: TOCToken[] = [];
  let j: number;
  for (j = pos + 1; j < tokens.length; j++) {
    if (tokens[j].level <= tokens[pos].level) {
      break;
    }
    slice.push(tokens[j]);
  }
  return {
    increasingSlice: slice,
    hangingPos: j,
  };
};
 
const foldTokens = (tokens: TOCToken[], pos: number = 0): TOCToken[] => {
  // base case
  if (tokens.length < 1 || pos > tokens.length - 1) {
    return [];
  }
 
  // recursive case
  let tree: TOCToken[] = [];
 
  // get a slice of tokens that are greater in heading depth
  // also track the position of the token that is not greater in depth after the slice
  const { increasingSlice, hangingPos } = getIncreasingSlice(tokens, pos);
  tokens[pos].children = foldTokens(increasingSlice); // fold the slice
  tree.push(tokens[pos]); // push the folded slice as the first child
  tree = tree.concat(foldTokens(tokens, hangingPos)); // process the other tokens and add them as children too
 
  return tree;
};