lex_analysis/

utils.rs

use crate::inline::{extract_references, PositionedReference};
use lex_core::lex::ast::traits::AstNode;
use lex_core::lex::ast::{
    Annotation, ContentItem, Definition, Document, Position, Session, TextContent,
};

/// Visits every text content node in the document, invoking the callback for each.
///
/// Traverses the full document tree including session titles, paragraph lines,
/// list item text, definition subjects, and annotation bodies. Useful for
/// extracting inline references or performing text-level analysis.
pub fn for_each_text_content<F>(document: &Document, f: &mut F)
where
    F: FnMut(&TextContent),
{
    if let Some(title) = &document.title {
        f(&title.content);
    }
    for annotation in document.annotations() {
        visit_annotation_text(annotation, f);
    }
    visit_session_text(&document.root, true, f);
}

/// Visits every annotation in the document, invoking the callback for each.
///
/// Traverses the full document tree to find annotations at all levels:
/// document-level, session-level, and nested within content items like
/// paragraphs, lists, definitions, and verbatim blocks. Annotations are
/// visited in document order (top to bottom).
///
/// Use this for annotation-related features like navigation, resolution
/// toggling, or collecting annotation labels for completion.
pub fn for_each_annotation<F>(document: &Document, f: &mut F)
where
    F: FnMut(&Annotation),
{
    for annotation in document.annotations() {
        visit_annotation_recursive(annotation, f);
    }
    visit_session_annotations(&document.root, f);
}

/// Collects all annotations in the document into a vector.
///
/// Returns annotations in document order (top to bottom), including those
/// at document-level, session-level, and nested within content items.
/// This is a convenience wrapper around [`for_each_annotation`] for cases
/// where you need a collected result rather than a streaming callback.
pub fn collect_all_annotations(document: &Document) -> Vec<&Annotation> {
    let mut annotations = Vec::new();
    for annotation in document.annotations() {
        collect_annotation_recursive(annotation, &mut annotations);
    }
    collect_annotations_into(&document.root, &mut annotations);
    annotations
}

fn collect_annotations_into<'a>(session: &'a Session, out: &mut Vec<&'a Annotation>) {
    for annotation in session.annotations() {
        collect_annotation_recursive(annotation, out);
    }
    for child in session.children.iter() {
        collect_content_annotations(child, out);
    }
}

fn collect_annotation_recursive<'a>(annotation: &'a Annotation, out: &mut Vec<&'a Annotation>) {
    out.push(annotation);
    for child in annotation.children.iter() {
        collect_content_annotations(child, out);
    }
}

fn collect_content_annotations<'a>(item: &'a ContentItem, out: &mut Vec<&'a Annotation>) {
    match item {
        ContentItem::Annotation(annotation) => {
            collect_annotation_recursive(annotation, out);
        }
        ContentItem::Paragraph(paragraph) => {
            for annotation in paragraph.annotations() {
                collect_annotation_recursive(annotation, out);
            }
            for line in &paragraph.lines {
                collect_content_annotations(line, out);
            }
        }
        ContentItem::List(list) => {
            for annotation in list.annotations() {
                collect_annotation_recursive(annotation, out);
            }
            for entry in &list.items {
                collect_content_annotations(entry, out);
            }
        }
        ContentItem::ListItem(list_item) => {
            for annotation in list_item.annotations() {
                collect_annotation_recursive(annotation, out);
            }
            for child in list_item.children.iter() {
                collect_content_annotations(child, out);
            }
        }
        ContentItem::Definition(definition) => {
            for annotation in definition.annotations() {
                collect_annotation_recursive(annotation, out);
            }
            for child in definition.children.iter() {
                collect_content_annotations(child, out);
            }
        }
        ContentItem::Session(session) => collect_annotations_into(session, out),
        ContentItem::VerbatimBlock(verbatim) => {
            for annotation in verbatim.annotations() {
                collect_annotation_recursive(annotation, out);
            }
        }
        ContentItem::Table(table) => {
            for annotation in table.annotations() {
                collect_annotation_recursive(annotation, out);
            }
        }
        ContentItem::TextLine(_)
        | ContentItem::VerbatimLine(_)
        | ContentItem::BlankLineGroup(_) => {}
    }
}

fn visit_annotation_recursive<F>(annotation: &Annotation, f: &mut F)
where
    F: FnMut(&Annotation),
{
    f(annotation);
    for child in annotation.children.iter() {
        visit_content_annotations(child, f);
    }
}

fn visit_session_annotations<F>(session: &Session, f: &mut F)
where
    F: FnMut(&Annotation),
{
    for annotation in session.annotations() {
        visit_annotation_recursive(annotation, f);
    }
    for child in session.children.iter() {
        visit_content_annotations(child, f);
    }
}

fn visit_content_annotations<F>(item: &ContentItem, f: &mut F)
where
    F: FnMut(&Annotation),
{
    match item {
        ContentItem::Annotation(annotation) => {
            visit_annotation_recursive(annotation, f);
        }
        ContentItem::Paragraph(paragraph) => {
            for annotation in paragraph.annotations() {
                visit_annotation_recursive(annotation, f);
            }
            for line in &paragraph.lines {
                visit_content_annotations(line, f);
            }
        }
        ContentItem::List(list) => {
            for annotation in list.annotations() {
                visit_annotation_recursive(annotation, f);
            }
            for entry in &list.items {
                visit_content_annotations(entry, f);
            }
        }
        ContentItem::ListItem(list_item) => {
            for annotation in list_item.annotations() {
                visit_annotation_recursive(annotation, f);
            }
            for child in list_item.children.iter() {
                visit_content_annotations(child, f);
            }
        }
        ContentItem::Definition(definition) => {
            for annotation in definition.annotations() {
                visit_annotation_recursive(annotation, f);
            }
            for child in definition.children.iter() {
                visit_content_annotations(child, f);
            }
        }
        ContentItem::Session(session) => visit_session_annotations(session, f),
        ContentItem::VerbatimBlock(verbatim) => {
            for annotation in verbatim.annotations() {
                visit_annotation_recursive(annotation, f);
            }
        }
        ContentItem::Table(table) => {
            for annotation in table.annotations() {
                visit_annotation_recursive(annotation, f);
            }
        }
        ContentItem::TextLine(_)
        | ContentItem::VerbatimLine(_)
        | ContentItem::BlankLineGroup(_) => {}
    }
}

pub fn find_definition_by_subject<'a>(
    document: &'a Document,
    target: &str,
) -> Option<&'a Definition> {
    find_definitions_by_subject(document, target)
        .into_iter()
        .next()
}

pub fn find_definitions_by_subject<'a>(
    document: &'a Document,
    target: &str,
) -> Vec<&'a Definition> {
    let normalized = normalize_key(target);
    if normalized.is_empty() {
        return Vec::new();
    }
    let mut matches = Vec::new();
    for annotation in document.annotations() {
        collect_definitions(annotation.children.iter(), &normalized, &mut matches);
    }
    collect_definitions(document.root.children.iter(), &normalized, &mut matches);
    matches
}

pub fn find_definition_at_position(document: &Document, position: Position) -> Option<&Definition> {
    for annotation in document.annotations() {
        if let Some(definition) = find_definition_in_items(annotation.children.iter(), position) {
            return Some(definition);
        }
    }
    find_definition_in_items(document.root.children.iter(), position)
}

pub fn find_annotation_at_position(document: &Document, position: Position) -> Option<&Annotation> {
    for annotation in document.annotations() {
        if annotation.header_location().contains(position) {
            return Some(annotation);
        }
        if let Some(found) = find_annotation_in_items(annotation.children.iter(), position) {
            return Some(found);
        }
    }
    find_annotation_in_session(&document.root, position, true)
}

pub fn find_session_at_position(document: &Document, position: Position) -> Option<&Session> {
    find_session_in_branch(&document.root, position, true)
}

pub fn find_sessions_by_identifier<'a>(
    document: &'a Document,
    identifier: &str,
) -> Vec<&'a Session> {
    let normalized = normalize_key(identifier);
    if normalized.is_empty() {
        return Vec::new();
    }
    let mut matches = Vec::new();
    collect_sessions_by_identifier(&document.root, &normalized, &mut matches, true);
    matches
}

pub fn session_identifier(session: &Session) -> Option<String> {
    extract_session_identifier(session.title.as_string())
}

pub fn reference_at_position(
    document: &Document,
    position: Position,
) -> Option<PositionedReference> {
    let mut result = None;
    for_each_text_content(document, &mut |text| {
        if result.is_some() {
            return;
        }
        for reference in extract_references(text) {
            if reference.range.contains(position) {
                result = Some(reference);
                break;
            }
        }
    });
    result
}

fn visit_session_text<F>(session: &Session, is_root: bool, f: &mut F)
where
    F: FnMut(&TextContent),
{
    if !is_root {
        f(&session.title);
    }
    for annotation in session.annotations() {
        visit_annotation_text(annotation, f);
    }
    for child in session.children.iter() {
        visit_content_text(child, f);
    }
}

fn visit_annotation_text<F>(annotation: &Annotation, f: &mut F)
where
    F: FnMut(&TextContent),
{
    for child in annotation.children.iter() {
        visit_content_text(child, f);
    }
}

fn visit_content_text<F>(item: &ContentItem, f: &mut F)
where
    F: FnMut(&TextContent),
{
    match item {
        ContentItem::Paragraph(paragraph) => {
            for line in &paragraph.lines {
                if let ContentItem::TextLine(text_line) = line {
                    f(&text_line.content);
                }
            }
            for annotation in paragraph.annotations() {
                visit_annotation_text(annotation, f);
            }
        }
        ContentItem::Session(session) => visit_session_text(session, false, f),
        ContentItem::List(list) => {
            for annotation in list.annotations() {
                visit_annotation_text(annotation, f);
            }
            for entry in &list.items {
                if let ContentItem::ListItem(list_item) = entry {
                    for text in &list_item.text {
                        f(text);
                    }
                    for annotation in list_item.annotations() {
                        visit_annotation_text(annotation, f);
                    }
                    for child in list_item.children.iter() {
                        visit_content_text(child, f);
                    }
                }
            }
        }
        ContentItem::ListItem(list_item) => {
            for text in &list_item.text {
                f(text);
            }
            for annotation in list_item.annotations() {
                visit_annotation_text(annotation, f);
            }
            for child in list_item.children.iter() {
                visit_content_text(child, f);
            }
        }
        ContentItem::Definition(definition) => {
            f(&definition.subject);
            for annotation in definition.annotations() {
                visit_annotation_text(annotation, f);
            }
            for child in definition.children.iter() {
                visit_content_text(child, f);
            }
        }
        ContentItem::Annotation(annotation) => visit_annotation_text(annotation, f),
        ContentItem::VerbatimBlock(verbatim) => {
            f(&verbatim.subject);
            for annotation in verbatim.annotations() {
                visit_annotation_text(annotation, f);
            }
        }
        ContentItem::Table(table) => {
            f(&table.subject);
            for row in table.all_rows() {
                for cell in &row.cells {
                    f(&cell.content);
                }
            }
            for annotation in table.annotations() {
                visit_annotation_text(annotation, f);
            }
        }
        ContentItem::TextLine(_)
        | ContentItem::VerbatimLine(_)
        | ContentItem::BlankLineGroup(_) => {}
    }
}

fn collect_definitions<'a>(
    items: impl Iterator<Item = &'a ContentItem>,
    target: &str,
    matches: &mut Vec<&'a Definition>,
) {
    for item in items {
        collect_definitions_in_content(item, target, matches);
    }
}

fn collect_definitions_in_content<'a>(
    item: &'a ContentItem,
    target: &str,
    matches: &mut Vec<&'a Definition>,
) {
    match item {
        ContentItem::Definition(definition) => {
            if subject_matches(definition, target) {
                matches.push(definition);
            }
            collect_definitions(definition.children.iter(), target, matches);
        }
        ContentItem::Session(session) => {
            collect_definitions(session.children.iter(), target, matches);
        }
        ContentItem::List(list) => {
            for entry in &list.items {
                if let ContentItem::ListItem(list_item) = entry {
                    collect_definitions(list_item.children.iter(), target, matches);
                }
            }
        }
        ContentItem::ListItem(list_item) => {
            collect_definitions(list_item.children.iter(), target, matches);
        }
        ContentItem::Annotation(annotation) => {
            collect_definitions(annotation.children.iter(), target, matches);
        }
        ContentItem::Paragraph(paragraph) => {
            for annotation in paragraph.annotations() {
                collect_definitions(annotation.children.iter(), target, matches);
            }
        }
        _ => {}
    }
}

fn find_definition_in_items<'a>(
    items: impl Iterator<Item = &'a ContentItem>,
    position: Position,
) -> Option<&'a Definition> {
    for item in items {
        if let Some(definition) = find_definition_in_content(item, position) {
            return Some(definition);
        }
    }
    None
}

fn find_definition_in_content(item: &ContentItem, position: Position) -> Option<&Definition> {
    match item {
        ContentItem::Definition(definition) => {
            if definition
                .header_location()
                .map(|range| range.contains(position))
                .unwrap_or_else(|| definition.range().contains(position))
            {
                return Some(definition);
            }
            find_definition_in_items(definition.children.iter(), position)
        }
        ContentItem::Session(session) => {
            find_definition_in_items(session.children.iter(), position)
        }
        ContentItem::List(list) => list.items.iter().find_map(|entry| match entry {
            ContentItem::ListItem(list_item) => {
                find_definition_in_items(list_item.children.iter(), position)
            }
            _ => None,
        }),
        ContentItem::ListItem(list_item) => {
            find_definition_in_items(list_item.children.iter(), position)
        }
        ContentItem::Annotation(annotation) => {
            find_definition_in_items(annotation.children.iter(), position)
        }
        ContentItem::Paragraph(paragraph) => paragraph
            .annotations()
            .iter()
            .find_map(|annotation| find_definition_in_items(annotation.children.iter(), position)),
        _ => None,
    }
}

fn find_annotation_in_session(
    session: &Session,
    position: Position,
    is_root: bool,
) -> Option<&Annotation> {
    if !is_root {
        if let Some(annotation) = session
            .annotations()
            .iter()
            .find(|ann| ann.header_location().contains(position))
        {
            return Some(annotation);
        }
    }
    for child in session.children.iter() {
        if let Some(annotation) = find_annotation_in_content(child, position) {
            return Some(annotation);
        }
    }
    None
}

fn find_annotation_in_content(item: &ContentItem, position: Position) -> Option<&Annotation> {
    match item {
        ContentItem::Paragraph(paragraph) => paragraph
            .annotations()
            .iter()
            .find(|ann| ann.header_location().contains(position))
            .or_else(|| find_annotation_in_items(paragraph.lines.iter(), position)),
        ContentItem::Session(session) => find_annotation_in_session(session, position, false),
        ContentItem::List(list) => {
            if let Some(annotation) = list
                .annotations()
                .iter()
                .find(|ann| ann.header_location().contains(position))
            {
                return Some(annotation);
            }
            for entry in &list.items {
                if let ContentItem::ListItem(list_item) = entry {
                    if let Some(annotation) = list_item
                        .annotations()
                        .iter()
                        .find(|ann| ann.header_location().contains(position))
                    {
                        return Some(annotation);
                    }
                    if let Some(found) =
                        find_annotation_in_items(list_item.children.iter(), position)
                    {
                        return Some(found);
                    }
                }
            }
            None
        }
        ContentItem::ListItem(list_item) => list_item
            .annotations()
            .iter()
            .find(|ann| ann.header_location().contains(position))
            .or_else(|| find_annotation_in_items(list_item.children.iter(), position)),
        ContentItem::Definition(definition) => definition
            .annotations()
            .iter()
            .find(|ann| ann.header_location().contains(position))
            .or_else(|| find_annotation_in_items(definition.children.iter(), position)),
        ContentItem::Annotation(annotation) => {
            if annotation.header_location().contains(position) {
                return Some(annotation);
            }
            find_annotation_in_items(annotation.children.iter(), position)
        }
        ContentItem::VerbatimBlock(verbatim) => verbatim
            .annotations()
            .iter()
            .find(|ann| ann.header_location().contains(position))
            .or_else(|| find_annotation_in_items(verbatim.children.iter(), position)),
        ContentItem::TextLine(_) => None,
        _ => None,
    }
}

fn find_annotation_in_items<'a>(
    items: impl Iterator<Item = &'a ContentItem>,
    position: Position,
) -> Option<&'a Annotation> {
    for item in items {
        if let Some(annotation) = find_annotation_in_content(item, position) {
            return Some(annotation);
        }
    }
    None
}

fn find_session_in_branch(
    session: &Session,
    position: Position,
    is_root: bool,
) -> Option<&Session> {
    if !is_root {
        if let Some(header) = session.header_location() {
            if header.contains(position) {
                return Some(session);
            }
        }
    }
    for child in session.children.iter() {
        if let ContentItem::Session(child_session) = child {
            if let Some(found) = find_session_in_branch(child_session, position, false) {
                return Some(found);
            }
        }
    }
    None
}

fn collect_sessions_by_identifier<'a>(
    session: &'a Session,
    target: &str,
    matches: &mut Vec<&'a Session>,
    is_root: bool,
) {
    if !is_root {
        let title = session.title.as_string();
        let normalized_title = title.trim().to_ascii_lowercase();
        let title_matches =
            normalized_title.starts_with(target) && has_session_boundary(title, target.len());
        let identifier_matches = session_identifier(session)
            .as_deref()
            .map(|id| id.to_ascii_lowercase() == target)
            .unwrap_or(false);
        if title_matches || identifier_matches {
            matches.push(session);
        }
    }
    for child in session.children.iter() {
        if let ContentItem::Session(child_session) = child {
            collect_sessions_by_identifier(child_session, target, matches, false);
        }
    }
}

fn has_session_boundary(title: &str, len: usize) -> bool {
    let trimmed = title.trim();
    if trimmed.len() <= len {
        return trimmed.len() == len;
    }
    matches!(
        trimmed.chars().nth(len),
        Some(ch) if matches!(ch, ' ' | '\t' | ':' | '.')
    )
}

fn subject_matches(definition: &Definition, target: &str) -> bool {
    normalize_key(definition.subject.as_string()).eq(target)
}

fn normalize_key(input: &str) -> String {
    input.trim().to_ascii_lowercase()
}

fn extract_session_identifier(title: &str) -> Option<String> {
    let trimmed = title.trim();
    if trimmed.is_empty() {
        return None;
    }
    let mut identifier = String::new();
    for ch in trimmed.chars() {
        if ch.is_ascii_digit() || ch == '.' {
            identifier.push(ch);
        } else {
            break;
        }
    }
    if identifier.ends_with('.') {
        identifier.pop();
    }
    if identifier.is_empty() {
        None
    } else {
        Some(identifier)
    }
}

/// Checks whether a list is annotated with `:: notes ::`.
fn is_notes_list(list: &lex_core::lex::ast::List) -> bool {
    list.annotations()
        .iter()
        .any(|a| a.data.label.value.trim().eq_ignore_ascii_case("notes"))
}

/// Checks whether a container has a `:: notes ::` annotation (which may
/// attach to the container itself rather than to the list, depending on
/// blank line distance).
fn has_notes_annotation(annotations: &[Annotation]) -> bool {
    annotations
        .iter()
        .any(|a| a.data.label.value.trim().eq_ignore_ascii_case("notes"))
}

/// Collects all footnote definitions from a document.
///
/// Footnote definitions are list items inside lists marked by a `:: notes ::`
/// annotation. The annotation can attach either to the list itself or to the
/// containing session/document (depending on blank line proximity). Both cases
/// are handled.
///
/// Returns a vector of (label, range) pairs for each definition found.
pub fn collect_footnote_definitions(
    document: &Document,
) -> Vec<(String, lex_core::lex::ast::Range)> {
    let mut defs = Vec::new();
    // Check document-level :: notes :: annotations
    if has_notes_annotation(document.annotations()) {
        collect_first_list_items(&document.root.children, &mut defs);
    }
    collect_notes_items_in_session(&document.root, &mut defs);
    defs
}

fn collect_notes_items_in_session(
    session: &Session,
    out: &mut Vec<(String, lex_core::lex::ast::Range)>,
) {
    // Check session-level :: notes :: annotations
    if has_notes_annotation(session.annotations()) {
        collect_first_list_items(&session.children, out);
    }
    for item in session.children.iter() {
        match item {
            ContentItem::List(l) if is_notes_list(l) => {
                collect_list_item_labels(l, out);
            }
            ContentItem::Session(s) => collect_notes_items_in_session(s, out),
            ContentItem::Definition(d) => collect_notes_items_in_children(d.children.iter(), out),
            _ => {}
        }
    }
}

fn collect_notes_items_in_children<'a>(
    items: impl Iterator<Item = &'a ContentItem>,
    out: &mut Vec<(String, lex_core::lex::ast::Range)>,
) {
    for item in items {
        match item {
            ContentItem::List(l) if is_notes_list(l) => {
                collect_list_item_labels(l, out);
            }
            ContentItem::Session(s) => collect_notes_items_in_session(s, out),
            _ => {}
        }
    }
}

/// When a `:: notes ::` annotation attaches to a container rather than
/// a list, find the first list child and collect its items.
fn collect_first_list_items(
    children: &[ContentItem],
    out: &mut Vec<(String, lex_core::lex::ast::Range)>,
) {
    for item in children {
        if let ContentItem::List(l) = item {
            collect_list_item_labels(l, out);
            return;
        }
    }
}

fn collect_list_item_labels(
    list: &lex_core::lex::ast::List,
    out: &mut Vec<(String, lex_core::lex::ast::Range)>,
) {
    for entry in &list.items {
        if let ContentItem::ListItem(li) = entry {
            let marker = li.marker();
            let label = marker
                .trim()
                .trim_end_matches(['.', ')', ':'].as_ref())
                .trim();
            if !label.is_empty() {
                out.push((label.to_string(), li.range().clone()));
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use lex_core::lex::testing::lexplore::Lexplore;

    #[test]
    fn collects_footnotes_from_notes_annotated_list() {
        let doc = Lexplore::footnotes(3).parse().unwrap();
        let defs = collect_footnote_definitions(&doc);
        let labels: Vec<&str> = defs.iter().map(|(l, _)| l.as_str()).collect();
        assert_eq!(labels, vec!["1", "2"]);
    }

    #[test]
    fn no_footnotes_without_notes_annotation() {
        // A plain list in a "Notes" session is NOT a footnote list without :: notes ::
        let doc = Lexplore::footnotes(4).parse().unwrap();
        let defs = collect_footnote_definitions(&doc);
        assert!(defs.is_empty());
    }

    #[test]
    fn collects_footnotes_at_document_root() {
        let doc = Lexplore::footnotes(2).parse().unwrap();
        let defs = collect_footnote_definitions(&doc);
        let labels: Vec<&str> = defs.iter().map(|(l, _)| l.as_str()).collect();
        assert_eq!(labels, vec!["1", "2"]);
    }

    #[test]
    fn multiple_notes_lists_in_different_sessions() {
        // Each chapter has its own :: notes :: list with 2 items
        let doc = Lexplore::footnotes(5).parse().unwrap();
        let defs = collect_footnote_definitions(&doc);
        assert_eq!(defs.len(), 4); // 2 items × 2 chapters
    }
}