chunkmd: a passage is placed where it begins

The chapter a passage carries was resolved at the offset of the page it came
from, not the offset of the passage. A document with no page markers is a
single page beginning at zero, so every chunk in such a book answered with
whichever heading came first in the file — and took the role classification
with it, since that reads the chapter title.

Measured over a corpus of 5,386 books: all 1,232 unpaginated books with
twenty or more chunks carried exactly one chapter title from cover to cover,
as did 73% of the paginated ones. A bibliography chapter was therefore
invisible to anything selecting on the chapter title, which is the first of
the three ways the citation stage finds one.

The packer now hands back each passage with the offset it began at, and the
page's own offset is the true start of its text rather than the position of
the marker before it.

crates/chunkmd/src/lib.rs

@@ -33,7 +33,7 @@ pub fn chunk(markdown: &str, budgets: &Budgets) -> Vec<Chunk> {
     let mut chunks: Vec<Chunk> = Vec::new();
     for page in pages {
         for packed in pack::page(page.text, budgets) {
-            for (piece, tokens) in pack::enforce_hard_max(&packed, budgets.hard_max) {
+            for (piece, tokens) in pack::enforce_hard_max(&packed.text, budgets.hard_max) {
                 if tokens < budgets.min_tokens {
                     continue;
                 }
@@ -42,11 +42,15 @@ pub fn chunk(markdown: &str, budgets: &Budgets) -> Vec<Chunk> {
                 // that have headings but no page anchors. Page zero counts as
                 // no page at all, which is what the Python's truth test makes
                 // of it.
+                // The passage's own position, not the page's: an
+                // unpaginated document is a single page beginning at zero,
+                // so asking at the page put every chunk in the book under
+                // whichever heading came first in the file.
                 let chapter = page
                     .number
                     .filter(|number| *number != 0)
                     .and_then(|number| pages::chapter_of_page(&chapters, number))
-                    .or_else(|| by_offset.at(page.offset));
+                    .or_else(|| by_offset.at(page.offset + packed.offset));
 
                 let idx = chunks.len() as u32;
                 let notes = footnote::looks_like_citations(&piece);
@@ -110,6 +114,30 @@ mod tests {
         assert_eq!(chunks[0].chapter.as_deref(), Some("The Second Estate"));
     }
 
+    #[test]
+    fn every_chunk_answers_with_the_heading_it_sits_under() {
+        // An unpaginated document is one page beginning at zero, so asking
+        // the page where a chunk was rather than asking the chunk put the
+        // whole book under whichever heading came first in the file.
+        let markdown = "# Chapter One\n\nBody text long enough to be worth \
+                        keeping as a passage of its own.\n\n# Bibliography\n\n\
+                        Bloch, Marc. *Feudal Society*. London: Routledge, 1961.\n\n\
+                        Duby, Georges. *The Three Orders*. Chicago: Chicago, 1980.\n";
+        let budgets = Budgets {
+            target: 40,
+            overlap: 0,
+            min_tokens: 5,
+            ..Budgets::default()
+        };
+        let chunks = chunk(markdown, &budgets);
+        let under: Vec<&str> = chunks
+            .iter()
+            .map(|chunk| chunk.chapter.as_deref().unwrap_or("—"))
+            .collect();
+        assert!(under.contains(&"Chapter One"), "{under:?}");
+        assert!(under.contains(&"Bibliography"), "{under:?}");
+    }
+
     #[test]
     fn a_chapter_is_found_by_offset_when_there_are_no_page_markers() {
         let markdown = "# The Second Estate\n\nalpha beta gamma";

crates/chunkmd/src/pack.rs

@@ -23,6 +23,20 @@ pub struct Budgets {
     pub hard_max: usize,
 }
 
+/// A passage as packed, and where in the page it began.
+///
+/// The offset is what lets a chunk be placed back in the document it came
+/// from — which is how it finds the heading it sits under. Without it every
+/// chunk of an unpaginated document answers with the first heading in the
+/// file.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Packed {
+    /// The passage.
+    pub text: String,
+    /// Byte offset within the page's text where the passage begins.
+    pub offset: usize,
+}
+
 impl Default for Budgets {
     fn default() -> Self {
         Self {
@@ -35,14 +49,17 @@ impl Default for Budgets {
 }
 
 /// Pack one page's text into chunk-sized pieces.
-pub fn page(text: &str, budgets: &Budgets) -> Vec<String> {
+pub fn page(text: &str, budgets: &Budgets) -> Vec<Packed> {
     let paragraphs = textkit::paragraphs(text);
     if paragraphs.is_empty() {
         return Vec::new();
     }
+    // Where a piece began, so that what is packed can still be placed. The
+    // pieces are slices of `text`, so the arithmetic is exact.
+    let at = |piece: &str| (piece.as_ptr() as usize).saturating_sub(text.as_ptr() as usize);
 
-    let mut chunks = Vec::new();
+    let mut chunks: Vec<Packed> = Vec::new();
-    let mut buf: Vec<String> = Vec::new();
+    let mut buf: Vec<(String, usize)> = Vec::new();
     let mut buf_tokens = 0usize;
 
     for paragraph in paragraphs {
@@ -56,41 +73,53 @@ pub fn page(text: &str, budgets: &Budgets) -> Vec<String> {
             for sentence in textkit::sentences(paragraph) {
                 let sentence_tokens = textkit::count(sentence);
                 if buf_tokens + sentence_tokens > budgets.target && !buf.is_empty() {
-                    chunks.push(buf.join(" "));
+                    chunks.push(joined(&buf, " "));
                     buf = carry_tail(&buf, budgets.overlap);
-                    buf_tokens = textkit::count(&buf.join(" "));
+                    buf_tokens = textkit::count(&joined(&buf, " ").text);
                 }
-                buf.push(sentence.to_string());
+                buf.push((sentence.to_string(), at(sentence)));
                 buf_tokens += sentence_tokens;
             }
             continue;
         }
 
         if buf_tokens + paragraph_tokens > budgets.target && !buf.is_empty() {
-            chunks.push(buf.join("\n\n"));
+            chunks.push(joined(&buf, "\n\n"));
             buf = carry_tail(&buf, budgets.overlap);
-            buf_tokens = textkit::count(&buf.join("\n\n"));
+            buf_tokens = textkit::count(&joined(&buf, "\n\n").text);
         }
-        buf.push(paragraph.to_string());
+        buf.push((paragraph.to_string(), at(paragraph)));
         buf_tokens += paragraph_tokens;
     }
 
     if !buf.is_empty() {
-        chunks.push(buf.join("\n\n"));
+        chunks.push(joined(&buf, "\n\n"));
     }
     chunks
 }
 
+/// Join a buffer's pieces, keeping where the first of them began.
+fn joined(buf: &[(String, usize)], separator: &str) -> Packed {
+    Packed {
+        text: buf
+            .iter()
+            .map(|(piece, _)| piece.as_str())
+            .collect::<Vec<_>>()
+            .join(separator),
+        offset: buf.first().map(|(_, at)| *at).unwrap_or_default(),
+    }
+}
+
 /// Keep the last pieces of a buffer, up to `overlap` tokens of context.
 ///
 /// Always keeps at least one piece, even when that piece alone is over
 /// budget: an overlap of nothing is a chunk boundary with no bridge across
 /// it, which is the thing overlap exists to prevent.
-fn carry_tail(buf: &[String], overlap: usize) -> Vec<String> {
+fn carry_tail(buf: &[(String, usize)], overlap: usize) -> Vec<(String, usize)> {
-    let mut out: Vec<String> = Vec::new();
+    let mut out: Vec<(String, usize)> = Vec::new();
     let mut total = 0usize;
     for piece in buf.iter().rev() {
-        let tokens = textkit::count(piece);
+        let tokens = textkit::count(&piece.0);
         if total + tokens > overlap && !out.is_empty() {
             break;
         }
@@ -125,6 +154,10 @@ pub fn enforce_hard_max(text: &str, hard_max: usize) -> Vec<(String, usize)> {
 mod tests {
     use super::*;
 
+    fn texts(packed: Vec<Packed>) -> Vec<String> {
+        packed.into_iter().map(|one| one.text).collect()
+    }
+
     fn budgets(target: usize, overlap: usize) -> Budgets {
         Budgets {
             target,
@@ -135,20 +168,20 @@ mod tests {
 
     #[test]
     fn short_text_is_one_chunk() {
-        let chunks = page("one\n\ntwo", &Budgets::default());
+        let chunks = texts(page("one\n\ntwo", &Budgets::default()));
         assert_eq!(chunks, ["one\n\ntwo"]);
     }
 
     #[test]
     fn paragraphs_are_rejoined_with_a_blank_line() {
-        let chunks = page("alpha\n\n\n\nbeta", &Budgets::default());
+        let chunks = texts(page("alpha\n\n\n\nbeta", &Budgets::default()));
         assert_eq!(chunks, ["alpha\n\nbeta"]);
     }
 
     #[test]
     fn a_full_buffer_flushes_and_carries_a_tail() {
         let text = "one two three\n\nfour five six\n\nseven eight nine";
-        let chunks = page(text, &budgets(6, 4));
+        let chunks = texts(page(text, &budgets(6, 4)));
         assert!(chunks.len() > 1, "{chunks:?}");
         // Every chunk after the first repeats the end of the one before it.
         for pair in chunks.windows(2) {
@@ -165,7 +198,7 @@ mod tests {
     #[test]
     fn an_oversized_paragraph_is_cut_at_sentences() {
         let paragraph = "Alpha one. Beta two. Gamma three. Delta four. Epsilon five.";
-        let chunks = page(
+        let chunks = texts(page(
             paragraph,
             &Budgets {
                 target: 6,
@@ -173,7 +206,7 @@ mod tests {
                 min_tokens: 1,
                 hard_max: 8,
             },
-        );
+        ));
         assert!(chunks.len() > 1, "{chunks:?}");
         assert!(chunks[0].contains(' '));
     }
@@ -185,7 +218,7 @@ mod tests {
 
     #[test]
     fn the_tail_keeps_one_piece_even_when_it_is_over_budget() {
-        let buf = vec!["a longer paragraph than the overlap allows".to_string()];
+        let buf = vec![("a longer paragraph than the overlap allows".to_string(), 0)];
         assert_eq!(carry_tail(&buf, 1).len(), 1);
     }
 

crates/chunkmd/src/pages.rs

@@ -25,8 +25,10 @@ pub struct Page<'a> {
     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
+    /// Byte offset in the source markdown where `text` begins, used to find
-    /// to find the chapter when there are no page markers to go on.
+    /// the chapter a passage sits under. Exact rather than approximate: it is
+    /// added to a passage's own offset within the page, and a heading missed
+    /// by a few bytes is the wrong chapter for everything after it.
     pub offset: usize,
 }
 
@@ -51,6 +53,12 @@ pub struct Chapter {
 /// 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.
+/// Where `part` begins inside `whole`, which it must be a slice of.
+fn offset_of(whole: &str, part: &str) -> usize {
+    (part.as_ptr() as usize).saturating_sub(whole.as_ptr() as usize)
+}
+
+/// Cut `markdown` at its page markers.
 pub fn split_by_page(markdown: &str) -> Vec<Page<'_>> {
     let markers: Vec<_> = PAGE.find_iter(markdown).collect();
     if markers.is_empty() {
@@ -67,7 +75,7 @@ pub fn split_by_page(markdown: &str) -> Vec<Page<'_>> {
         pages.push(Page {
             number: None,
             text: head,
-            offset: 0,
+            offset: offset_of(markdown, head),
         });
     }
     for (at, marker) in markers.iter().enumerate() {
@@ -80,7 +88,7 @@ pub fn split_by_page(markdown: &str) -> Vec<Page<'_>> {
             pages.push(Page {
                 number: Some(textkit::decimal(&markdown[marker.start()..marker.end()])),
                 text,
-                offset: start,
+                offset: offset_of(markdown, text),
             });
         }
     }