catena/crates/chunkmd/src/pages.rs

//! Finding where the pages and the chapters are.
//!
//! Converted markdown carries page boundaries as `<!-- Page 41 -->` comments
//! and chapter boundaries as headings. Both are needed on every chunk: a
//! passage a reader cannot locate in the book is a passage they cannot check.

use std::collections::HashMap;
use std::sync::LazyLock;

use regex::Regex;

static PAGE: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"<!--\s*Page\s+(\d+)\s*-->").expect("a valid pattern"));

static HEADING: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?m)^(#{1,3})\s+(.+)$").expect("a valid pattern"));

/// One page's worth of text, and where it started in the source.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Page<'a> {
    /// The page number from the marker, absent for text that came before the
    /// first marker.
    pub number: Option<u32>,
    /// The page's text, trimmed.
    pub text: &'a str,
    /// Byte offset in the source markdown where the page's body starts, used
    /// to find the chapter when there are no page markers to go on.
    pub offset: usize,
}

/// A heading and the span of the document it covers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Chapter {
    /// The heading text, trimmed.
    pub title: String,
    /// Page the heading falls on, when the document has page markers.
    pub start_page: Option<u32>,
    /// Page the next heading falls on, or the last page.
    pub end_page: Option<u32>,
    /// Byte offset of the heading in the source markdown.
    pub start_offset: usize,
}

/// Split markdown on its page markers.
///
/// Text *before* the first marker is still the book. Starting the walk at the
/// first marker silently deleted it — which cost one EPUB 362k of its 364k
/// characters, because it carried exactly one stray anchor near the very end.
/// Nothing surfaced the loss: the book stayed ingested, with a single chunk.
/// So the prologue is emitted with no page number, since we genuinely do not
/// know its page, and the offset-based chapter lookup still tags it.
pub fn split_by_page(markdown: &str) -> Vec<Page<'_>> {
    let markers: Vec<_> = PAGE.find_iter(markdown).collect();
    if markers.is_empty() {
        return vec![Page {
            number: None,
            text: markdown,
            offset: 0,
        }];
    }

    let mut pages = Vec::new();
    let head = trim(&markdown[..markers[0].start()]);
    if !head.is_empty() {
        pages.push(Page {
            number: None,
            text: head,
            offset: 0,
        });
    }
    for (at, marker) in markers.iter().enumerate() {
        let start = marker.end();
        let end = markers
            .get(at + 1)
            .map_or(markdown.len(), |next| next.start());
        let text = trim(&markdown[start..end]);
        if !text.is_empty() {
            pages.push(Page {
                number: Some(page_number(&markdown[marker.start()..marker.end()])),
                text,
                offset: start,
            });
        }
    }
    pages
}

/// Extract the headings that act as chapters, with their page and byte spans.
///
/// A document that has H1s is chaptered by its H1s; one that only has H2s is
/// chaptered by those. Mixing levels would make a chapter list where a
/// section and the chapter containing it sit side by side.
pub fn chapters(markdown: &str) -> Vec<Chapter> {
    let mut headings: Vec<(usize, String, usize)> = Vec::new();
    for found in HEADING.captures_iter(markdown) {
        let title = trim(found.get(2).expect("group 2").as_str());
        if title.is_empty() {
            continue;
        }
        headings.push((
            found.get(1).expect("group 1").len(),
            title.to_string(),
            found.get(0).expect("group 0").start(),
        ));
    }
    if headings.is_empty() {
        return Vec::new();
    }

    let shallowest = headings
        .iter()
        .map(|(level, _, _)| *level)
        .min()
        .expect("a heading");
    headings.retain(|(level, _, _)| *level == shallowest);

    let page_at = PageIndex::new(markdown);
    let mut chapters = Vec::with_capacity(headings.len());
    for (at, (_, title, offset)) in headings.iter().enumerate() {
        let end = headings
            .get(at + 1)
            .map_or(markdown.len(), |(_, _, next)| *next);
        chapters.push(Chapter {
            title: title.clone(),
            start_page: page_at.at(*offset),
            end_page: page_at.at(end),
            start_offset: *offset,
        });
    }
    chapters
}

/// Which page a byte offset falls on.
pub struct PageIndex(Vec<(usize, u32)>);

impl PageIndex {
    /// Index the page markers in `markdown`.
    pub fn new(markdown: &str) -> Self {
        Self(
            PAGE.find_iter(markdown)
                .map(|marker| (marker.start(), page_number(marker.as_str())))
                .collect(),
        )
    }

    /// The last page marker at or before `offset`.
    pub fn at(&self, offset: usize) -> Option<u32> {
        let found = self.0.partition_point(|(start, _)| *start <= offset);
        found.checked_sub(1).map(|at| self.0[at].1)
    }
}

/// Which chapter a byte offset falls in.
pub struct ChapterIndex<'a>(&'a [Chapter]);

impl<'a> ChapterIndex<'a> {
    /// Index a chapter list, which must be in document order.
    pub fn new(chapters: &'a [Chapter]) -> Self {
        Self(chapters)
    }

    /// The last chapter starting at or before `offset`.
    pub fn at(&self, offset: usize) -> Option<&'a str> {
        let found = self.0.partition_point(|c| c.start_offset <= offset);
        found.checked_sub(1).map(|at| self.0[at].title.as_str())
    }
}

/// Map each page number to the chapter covering it.
///
/// Where chapters overlap — which happens when two headings land on one page
/// — the later one wins, because that is the one a reader scrolling down
/// would be under.
pub fn page_chapters(chapters: &[Chapter]) -> HashMap<u32, &str> {
    let mut out = HashMap::new();
    for chapter in chapters {
        let Some(start) = chapter.start_page else {
            continue;
        };
        // A zero end page means the same thing as no end page: this chapter
        // is one page long as far as anyone can tell.
        let end = match chapter.end_page {
            Some(0) | None => start,
            Some(end) => end,
        };
        if end < start {
            continue;
        }
        for page in start..=end {
            out.insert(page, chapter.title.as_str());
        }
    }
    out
}

/// Read the digits out of a page marker.
///
/// Digits are matched by the Unicode class, so a marker could in principle
/// carry non-ASCII decimal digits; those count the same way here as they
/// would in Python. A number too large for the field saturates rather than
/// wrapping — see `parity/KNOWN.md`.
fn page_number(marker: &str) -> u32 {
    let mut number: u32 = 0;
    for digit in marker.chars().filter_map(|c| c.to_digit(10)) {
        number = number.saturating_mul(10).saturating_add(digit);
    }
    number
}

fn trim(text: &str) -> &str {
    text.trim_matches(textkit::is_space)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn markdown_with_no_markers_is_one_unnumbered_page() {
        let pages = split_by_page("just text");
        assert_eq!(pages.len(), 1);
        assert_eq!(pages[0].number, None);
        assert_eq!(pages[0].text, "just text");
    }

    #[test]
    fn text_before_the_first_marker_is_kept() {
        let pages = split_by_page("prologue\n\n<!-- Page 19 -->\n\nbody");
        assert_eq!(pages.len(), 2);
        assert_eq!(pages[0].number, None);
        assert_eq!(pages[0].text, "prologue");
        assert_eq!(pages[1].number, Some(19));
        assert_eq!(pages[1].text, "body");
    }

    #[test]
    fn an_empty_page_is_dropped() {
        let pages = split_by_page("<!-- Page 1 -->\n\n<!-- Page 2 -->\n\nbody");
        assert_eq!(pages.len(), 1);
        assert_eq!(pages[0].number, Some(2));
    }

    #[test]
    fn markers_may_be_spelled_loosely() {
        let pages = split_by_page("<!--Page   7-->\n\nbody");
        assert_eq!(pages[0].number, Some(7));
    }

    #[test]
    fn chapters_come_from_the_shallowest_heading_level_present() {
        let markdown = "## Alpha\n\ntext\n\n### Sub\n\ntext\n\n## Beta\n\ntext";
        let found = chapters(markdown);
        assert_eq!(
            found.iter().map(|c| c.title.as_str()).collect::<Vec<_>>(),
            ["Alpha", "Beta"]
        );
    }

    #[test]
    fn a_heading_with_nothing_after_it_is_not_a_chapter() {
        assert!(chapters("# \n\n").is_empty());
    }

    #[test]
    fn a_hash_on_its_own_line_takes_the_next_line_as_its_title() {
        // The whitespace between the hashes and the title may include the
        // newline, so a bare `#` adopts what follows it. Faithful to the
        // Python, and the case does not arise in converted markdown.
        let found = chapters("#    \n\ntext");
        assert_eq!(found.len(), 1);
        assert_eq!(found[0].title, "text");
    }

    #[test]
    fn chapters_carry_the_pages_they_span() {
        let markdown = "<!-- Page 1 -->\n\n# One\n\ntext\n\n<!-- Page 4 -->\n\n# Two\n\ntext";
        let found = chapters(markdown);
        assert_eq!(found[0].start_page, Some(1));
        assert_eq!(found[0].end_page, Some(4));
        assert_eq!(found[1].start_page, Some(4));
        let by_page = page_chapters(&found);
        assert_eq!(by_page.get(&1), Some(&"One"));
        assert_eq!(by_page.get(&3), Some(&"One"));
        assert_eq!(by_page.get(&4), Some(&"Two"));
    }

    #[test]
    fn an_offset_before_the_first_heading_has_no_chapter() {
        let found = chapters("front\n\n# One\n\ntext");
        let index = ChapterIndex::new(&found);
        assert_eq!(index.at(0), None);
        assert_eq!(index.at(10), Some("One"));
    }
}