use super::tables::TableShape;
use super::text_extractor::{DocParagraph, ParagraphType};
#[derive(Debug, Clone, Default, PartialEq)]
pub(crate) struct ChunkContext {
pub page_number: Option<usize>,
pub section_heading: Option<String>,
pub section_heading_level: Option<u8>,
pub heading_path: Vec<String>,
pub list_level: Option<u8>,
pub table: Option<TableShape>,
}
impl ChunkContext {
pub fn heading_path_string(&self) -> Option<String> {
(!self.heading_path.is_empty()).then(|| self.heading_path.join(" > "))
}
}
#[derive(Default)]
struct SectionTracker {
stack: Vec<(u8, String)>,
}
impl SectionTracker {
fn observe(&mut self, p: &DocParagraph) {
let ParagraphType::Heading(level) = p.paragraph_type else {
return;
};
let content = p.content.trim();
if content.is_empty() {
return;
}
while self.stack.last().map(|(l, _)| *l >= level).unwrap_or(false) {
self.stack.pop();
}
self.stack.push((level, content.to_string()));
}
fn context_for(&self, p: &DocParagraph) -> ChunkContext {
ChunkContext {
page_number: p.page_index.map(|i| i + 1),
section_heading: self.stack.last().map(|(_, h)| h.clone()),
section_heading_level: self.stack.last().map(|(l, _)| *l),
heading_path: self.stack.iter().map(|(_, h)| h.clone()).collect(),
list_level: p.list_level,
table: p.table,
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct Positioned {
pub paragraph: DocParagraph,
pub context: ChunkContext,
}
pub(crate) fn position(paragraphs: Vec<DocParagraph>) -> Vec<Positioned> {
let mut tracker = SectionTracker::default();
paragraphs
.into_iter()
.map(|paragraph| {
tracker.observe(¶graph);
let context = tracker.context_for(¶graph);
Positioned { paragraph, context }
})
.collect()
}
pub(crate) fn context_of(window: &[Positioned]) -> ChunkContext {
window.first().map(|p| p.context.clone()).unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
fn heading(level: u8, text: &str) -> DocParagraph {
DocParagraph::plain(text.to_string(), ParagraphType::Heading(level), None)
}
fn body(text: &str) -> DocParagraph {
DocParagraph::plain(text.to_string(), ParagraphType::Normal, None)
}
#[test]
fn a_nested_heading_extends_the_breadcrumb() {
let positioned = position(vec![
heading(1, "Procedure"),
heading(2, "Review Stage"),
body("Instructions about final paper submissions."),
]);
assert_eq!(
positioned[2].context.heading_path_string().as_deref(),
Some("Procedure > Review Stage")
);
assert_eq!(
positioned[2].context.section_heading.as_deref(),
Some("Review Stage")
);
assert_eq!(positioned[2].context.section_heading_level, Some(2));
}
#[test]
fn a_heading_is_its_own_section() {
let positioned = position(vec![heading(1, "Procedure"), heading(2, "Review Stage")]);
assert_eq!(
positioned[1].context.heading_path,
vec!["Procedure", "Review Stage"]
);
}
#[test]
fn a_sibling_heading_pops_the_deeper_ones() {
let positioned = position(vec![
heading(1, "First"),
heading(2, "Sub"),
heading(3, "Subsub"),
heading(1, "Second"),
body("text"),
]);
assert_eq!(positioned[4].context.heading_path, vec!["Second"]);
}
#[test]
fn a_document_without_headings_has_no_breadcrumb() {
let positioned = position(vec![body("just text")]);
assert_eq!(positioned[0].context.heading_path_string(), None);
assert_eq!(positioned[0].context.section_heading, None);
}
}