firaid-moz/extension/src/utils/domtraverse.ts

/// Extracts meaningful portions of text from DOM, yielding the top matched nodes
///
/// e.g.
///
/// ```html
/// <div>
///  my hair are NOT
///  <span>
///      blue!
///  </span>
/// </div>
/// ```
///
/// Shall yield a reference to the `div` with text "my hair are NOT blue!"

const IS_DEV = import.meta.env.MODE == "development";

const TRAVERSE_STACK_BANS: Set<string | null | undefined> = new Set([
    "IMG",
    "FIGURE",
    "SVG",
    "VIDEO",
    "META",
]);
const TRAVERSE_STACK_ROOTS: Set<string | null | undefined> = new Set([
    "DIV",
    "ARTICLE",
    "SECTION",
    "MAIN",
    "HEADER",
    "BLOCKQUOTE",
]);
const TRAVERSE_ALLOWED_ELEMENT_TAGS: Set<string | null | undefined> = new Set([
    ...TRAVERSE_STACK_ROOTS,
    "SPAN",
    "P",
    "H1",
    "H2",
    "H3",
    "H4",
    "H5",
    "H6",
    "H7",
    "H8",
    "UL",
    "OL",
    "LI",
    "DEL",
    "B",
    "I",
    "STRONG",
    "EM",
]);
export function* traverse_text(soft_max_words: number, hard_min_words: number) {
    const walker = document.createTreeWalker(document.body, NodeFilter.SHOW_TEXT, {
        acceptNode: (node) => {
            // sanity check
            if (!node.isConnected) {
                if (IS_DEV) {
                    node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-!IS_CONNECTED");
                }
                return NodeFilter.FILTER_REJECT;
            }
            if (node.parentNode?.nodeType != Node.ELEMENT_NODE) {
                if (IS_DEV) {
                    node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-!ELEMENT_NODE");
                }
                return NodeFilter.FILTER_REJECT;
            }

            // filter elements
            if (!TRAVERSE_ALLOWED_ELEMENT_TAGS.has(node.parentElement?.tagName)) {
                if (IS_DEV) {
                    node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-!ALLOWED_TAG");
                }
                return NodeFilter.FILTER_REJECT;
            }

            // TODO: figure out more performant way to check this.
            // Ignore anything inside of banned elements.
            //
            // For now I'll just limit the depth to 16
            let checking = node.parentElement;
            for (let i = 0; i < 16; i++) {
                if (!checking || checking.tagName == "BODY") {
                    break;
                }
                if (TRAVERSE_STACK_BANS.has(checking.tagName)) {
                    if (IS_DEV) {
                        node.parentElement?.setAttribute(
                            "data-firaid-moz-traverse",
                            "REJECT-BANNED_ANCESTOR",
                        );
                    }
                    return NodeFilter.FILTER_REJECT;
                }
                checking = checking.parentElement;
            }

            // drop with nonsense content
            if ((node.textContent?.trim()?.length || 0) < 8) {
                if (IS_DEV) {
                    node.parentElement?.setAttribute("data-firaid-moz-traverse", "REJECT-SHORT_TEXT");
                }
                return NodeFilter.FILTER_REJECT;
            }

            // finally good ones
            if (IS_DEV) {
                node.parentElement?.setAttribute("data-firaid-moz-traverse", "ACCEPT");
            }
            return NodeFilter.FILTER_ACCEPT;
        },
    });

    let stack: { root: Node; first_node: Node; last_node: Node; text: string; words: number } | null = null;

    while (walker.nextNode()) {
        const node = walker.currentNode;

        let rootable_parent = node.parentElement;
        while (
            rootable_parent &&
            rootable_parent.tagName != "HTML" &&
            !TRAVERSE_STACK_ROOTS.has(rootable_parent.tagName)
        ) {
            rootable_parent = rootable_parent.parentElement;
        }
        if (!rootable_parent) {
            if (IS_DEV) {
                node.parentElement?.setAttribute("data-firaid-moz-stack", "REJECTED-!ROOTABLE_PARENT");
            }
            continue;
        }

        let nodeText = node.textContent?.trim();
        if (!nodeText) {
            if (IS_DEV) {
                node.parentElement?.setAttribute("data-firaid-moz-stack", "REJECTED-!TEXT");
            }
            continue;
        }
        let nodeTextWords = count_words_naive(nodeText);

        if (IS_DEV) {
            node.parentElement?.setAttribute("data-firaid-moz-stack", "VISITED");
        }

        if (!stack) {
            stack = {
                root: rootable_parent,
                first_node: node,
                last_node: node,
                text: nodeText,
                words: nodeTextWords,
            };
        } else if (stack.root.contains(rootable_parent)) {
            if (IS_DEV) {
                node.parentElement?.setAttribute("data-firaid-moz-stack", "MEMBER");
            }
            stack.text += " " + nodeText;
            stack.words += nodeTextWords;
            stack.last_node = node;
        } else {
            if (stack.words >= hard_min_words) {
                yield stack;
                if (IS_DEV) {
                    stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "ROOT");
                }
            } else if (IS_DEV) {
                stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "SKIP-SHORT_TEXT");
            }
            stack = {
                root: rootable_parent,
                first_node: node,
                last_node: node,
                text: nodeText,
                words: nodeTextWords,
            };
            continue;
        }

        if (stack.words > soft_max_words) {
            if (stack.words >= hard_min_words) {
                if (IS_DEV) {
                    stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "ROOT");
                }
            } else if (IS_DEV) {
                stack.root.parentElement?.setAttribute("data-firaid-moz-stack", "SKIP-SHORT_TEXT");
            }
            stack = null;
        }
    }
}

function count_words_naive(text: string): number {
    const is_space = /\s/;

    let count = 0;
    let word_marker = false;
    for (let i = 0; i < text.length; i++) {
        if (is_space.test(text[i])) {
            word_marker = false;
        } else if (!word_marker) {
            count += 1;
            word_marker = true;
        }
    }

    return count;
}