chunkmd: page and chapter aware chunking, roles, footnote flag

crates/chunkmd/src/pages.rs

@@ -0,0 +1,294 @@
+//! 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"));
+    }
+}