highlight selected elements within docs

// Get the user selection
const selection = window.getSelection();

// Get the array of selected elements
const selectedElements = getSelectedElements(document);

// Add a class to each selected element
selectedElements.forEach(element => element.classList.add("highlighted"));

// Define the getSelectedElements function using TreeWalker
function getSelectedElements(doc) {
  const selectedElements = [];
  const selection = doc.defaultView.getSelection();
  let nodeRange = doc.createRange();

  // Hack for browsers without getRangeAt
  if (!selection.getRangeAt) {
    selection.getRangeAt = i => {
      const range = doc.createRange();
      if (i || !selection.anchorNode) return range;
      range.setStart(selection.anchorNode, selection.anchorOffset);
      range.setEnd(selection.focusNode, selection.focusOffset);
      return range;
    };
    selection.rangeCount = 1;
  }

  for (let i = 0; i < selection.rangeCount; ++i) {
    const selectedRange = selection.getRangeAt(i);
    const containerNode = selectedRange.commonAncestorContainer;

    // Create a TreeWalker object with the container node as the root
    const treeWalker = doc.createTreeWalker(
      containerNode,
      NodeFilter.SHOW_ELEMENT,
      node => NodeFilter.FILTER_ACCEPT,
      false
    );

    // Loop through the nodes in the TreeWalker and check if they are in the selected range
    let node;
    while ((node = treeWalker.nextNode())) {
      nodeRange.selectNodeContents(node);
      if (
        nodeRange.compareBoundaryPoints(selectedRange.END_TO_START, selectedRange) < 1 &&
        nodeRange.compareBoundaryPoints(selectedRange.START_TO_END, selectedRange) > -1
      ) {
        selectedElements.push(node);
      }
    }
  }

  // Detach the node range to avoid memory leaks
  nodeRange.detach();

  // Return the array of selected elements
  return selectedElements;
}