diff --git a/Cargo.toml b/Cargo.toml index 22ae989..eea0ced 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ repository = "https://github.com/firecrawl/pdf-inspector" [dependencies] # PDF parsing -lopdf = { git = "https://github.com/firecrawl/lopdf", branch = "firecrawl/zlib-checksum-encrypted", features = ["rayon"] } +lopdf = { git = "https://github.com/J-F-Liu/lopdf", rev = "845cd3d4648d9cefb7b5def5fb387df61ba3a0e5", features = ["rayon"] } # Error handling thiserror = "2.0" diff --git a/src/extractor/content_stream.rs b/src/extractor/content_stream.rs index 2ce06e6..ff4a3c1 100644 --- a/src/extractor/content_stream.rs +++ b/src/extractor/content_stream.rs @@ -124,10 +124,18 @@ pub(crate) fn extract_page_text_items( let mut line_matrix = [1.0f32, 0.0, 0.0, 1.0, 0.0, 0.0]; let mut in_text_block = false; - // Marked content (ActualText) tracking - let mut marked_content_stack: Vec> = Vec::new(); + // Marked content tracking: (ActualText, MCID) per nesting level + struct MarkedContentEntry { + actual_text: Option, + mcid: Option, + } + let mut marked_content_stack: Vec = Vec::new(); let mut suppress_glyph_extraction = false; let mut actual_text_start_tm: Option<[f32; 6]> = None; // text matrix at BDC entry + /// Get the innermost MCID from the marked content stack. + fn current_mcid(stack: &[MarkedContentEntry]) -> Option { + stack.iter().rev().find_map(|e| e.mcid) + } for op in &content.operations { trace!("{} {:?}", op.operator, op.operands); @@ -292,6 +300,7 @@ pub(crate) fn extract_page_text_items( is_bold: is_bold_font(base_font), is_italic: is_italic_font(base_font), item_type: ItemType::Text, + mcid: current_mcid(&marked_content_stack), }); } } @@ -440,6 +449,7 @@ pub(crate) fn extract_page_text_items( is_bold: is_bold_font(base_font), is_italic: is_italic_font(base_font), item_type: ItemType::Text, + mcid: current_mcid(&marked_content_stack), }); } } @@ -496,6 +506,7 @@ pub(crate) fn extract_page_text_items( is_bold: is_bold_font(base_font), is_italic: is_italic_font(base_font), item_type: ItemType::Text, + mcid: current_mcid(&marked_content_stack), }); } } @@ -531,11 +542,15 @@ pub(crate) fn extract_page_text_items( } "BMC" => { // Begin Marked Content (no properties) - marked_content_stack.push(None); + marked_content_stack.push(MarkedContentEntry { + actual_text: None, + mcid: None, + }); } "BDC" => { - // Begin Marked Content with properties — extract ActualText + // Begin Marked Content with properties — extract ActualText and MCID let mut actual_text: Option = None; + let mut mcid: Option = None; if op.operands.len() >= 2 { let dict = match &op.operands[1] { Object::Dictionary(d) => Some(d.clone()), @@ -549,47 +564,56 @@ pub(crate) fn extract_page_text_items( _ => None, }; } + if let Ok(Object::Integer(id)) = d.get(b"MCID") { + mcid = Some(*id); + } } } if actual_text.is_some() { suppress_glyph_extraction = true; actual_text_start_tm = Some(text_matrix); } - marked_content_stack.push(actual_text); + marked_content_stack.push(MarkedContentEntry { actual_text, mcid }); } "EMC" => { // End Marked Content — emit ActualText item with correct width - if let Some(Some(at)) = marked_content_stack.pop() { - // Compute width from text matrix advancement during BDC..EMC - if let Some(start_tm) = actual_text_start_tm.take() { - let combined = multiply_matrices(&start_tm, &ctm); - let rendered_size = effective_font_size(current_font_size, &combined); - let (x, y) = (combined[4], combined[5]); - // Width in device space from text matrix delta - let delta_ts = text_matrix[4] - start_tm[4]; - let scale_x = start_tm[0] * ctm[0] + start_tm[1] * ctm[2]; - let width = (delta_ts * scale_x).abs(); - if !at.trim().is_empty() { - let base_font = font_base_names - .get(¤t_font) - .map(|s| s.as_str()) - .unwrap_or(¤t_font); - items.push(TextItem { - text: expand_ligatures(&at), - x, - y, - width, - height: rendered_size, - font: current_font.clone(), - font_size: rendered_size, - page: page_num, - is_bold: is_bold_font(base_font), - is_italic: is_italic_font(base_font), - item_type: ItemType::Text, - }); + if let Some(entry) = marked_content_stack.pop() { + if let Some(at) = entry.actual_text { + // Compute width from text matrix advancement during BDC..EMC + if let Some(start_tm) = actual_text_start_tm.take() { + let combined = multiply_matrices(&start_tm, &ctm); + let rendered_size = effective_font_size(current_font_size, &combined); + let (x, y) = (combined[4], combined[5]); + // Width in device space from text matrix delta + let delta_ts = text_matrix[4] - start_tm[4]; + let scale_x = start_tm[0] * ctm[0] + start_tm[1] * ctm[2]; + let width = (delta_ts * scale_x).abs(); + if !at.trim().is_empty() { + let base_font = font_base_names + .get(¤t_font) + .map(|s| s.as_str()) + .unwrap_or(¤t_font); + items.push(TextItem { + text: expand_ligatures(&at), + x, + y, + width, + height: rendered_size, + font: current_font.clone(), + font_size: rendered_size, + page: page_num, + is_bold: is_bold_font(base_font), + is_italic: is_italic_font(base_font), + item_type: ItemType::Text, + mcid: entry + .mcid + .or_else(|| current_mcid(&marked_content_stack)), + }); + } } + suppress_glyph_extraction = + marked_content_stack.iter().any(|e| e.actual_text.is_some()); } - suppress_glyph_extraction = marked_content_stack.iter().any(|a| a.is_some()); } } "re" => { diff --git a/src/extractor/layout.rs b/src/extractor/layout.rs index 471ec81..312c6ca 100644 --- a/src/extractor/layout.rs +++ b/src/extractor/layout.rs @@ -796,6 +796,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, } } diff --git a/src/extractor/links.rs b/src/extractor/links.rs index f7e9090..74a1fae 100644 --- a/src/extractor/links.rs +++ b/src/extractor/links.rs @@ -79,6 +79,7 @@ pub fn extract_page_links(doc: &Document, page_id: ObjectId, page_num: u32) -> V is_bold: false, is_italic: false, item_type: ItemType::Link(url), + mcid: None, }); } } @@ -316,5 +317,6 @@ pub(crate) fn walk_form_fields( is_bold: false, is_italic: false, item_type: ItemType::FormField, + mcid: None, }); } diff --git a/src/extractor/mod.rs b/src/extractor/mod.rs index 767506f..2b85eed 100644 --- a/src/extractor/mod.rs +++ b/src/extractor/mod.rs @@ -319,6 +319,7 @@ pub(crate) fn merge_text_items(items: Vec) -> Vec { is_bold: first.is_bold, is_italic: first.is_italic, item_type: first.item_type.clone(), + mcid: first.mcid, }); i = j; @@ -359,6 +360,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "World".into(), @@ -372,6 +374,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "Next line".into(), @@ -385,6 +388,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, ]; @@ -438,6 +442,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "Prague".into(), @@ -451,6 +456,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "Rules".into(), @@ -464,6 +470,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, ]; @@ -488,6 +495,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "A".into(), @@ -501,6 +509,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "V".into(), @@ -514,6 +523,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, ]; @@ -540,6 +550,7 @@ mod tests { is_bold: true, is_italic: false, item_type: ItemType::Text, + mcid: None, } } @@ -573,6 +584,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, } } @@ -607,6 +619,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "履行義務".into(), @@ -620,6 +633,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "を識別す".into(), @@ -633,6 +647,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, ]; @@ -654,6 +669,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, } } @@ -765,6 +781,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "\u{05D1}".into(), // bet at x=200 (rightmost) @@ -778,6 +795,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, ]; sort_line_items(&mut items); @@ -801,6 +819,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, TextItem { text: "World".into(), @@ -814,6 +833,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }, ]; sort_line_items(&mut items); @@ -853,6 +873,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }], }; @@ -896,6 +917,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }], }; @@ -939,6 +961,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, }], }; diff --git a/src/extractor/xobjects.rs b/src/extractor/xobjects.rs index c9e5da1..e011b88 100644 --- a/src/extractor/xobjects.rs +++ b/src/extractor/xobjects.rs @@ -408,6 +408,7 @@ fn extract_form_xobject_text_inner( is_bold: is_bold_font(base_font), is_italic: is_italic_font(base_font), item_type: ItemType::Text, + mcid: None, }); } } @@ -549,6 +550,7 @@ fn extract_form_xobject_text_inner( is_bold: is_bold_font(base_font), is_italic: is_italic_font(base_font), item_type: ItemType::Text, + mcid: None, }); } } diff --git a/src/lib.rs b/src/lib.rs index 7011967..4acabe5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,6 +28,7 @@ pub mod extractor; pub mod glyph_names; pub mod markdown; pub mod process_mode; +pub mod structure_tree; pub mod tables; pub mod text_utils; pub mod tounicode; @@ -273,20 +274,21 @@ pub fn process_pdf_mem_with_config( /// are combined here, but lopdf loads the full doc in `load()` so we extract /// page count from it directly to avoid the metadata-only round-trip. fn load_document_from_path>(path: P) -> Result<(Document, u32), PdfError> { - let doc = match Document::load(&path) { - Ok(d) => d, - Err(ref e) if is_encrypted_lopdf_error(e) => Document::load_with_password(&path, "")?, - Err(e) => return Err(e.into()), - }; - let page_count = doc.get_pages().len() as u32; - Ok((doc, page_count)) + let buffer = std::fs::read(&path)?; + load_document_from_mem(&buffer) } /// Load a PDF from a memory buffer. fn load_document_from_mem(buffer: &[u8]) -> Result<(Document, u32), PdfError> { - let doc = match Document::load_mem(buffer) { + // Fix malformed struct element names before parsing. Some PDF generators + // write bare names (/S Code) instead of proper PDF names (/S /Code), which + // causes lopdf to silently drop the entire object. + let fixed = structure_tree::fix_bare_struct_names(buffer); + let buf = fixed.as_ref(); + + let doc = match Document::load_mem(buf) { Ok(d) => d, - Err(ref e) if is_encrypted_lopdf_error(e) => Document::load_mem_with_password(buffer, "")?, + Err(ref e) if is_encrypted_lopdf_error(e) => Document::load_mem_with_password(buf, "")?, Err(e) => return Err(e.into()), }; let page_count = doc.get_pages().len() as u32; @@ -350,6 +352,22 @@ fn process_document( Some(extracted?) }; + // Parse structure tree for tagged PDFs (reuses the loaded document) + let struct_roles = structure_tree::StructTree::from_doc(&doc).and_then(|tree| { + let page_ids = doc.get_pages(); + let roles = tree.mcid_to_roles(&page_ids); + if roles.is_empty() { + None + } else { + log::debug!( + "structure tree: {} pages with MCID roles, {} total MCIDs", + roles.len(), + tree.mcid_count() + ); + Some(roles) + } + }); + let (markdown, layout, has_encoding_issues) = match extracted { Some(((items, rects, lines), page_thresholds)) => { let layout = compute_layout_complexity(&items, &rects, &lines); @@ -363,6 +381,7 @@ fn process_document( &rects, &lines, &page_thresholds, + struct_roles.as_ref(), )) }; diff --git a/src/markdown/convert.rs b/src/markdown/convert.rs index 42fb107..1fef66f 100644 --- a/src/markdown/convert.rs +++ b/src/markdown/convert.rs @@ -2,6 +2,7 @@ use std::collections::HashSet; +use crate::structure_tree::StructRole; use crate::types::TextLine; use super::analysis::{ @@ -13,6 +14,50 @@ use super::postprocess::clean_markdown; use super::preprocess::{merge_drop_caps, merge_heading_lines}; use super::MarkdownOptions; +/// Resolve the dominant structure role for a text line by looking up its items' MCIDs. +/// +/// Returns the first non-container role found (skipping Document/Part/Sect/Div/NonStruct/Span). +/// These wrapper roles don't carry useful semantic info for markdown generation. +fn resolve_line_struct_role( + line: &TextLine, + struct_roles: &std::collections::HashMap>, +) -> Option { + let page_roles = struct_roles.get(&line.page)?; + for item in &line.items { + if let Some(mcid) = item.mcid { + if let Some(role) = page_roles.get(&mcid) { + match role { + // Skip container/wrapper roles — not useful for line classification + StructRole::Document + | StructRole::Part + | StructRole::Art + | StructRole::Sect + | StructRole::Div + | StructRole::NonStruct + | StructRole::Span + | StructRole::Private => continue, + _ => return Some(role.clone()), + } + } + } + } + None +} + +/// Map a StructRole heading variant to a markdown heading level (1–6). +fn struct_role_heading_level(role: &StructRole) -> Option { + match role { + StructRole::H => Some(1), // Generic heading → H1 + StructRole::H1 => Some(1), + StructRole::H2 => Some(2), + StructRole::H3 => Some(3), + StructRole::H4 => Some(4), + StructRole::H5 => Some(5), + StructRole::H6 => Some(6), + _ => None, + } +} + /// Merge continuation tables that span across page breaks. /// /// When consecutive pages each have exactly one table with the same number of columns @@ -182,6 +227,9 @@ pub(super) fn to_markdown_from_lines_with_tables_and_images( page_tables: std::collections::HashMap>, page_images: std::collections::HashMap>, band_split_pages: &HashSet, + struct_roles: Option< + &std::collections::HashMap>, + >, ) -> String { if lines.is_empty() && page_tables.is_empty() && page_images.is_empty() { return String::new(); @@ -215,6 +263,7 @@ pub(super) fn to_markdown_from_lines_with_tables_and_images( let mut in_list = false; let mut in_paragraph = false; let mut last_list_x: Option = None; + let mut in_code_block = false; let mut prev_had_dot_leaders = false; let mut inserted_tables: HashSet<(u32, usize)> = HashSet::new(); let mut inserted_images: HashSet<(u32, usize)> = HashSet::new(); @@ -233,6 +282,10 @@ pub(super) fn to_markdown_from_lines_with_tables_and_images( if line.page != current_page { // Flush current page's remaining tables and images if current_page > 0 { + if in_code_block { + output.push_str("```\n"); + in_code_block = false; + } flush_page_tables_and_images( current_page, &page_tables, @@ -352,7 +405,25 @@ pub(super) fn to_markdown_from_lines_with_tables_and_images( // Detect figure/table captions and source citations // These should be on their own line followed by a paragraph break - if is_caption_line(plain_trimmed) { + let struct_role = struct_roles.and_then(|roles| resolve_line_struct_role(&line, roles)); + + // Determine if this line is code (struct-tree or font-based) for block accumulation + let is_code_line = struct_role + .as_ref() + .is_some_and(|r| matches!(r, StructRole::Code)) + || (options.detect_code && line.items.iter().any(|i| is_monospace_font(&i.font))); + + // Close code block when transitioning to non-code + if in_code_block && !is_code_line { + output.push_str("```\n"); + in_code_block = false; + } + + if struct_role + .as_ref() + .is_some_and(|r| matches!(r, StructRole::Caption)) + || is_caption_line(plain_trimmed) + { if in_paragraph { output.push_str("\n\n"); in_paragraph = false; @@ -362,27 +433,48 @@ pub(super) fn to_markdown_from_lines_with_tables_and_images( continue; } - // Detect headers by font size - // Note: Headers typically shouldn't have bold markers since they're already emphasized - // Skip very short text (drop caps/labels) and very long text (body paragraphs) - if options.detect_headers + // Detect headers: structure-tree headings win, then font-size heuristics. + // Structure roles ADD headings (e.g. same-size text tagged H2) but do NOT + // suppress headings that the font heuristic would detect (some tagged PDFs + // mark obvious headings as P or Span). + let struct_heading = struct_role.as_ref().and_then(struct_role_heading_level); + let heuristic_heading = if options.detect_headers && plain_trimmed.len() > 3 && plain_trimmed.split_whitespace().count() <= 15 { let line_font_size = line.items.first().map(|i| i.font_size).unwrap_or(base_size); - if let Some(header_level) = - detect_header_level(line_font_size, base_size, &heading_tiers) - { - if in_paragraph { - output.push_str("\n\n"); - in_paragraph = false; - } - let prefix = "#".repeat(header_level); - // Use plain text for headers to avoid redundant formatting - output.push_str(&format!("{} {}\n\n", prefix, plain_trimmed)); - in_list = false; - continue; + detect_header_level(line_font_size, base_size, &heading_tiers) + } else { + None + }; + + if let Some(level) = struct_heading.or(heuristic_heading) { + if in_paragraph { + output.push_str("\n\n"); + in_paragraph = false; } + let prefix = "#".repeat(level); + // Use plain text for headers to avoid redundant formatting + output.push_str(&format!("{} {}\n\n", prefix, plain_trimmed)); + in_list = false; + continue; + } + + // Structure-tree list item (LI only — LBody is a continuation, not a new item) + if struct_role + .as_ref() + .is_some_and(|r| matches!(r, StructRole::LI)) + && !is_list_item(plain_trimmed) + { + if in_paragraph { + output.push_str("\n\n"); + in_paragraph = false; + } + output.push_str(&format!("- {}", trimmed)); + output.push('\n'); + in_list = true; + last_list_x = line.items.first().map(|i| i.x); + continue; } // Detect list items @@ -428,18 +520,32 @@ pub(super) fn to_markdown_from_lines_with_tables_and_images( } } - // Detect code blocks by font - if options.detect_code { - let is_mono = line.items.iter().any(|i| is_monospace_font(&i.font)); - if is_mono { - if in_paragraph { - output.push_str("\n\n"); - in_paragraph = false; - } - // Use plain text for code blocks - output.push_str(&format!("```\n{}\n```\n", plain_trimmed)); - continue; + // Structure-tree block quote + if struct_role + .as_ref() + .is_some_and(|r| matches!(r, StructRole::BlockQuote)) + { + if in_paragraph { + output.push_str("\n\n"); + in_paragraph = false; } + output.push_str(&format!("> {}\n", trimmed)); + continue; + } + + // Code block accumulation (struct-tree Code role or monospace font) + if is_code_line { + if in_paragraph { + output.push_str("\n\n"); + in_paragraph = false; + } + if !in_code_block { + output.push_str("```\n"); + in_code_block = true; + } + output.push_str(plain_trimmed); + output.push('\n'); + continue; } // Regular text - join lines within same paragraph with space @@ -456,6 +562,11 @@ pub(super) fn to_markdown_from_lines_with_tables_and_images( prev_had_dot_leaders = cur_dot_leaders; } + // Close any trailing code block + if in_code_block { + output.push_str("```\n"); + } + // Flush current page and any remaining pages with tables/images // (handles table-only pages after the last text line, and trailing image-only pages) flush_page_tables_and_images( @@ -680,3 +791,271 @@ pub fn to_markdown_from_lines(lines: Vec, options: MarkdownOptions) -> // Clean up and post-process clean_markdown(output, &options) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::structure_tree::StructRole; + use crate::types::TextItem; + use std::collections::HashMap; + + fn make_item(text: &str, page: u32, mcid: Option) -> TextItem { + TextItem { + text: text.to_string(), + x: 72.0, + y: 700.0, + width: 100.0, + height: 12.0, + font: "Helvetica".to_string(), + font_size: 12.0, + page, + is_bold: false, + is_italic: false, + item_type: crate::types::ItemType::Text, + mcid, + } + } + + fn make_line(items: Vec) -> TextLine { + let y = items.first().map(|i| i.y).unwrap_or(0.0); + let page = items.first().map(|i| i.page).unwrap_or(1); + TextLine { + items, + y, + page, + adaptive_threshold: 0.10, + } + } + + #[test] + fn test_struct_role_heading() { + let lines = vec![ + make_line(vec![make_item("Introduction", 1, Some(0))]), + make_line(vec![{ + let mut item = make_item("Body text here.", 1, Some(1)); + item.y = 680.0; + item + }]), + ]; + + let mut page_roles = HashMap::new(); + page_roles.insert(0i64, StructRole::H1); + page_roles.insert(1i64, StructRole::P); + let mut roles = HashMap::new(); + roles.insert(1u32, page_roles); + + let md = to_markdown_from_lines_with_tables_and_images( + lines, + MarkdownOptions::default(), + HashMap::new(), + HashMap::new(), + &std::collections::HashSet::new(), + Some(&roles), + ); + + assert!( + md.contains("# Introduction"), + "Should have H1 heading: {md}" + ); + assert!( + md.contains("Body text here."), + "Should have body text: {md}" + ); + } + + #[test] + fn test_struct_role_list_item() { + let lines = vec![make_line(vec![make_item("First item", 1, Some(0))])]; + + let mut page_roles = HashMap::new(); + page_roles.insert(0i64, StructRole::LI); + let mut roles = HashMap::new(); + roles.insert(1u32, page_roles); + + let md = to_markdown_from_lines_with_tables_and_images( + lines, + MarkdownOptions::default(), + HashMap::new(), + HashMap::new(), + &std::collections::HashSet::new(), + Some(&roles), + ); + + assert!( + md.contains("- First item"), + "Should format as list item: {md}" + ); + } + + #[test] + fn test_struct_role_blockquote() { + let lines = vec![make_line(vec![make_item("Quoted text", 1, Some(0))])]; + + let mut page_roles = HashMap::new(); + page_roles.insert(0i64, StructRole::BlockQuote); + let mut roles = HashMap::new(); + roles.insert(1u32, page_roles); + + let md = to_markdown_from_lines_with_tables_and_images( + lines, + MarkdownOptions::default(), + HashMap::new(), + HashMap::new(), + &std::collections::HashSet::new(), + Some(&roles), + ); + + assert!( + md.contains("> Quoted text"), + "Should format as blockquote: {md}" + ); + } + + #[test] + fn test_struct_role_heading_levels() { + let mcids = vec![ + (StructRole::H1, "Title"), + (StructRole::H2, "Section"), + (StructRole::H3, "Subsection"), + ]; + + let mut lines = Vec::new(); + let mut page_roles = HashMap::new(); + for (i, (role, text)) in mcids.iter().enumerate() { + let mut item = make_item(text, 1, Some(i as i64)); + item.y = 700.0 - (i as f32 * 30.0); + lines.push(make_line(vec![item])); + page_roles.insert(i as i64, role.clone()); + } + + let mut roles = HashMap::new(); + roles.insert(1u32, page_roles); + + let md = to_markdown_from_lines_with_tables_and_images( + lines, + MarkdownOptions::default(), + HashMap::new(), + HashMap::new(), + &std::collections::HashSet::new(), + Some(&roles), + ); + + assert!(md.contains("# Title"), "H1 → #: {md}"); + assert!(md.contains("## Section"), "H2 → ##: {md}"); + assert!(md.contains("### Subsection"), "H3 → ###: {md}"); + } + + #[test] + fn test_no_struct_roles_falls_back_to_heuristics() { + let mut item = make_item("Big Title", 1, None); + item.font_size = 24.0; + item.height = 24.0; + + let lines = vec![ + make_line(vec![item]), + make_line(vec![{ + let mut body = make_item("Normal body text.", 1, None); + body.y = 660.0; + body + }]), + ]; + + let md = to_markdown_from_lines_with_tables_and_images( + lines, + MarkdownOptions::default(), + HashMap::new(), + HashMap::new(), + &std::collections::HashSet::new(), + None, + ); + + assert!( + md.contains("# Big Title"), + "Font heuristic should detect heading: {md}" + ); + } + + #[test] + fn test_resolve_line_struct_role_skips_containers() { + let mut page_roles = HashMap::new(); + page_roles.insert(0i64, StructRole::Div); + page_roles.insert(1i64, StructRole::H2); + let mut roles = HashMap::new(); + roles.insert(1u32, page_roles); + + let line = make_line(vec![ + make_item("Part ", 1, Some(0)), + make_item("Title", 1, Some(1)), + ]); + + let role = resolve_line_struct_role(&line, &roles); + assert_eq!(role, Some(StructRole::H2)); + } + + #[test] + fn test_struct_role_code() { + let lines = vec![make_line(vec![make_item("fn main() {}", 1, Some(0))])]; + + let mut page_roles = HashMap::new(); + page_roles.insert(0i64, StructRole::Code); + let mut roles = HashMap::new(); + roles.insert(1u32, page_roles); + + let md = to_markdown_from_lines_with_tables_and_images( + lines, + MarkdownOptions::default(), + HashMap::new(), + HashMap::new(), + &std::collections::HashSet::new(), + Some(&roles), + ); + + assert!( + md.contains("```\nfn main() {}\n```"), + "Should format as code block: {md}" + ); + } + + #[test] + fn test_struct_role_code_multiline_accumulation() { + let mut line1 = make_item("fn main() {", 1, Some(0)); + line1.y = 700.0; + let mut line2 = make_item(" println!(\"hello\");", 1, Some(1)); + line2.y = 688.0; + let mut line3 = make_item("}", 1, Some(2)); + line3.y = 676.0; + + let lines = vec![ + make_line(vec![line1]), + make_line(vec![line2]), + make_line(vec![line3]), + ]; + + let mut page_roles = HashMap::new(); + page_roles.insert(0i64, StructRole::Code); + page_roles.insert(1i64, StructRole::Code); + page_roles.insert(2i64, StructRole::Code); + let mut roles = HashMap::new(); + roles.insert(1u32, page_roles); + + let md = to_markdown_from_lines_with_tables_and_images( + lines, + MarkdownOptions::default(), + HashMap::new(), + HashMap::new(), + &std::collections::HashSet::new(), + Some(&roles), + ); + + // Should produce a single fenced block, not three separate ones + assert!( + md.contains("```\nfn main() {\nprintln!(\"hello\");\n}\n```"), + "Should accumulate consecutive code lines into one block: {md}" + ); + // Should NOT have adjacent fences + assert!( + !md.contains("```\n```"), + "Should not have adjacent close/open fences: {md}" + ); + } +} diff --git a/src/markdown/mod.rs b/src/markdown/mod.rs index 8a8e241..ed0cbd5 100644 --- a/src/markdown/mod.rs +++ b/src/markdown/mod.rs @@ -452,7 +452,7 @@ pub fn to_markdown_from_items_with_rects( options: MarkdownOptions, rects: &[crate::types::PdfRect], ) -> String { - to_markdown_from_items_with_rects_and_lines(items, options, rects, &[], &HashMap::new()) + to_markdown_from_items_with_rects_and_lines(items, options, rects, &[], &HashMap::new(), None) } /// Convert positioned text items to markdown, using rectangles and line segments for table detection. @@ -465,6 +465,7 @@ pub(crate) fn to_markdown_from_items_with_rects_and_lines( rects: &[crate::types::PdfRect], pdf_lines: &[crate::types::PdfLine], page_thresholds: &HashMap, + struct_roles: Option<&HashMap>>, ) -> String { use crate::tables::{ detect_tables, detect_tables_from_lines, detect_tables_from_rects, table_to_markdown, @@ -757,6 +758,40 @@ pub(crate) fn to_markdown_from_items_with_rects_and_lines( } } + // Check structure tree coverage on ALL text items (before table filtering) + // to decide whether to use structure-aware markdown generation. + let struct_roles_coverage_ok = struct_roles.is_some_and(|roles| { + let total = text_items.len(); + if total == 0 { + return false; + } + let tagged = text_items + .iter() + .filter(|item| { + item.mcid + .and_then(|mcid| { + roles + .get(&item.page) + .and_then(|page_roles| page_roles.get(&mcid)) + }) + .is_some() + }) + .count(); + let coverage = tagged as f32 / total as f32; + log::debug!( + "structure tree coverage: {}/{} items ({:.0}%)", + tagged, + total, + coverage * 100.0 + ); + coverage >= 0.5 + }); + let effective_struct_roles = if struct_roles_coverage_ok { + struct_roles + } else { + None + }; + // Filter out table items and process the rest let non_table_items: Vec = text_items .into_iter() @@ -841,6 +876,7 @@ pub(crate) fn to_markdown_from_items_with_rects_and_lines( page_tables, page_images, &band_split_page_set, + effective_struct_roles, ) } @@ -926,6 +962,7 @@ mod tests { is_bold: false, is_italic: false, item_type: crate::types::ItemType::Text, + mcid: None, } } diff --git a/src/structure_tree.rs b/src/structure_tree.rs new file mode 100644 index 0000000..d7e440a --- /dev/null +++ b/src/structure_tree.rs @@ -0,0 +1,1035 @@ +//! Tagged PDF structure tree parser. +//! +//! Reads the `/StructTreeRoot` from the document catalog and builds an +//! in-memory tree of [`StructElement`] nodes. Each leaf maps back to +//! content-stream marked content via MCID (Marked Content ID), which lets +//! downstream code attach semantic roles (heading, paragraph, table cell, +//! list item, …) to extracted [`TextItem`]s. + +use log::debug; +use lopdf::{Document, Object, ObjectId}; +use std::borrow::Cow; +use std::collections::HashMap; + +// ─── Standard structure types ──────────────────────────────────────── + +/// Standard PDF structure element types (ISO 32000-1, Table 333–340). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum StructRole { + Document, + Part, + Art, + Sect, + Div, + BlockQuote, + Caption, + TOC, + TOCI, + Index, + NonStruct, + Private, + // Heading & paragraph + H, + H1, + H2, + H3, + H4, + H5, + H6, + P, + // List + L, + LI, + Lbl, + LBody, + // Table + Table, + TR, + TH, + TD, + THead, + TBody, + TFoot, + // Inline + Span, + Quote, + Note, + Reference, + BibEntry, + Code, + Link, + Annot, + // Illustration + Figure, + Formula, + Form, + // Ruby / Warichu (CJK) + Ruby, + RB, + RT, + RP, + Warichu, + WT, + WP, + // Fallback + Other(String), +} + +impl StructRole { + fn from_name(name: &str) -> Self { + match name { + "Document" => Self::Document, + "Part" => Self::Part, + "Art" => Self::Art, + "Sect" => Self::Sect, + "Div" => Self::Div, + "BlockQuote" => Self::BlockQuote, + "Caption" => Self::Caption, + "TOC" => Self::TOC, + "TOCI" => Self::TOCI, + "Index" => Self::Index, + "NonStruct" => Self::NonStruct, + "Private" => Self::Private, + "H" => Self::H, + "H1" => Self::H1, + "H2" => Self::H2, + "H3" => Self::H3, + "H4" => Self::H4, + "H5" => Self::H5, + "H6" => Self::H6, + "P" => Self::P, + "L" => Self::L, + "LI" => Self::LI, + "Lbl" => Self::Lbl, + "LBody" => Self::LBody, + "Table" => Self::Table, + "TR" => Self::TR, + "TH" => Self::TH, + "TD" => Self::TD, + "THead" => Self::THead, + "TBody" => Self::TBody, + "TFoot" => Self::TFoot, + "Span" => Self::Span, + "Quote" => Self::Quote, + "Note" => Self::Note, + "Reference" => Self::Reference, + "BibEntry" => Self::BibEntry, + "Code" => Self::Code, + "Link" => Self::Link, + "Annot" => Self::Annot, + "Figure" => Self::Figure, + "Formula" => Self::Formula, + "Form" => Self::Form, + "Ruby" => Self::Ruby, + "RB" => Self::RB, + "RT" => Self::RT, + "RP" => Self::RP, + "Warichu" => Self::Warichu, + "WT" => Self::WT, + "WP" => Self::WP, + other => Self::Other(other.to_string()), + } + } + + /// Resolve a possibly-custom tag name through a role map. + fn from_name_with_role_map(name: &str, role_map: &HashMap) -> Self { + // Follow role map chain (max 8 hops to avoid cycles) + let mut current = name.to_string(); + for _ in 0..8 { + let role = Self::from_name(¤t); + if !matches!(role, Self::Other(_)) { + return role; + } + if let Some(mapped) = role_map.get(current.as_str()) { + current = mapped.clone(); + } else { + return role; + } + } + Self::Other(name.to_string()) + } +} + +// ─── Marked content reference ──────────────────────────────────────── + +/// A leaf reference linking a structure element to content-stream content. +#[derive(Debug, Clone)] +pub struct MarkedContentRef { + /// The Marked Content ID used in the content stream's `BDC`/`BMC`. + pub mcid: i64, + /// Page ObjectId this content belongs to (from `/Pg` key). + pub page_id: Option, +} + +// ─── Structure element ─────────────────────────────────────────────── + +/// A node in the PDF structure tree. +#[derive(Debug, Clone)] +pub struct StructElement { + /// Semantic role (H1, P, Table, TD, …). + pub role: StructRole, + /// Alternative text for figures / illustrations. + pub alt_text: Option, + /// Actual text override (e.g. for ligatures). + pub actual_text: Option, + /// Language override (e.g. "en-US"). + pub lang: Option, + /// Direct marked-content references (leaf content). + pub content_refs: Vec, + /// Child structure elements. + pub children: Vec, +} + +// ─── Structure tree (top level) ────────────────────────────────────── + +/// Parsed PDF structure tree. +/// +/// Built from `/StructTreeRoot` in the document catalog. Use +/// [`StructTree::from_doc`] to parse, then [`StructTree::mcid_to_roles`] +/// to get per-page MCID → role lookup tables. +#[derive(Debug, Clone)] +pub struct StructTree { + /// Root children (the top-level structure elements). + pub children: Vec, +} + +impl StructTree { + /// Attempt to parse the structure tree from a PDF document. + /// + /// Returns `None` if the PDF is not tagged (no `/StructTreeRoot`). + pub fn from_doc(doc: &Document) -> Option { + let catalog = doc.catalog().ok()?; + let struct_root_obj = catalog.get(b"StructTreeRoot").ok()?; + let struct_root = resolve_dict(doc, struct_root_obj)?; + + // Parse role map: custom tag → standard tag + let role_map = parse_role_map(doc, struct_root); + debug!("structure tree: {} role map entries", role_map.len()); + + // Parse child elements from /K + let children = parse_kids(doc, struct_root, &role_map, None, 0); + debug!("structure tree: {} top-level elements", children.len()); + + if children.is_empty() { + return None; + } + + Some(StructTree { children }) + } + + /// Build per-page MCID → StructRole lookup. + /// + /// Returns a map: page_number (1-indexed) → (MCID → StructRole). + /// The `page_ids` map should come from `doc.get_pages()`. + pub fn mcid_to_roles( + &self, + page_ids: &std::collections::BTreeMap, + ) -> HashMap> { + // Invert: ObjectId → page number + let obj_to_page: HashMap = + page_ids.iter().map(|(&num, &id)| (id, num)).collect(); + + let mut result: HashMap> = HashMap::new(); + self.collect_mcid_roles(&self.children, &obj_to_page, &mut result); + result + } + + fn collect_mcid_roles( + &self, + elements: &[StructElement], + obj_to_page: &HashMap, + result: &mut HashMap>, + ) { + for elem in elements { + for mcref in &elem.content_refs { + if let Some(page_id) = mcref.page_id { + if let Some(&page_num) = obj_to_page.get(&page_id) { + result + .entry(page_num) + .or_default() + .insert(mcref.mcid, elem.role.clone()); + } + } + } + self.collect_mcid_roles(&elem.children, obj_to_page, result); + } + } + + /// Count total marked-content references across the tree. + pub fn mcid_count(&self) -> usize { + fn count(elements: &[StructElement]) -> usize { + elements + .iter() + .map(|e| e.content_refs.len() + count(&e.children)) + .sum() + } + count(&self.children) + } + + /// Build a flat list of structure elements with their roles and MCIDs, + /// preserving document order. Useful for structure-aware markdown generation. + pub fn flatten(&self) -> Vec { + let mut out = Vec::new(); + flatten_recursive(&self.children, &mut out, 0); + out + } +} + +/// A flattened view of a structure element for linear traversal. +#[derive(Debug, Clone)] +pub struct FlatStructElement { + /// Semantic role. + pub role: StructRole, + /// Nesting depth (0 = top-level). + pub depth: usize, + /// Alt text (figures). + pub alt_text: Option, + /// Direct MCIDs with page ObjectIds. + pub content_refs: Vec, + /// Number of child elements (in the original tree). + pub child_count: usize, +} + +fn flatten_recursive(elements: &[StructElement], out: &mut Vec, depth: usize) { + for elem in elements { + out.push(FlatStructElement { + role: elem.role.clone(), + depth, + alt_text: elem.alt_text.clone(), + content_refs: elem.content_refs.clone(), + child_count: elem.children.len(), + }); + flatten_recursive(&elem.children, out, depth + 1); + } +} + +// ─── Parsing helpers ───────────────────────────────────────────────── + +/// Parse the `/RoleMap` dictionary (custom tag → standard tag). +fn parse_role_map(doc: &Document, struct_root: &lopdf::Dictionary) -> HashMap { + let mut map = HashMap::new(); + let Ok(rm_obj) = struct_root.get(b"RoleMap") else { + return map; + }; + let Some(rm_dict) = resolve_dict(doc, rm_obj) else { + return map; + }; + for (key, val) in rm_dict.iter() { + let key_str = String::from_utf8_lossy(key).to_string(); + if let Ok(name) = val.as_name() { + let val_str = String::from_utf8_lossy(name).to_string(); + map.insert(key_str, val_str); + } + } + map +} + +/// Max recursion depth for structure tree parsing (prevents stack overflow on +/// malformed PDFs). +const MAX_DEPTH: usize = 64; + +/// Parse child elements from a `/K` entry. +fn parse_kids( + doc: &Document, + dict: &lopdf::Dictionary, + role_map: &HashMap, + inherited_page: Option, + depth: usize, +) -> Vec { + if depth >= MAX_DEPTH { + return Vec::new(); + } + + let Ok(k_obj) = dict.get(b"K") else { + return Vec::new(); + }; + + // /Pg on this element (inherited by children) + let page_id = get_page_ref(doc, dict).or(inherited_page); + + match k_obj { + Object::Array(arr) => { + let mut children = Vec::new(); + for item in arr { + let resolved = resolve_obj(doc, item); + parse_kid(doc, resolved, role_map, page_id, depth, &mut children); + } + children + } + other => { + let resolved = resolve_obj(doc, other); + let mut children = Vec::new(); + parse_kid(doc, resolved, role_map, page_id, depth, &mut children); + children + } + } +} + +/// Parse a single child (either a struct element dict or an MCID integer). +fn parse_kid( + doc: &Document, + obj: &Object, + role_map: &HashMap, + inherited_page: Option, + depth: usize, + out: &mut Vec, +) { + match obj { + // Direct MCID integer — create a leaf wrapper + Object::Integer(mcid) => { + // This is a bare MCID at the struct-element level. + // We attach it to the parent element, so we create a wrapper struct element. + // Actually, bare MCIDs inside /K are content refs for the parent, + // not separate child elements. We handle this at the caller level. + // For now, create a minimal Span wrapper. + out.push(StructElement { + role: StructRole::Span, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![MarkedContentRef { + mcid: *mcid, + page_id: inherited_page, + }], + children: Vec::new(), + }); + } + Object::Dictionary(d) => { + parse_struct_element_dict(doc, d, role_map, inherited_page, depth, out); + } + Object::Stream(s) => { + // Some PDFs wrap struct elements in streams (rare) + parse_struct_element_dict(doc, &s.dict, role_map, inherited_page, depth, out); + } + _ => {} + } +} + +/// Parse a dictionary that could be either a struct element or a marked-content +/// reference (MCR) dictionary. +fn parse_struct_element_dict( + doc: &Document, + dict: &lopdf::Dictionary, + role_map: &HashMap, + inherited_page: Option, + depth: usize, + out: &mut Vec, +) { + if depth >= MAX_DEPTH { + return; + } + // Check if this is a marked-content reference dict (has /Type /MCR) + if is_mcr_dict(dict) { + if let Ok(Object::Integer(mcid)) = dict.get(b"MCID") { + let page_id = get_page_ref(doc, dict).or(inherited_page); + out.push(StructElement { + role: StructRole::Span, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![MarkedContentRef { + mcid: *mcid, + page_id, + }], + children: Vec::new(), + }); + } + return; + } + + // Check if this is an object reference dict (has /Type /OBJR) — skip these + if is_objr_dict(dict) { + return; + } + + // It's a struct element — parse its /S (structure type) + let role_name = match dict.get(b"S") { + Ok(s_obj) => { + let resolved = resolve_obj(doc, s_obj); + match resolved.as_name() { + Ok(name) => String::from_utf8_lossy(name).to_string(), + Err(_) => return, + } + } + Err(_) => return, + }; + + let role = StructRole::from_name_with_role_map(&role_name, role_map); + let page_id = get_page_ref(doc, dict).or(inherited_page); + + // Extract optional attributes + let alt_text = get_text_string(dict, b"Alt"); + let actual_text = get_text_string(dict, b"ActualText"); + let lang = get_text_string(dict, b"Lang"); + + // Parse children from /K + let mut content_refs = Vec::new(); + let mut children = Vec::new(); + + if let Ok(k_obj) = dict.get(b"K") { + let k_resolved = resolve_obj(doc, k_obj); + match k_resolved { + Object::Integer(mcid) => { + content_refs.push(MarkedContentRef { + mcid: *mcid, + page_id, + }); + } + Object::Array(arr) => { + for item in arr { + let resolved = resolve_obj(doc, item); + match resolved { + Object::Integer(mcid) => { + content_refs.push(MarkedContentRef { + mcid: *mcid, + page_id, + }); + } + Object::Dictionary(d) => { + if is_mcr_dict(d) { + if let Ok(Object::Integer(mcid)) = d.get(b"MCID") { + let pg = get_page_ref(doc, d).or(page_id); + content_refs.push(MarkedContentRef { + mcid: *mcid, + page_id: pg, + }); + } + } else if is_objr_dict(d) { + // Skip object references + } else { + parse_struct_element_dict( + doc, + d, + role_map, + page_id, + depth + 1, + &mut children, + ); + } + } + Object::Stream(s) => { + parse_struct_element_dict( + doc, + &s.dict, + role_map, + page_id, + depth + 1, + &mut children, + ); + } + _ => {} + } + } + } + Object::Dictionary(d) => { + if is_mcr_dict(d) { + if let Ok(Object::Integer(mcid)) = d.get(b"MCID") { + let pg = get_page_ref(doc, d).or(page_id); + content_refs.push(MarkedContentRef { + mcid: *mcid, + page_id: pg, + }); + } + } else { + parse_struct_element_dict(doc, d, role_map, page_id, depth + 1, &mut children); + } + } + _ => {} + } + } + + out.push(StructElement { + role, + alt_text, + actual_text, + lang, + content_refs, + children, + }); +} + +/// Check if dict has `/Type /MCR`. +fn is_mcr_dict(dict: &lopdf::Dictionary) -> bool { + dict.get(b"Type") + .ok() + .and_then(|o| o.as_name().ok()) + .is_some_and(|n| n == b"MCR") +} + +/// Check if dict has `/Type /OBJR`. +fn is_objr_dict(dict: &lopdf::Dictionary) -> bool { + dict.get(b"Type") + .ok() + .and_then(|o| o.as_name().ok()) + .is_some_and(|n| n == b"OBJR") +} + +/// Get the `/Pg` page reference from a dictionary. +fn get_page_ref(doc: &Document, dict: &lopdf::Dictionary) -> Option { + let pg = dict.get(b"Pg").ok()?; + match pg { + Object::Reference(id) => Some(*id), + _ => { + let resolved = resolve_obj(doc, pg); + if let Object::Reference(id) = resolved { + Some(*id) + } else { + None + } + } + } +} + +/// Extract a text string from a dictionary key (handles PDF text encoding). +fn get_text_string(dict: &lopdf::Dictionary, key: &[u8]) -> Option { + let obj = dict.get(key).ok()?; + match obj { + Object::String(bytes, _) => Some(crate::text_utils::decode_text_string(bytes)), + _ => None, + } +} + +/// Resolve an Object reference, returning the target object. +fn resolve_obj<'a>(doc: &'a Document, obj: &'a Object) -> &'a Object { + match obj { + Object::Reference(id) => doc.get_object(*id).unwrap_or(obj), + _ => obj, + } +} + +/// Resolve an Object to a dictionary (handling references). +fn resolve_dict<'a>(doc: &'a Document, obj: &'a Object) -> Option<&'a lopdf::Dictionary> { + match obj { + Object::Dictionary(d) => Some(d), + Object::Reference(id) => doc.get_dictionary(*id).ok(), + _ => None, + } +} + +// ─── PDF byte pre-processing ──────────────────────────────────────── + +/// Fix malformed structure element `/S` entries in raw PDF bytes. +/// +/// Some PDF generators (notably fpdf2) write bare names like `/S Code` +/// instead of the correct `/S /Code`. lopdf cannot parse dictionaries +/// containing bare tokens, so the entire object is silently dropped. +/// +/// This function scans for the pattern `/S ` inside struct +/// element dictionaries and prepends `/` to make them valid PDF names. +/// Returns `Cow::Borrowed` if no fixes were needed. +pub fn fix_bare_struct_names(buf: &[u8]) -> Cow<'_, [u8]> { + // Quick check: if no StructTreeRoot, nothing to fix + if !contains_bytes(buf, b"/StructTreeRoot") { + return Cow::Borrowed(buf); + } + + // Known struct type names that may appear as bare tokens. + // We only fix names that are valid PDF structure types to avoid + // false positives on arbitrary dictionary values. + const KNOWN_NAMES: &[&[u8]] = &[ + b"Document", + b"Part", + b"Art", + b"Sect", + b"Div", + b"BlockQuote", + b"Caption", + b"TOC", + b"TOCI", + b"Index", + b"NonStruct", + b"Private", + b"H", + b"H1", + b"H2", + b"H3", + b"H4", + b"H5", + b"H6", + b"P", + b"L", + b"LI", + b"Lbl", + b"LBody", + b"Table", + b"TR", + b"TH", + b"TD", + b"THead", + b"TBody", + b"TFoot", + b"Span", + b"Quote", + b"Note", + b"Reference", + b"BibEntry", + b"Code", + b"Link", + b"Annot", + b"Figure", + b"Formula", + b"Form", + b"Ruby", + b"RB", + b"RT", + b"RP", + b"Warichu", + b"WT", + b"WP", + ]; + + let pattern = b"/S "; + let mut result: Option> = None; + let mut pos = 0; + + while pos + pattern.len() < buf.len() { + let Some(idx) = find_bytes(&buf[pos..], pattern).map(|i| i + pos) else { + break; + }; + + let after = idx + pattern.len(); + // Check if the next char is already '/' (correct name) or not + if after < buf.len() && buf[after] == b'/' { + pos = after; + continue; + } + + // Try to match a known bare struct name at this position + let mut matched = false; + for name in KNOWN_NAMES { + let end = after + name.len(); + if end <= buf.len() + && &buf[after..end] == *name + // Must be followed by a delimiter (whitespace, newline, /, >) + && (end >= buf.len() || matches!(buf[end], b'\n' | b'\r' | b' ' | b'/' | b'>')) + { + // Found a bare name — lazily allocate output buffer + let out = result.get_or_insert_with(|| buf[..after].to_vec()); + // Append everything from last position up to the bare name + if out.len() < after { + out.extend_from_slice(&buf[out.len()..after]); + } + out.push(b'/'); + out.extend_from_slice(name); + pos = end; + matched = true; + debug!( + "fix_bare_struct_names: patched /S {} → /S /{}", + String::from_utf8_lossy(name), + String::from_utf8_lossy(name) + ); + break; + } + } + + if !matched { + pos = after; + } + } + + match result { + Some(mut out) => { + // Append remaining bytes + if out.len() < buf.len() { + out.extend_from_slice(&buf[out.len()..]); + } + Cow::Owned(out) + } + None => Cow::Borrowed(buf), + } +} + +fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option { + haystack.windows(needle.len()).position(|w| w == needle) +} + +fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool { + find_bytes(haystack, needle).is_some() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_struct_role_from_name() { + assert_eq!(StructRole::from_name("H1"), StructRole::H1); + assert_eq!(StructRole::from_name("P"), StructRole::P); + assert_eq!(StructRole::from_name("Table"), StructRole::Table); + assert_eq!(StructRole::from_name("TD"), StructRole::TD); + assert_eq!( + StructRole::from_name("CustomTag"), + StructRole::Other("CustomTag".to_string()) + ); + } + + #[test] + fn test_struct_role_with_role_map() { + let mut role_map = HashMap::new(); + role_map.insert("Heading1".to_string(), "H1".to_string()); + role_map.insert("Body".to_string(), "P".to_string()); + // Chain: MyTag → Heading1 → H1 + role_map.insert("MyTag".to_string(), "Heading1".to_string()); + + assert_eq!( + StructRole::from_name_with_role_map("Heading1", &role_map), + StructRole::H1 + ); + assert_eq!( + StructRole::from_name_with_role_map("Body", &role_map), + StructRole::P + ); + assert_eq!( + StructRole::from_name_with_role_map("MyTag", &role_map), + StructRole::H1 + ); + // Standard names bypass the map + assert_eq!( + StructRole::from_name_with_role_map("H2", &role_map), + StructRole::H2 + ); + } + + #[test] + fn test_struct_role_role_map_cycle() { + // A→B→A cycle should not infinite-loop + let mut role_map = HashMap::new(); + role_map.insert("A".to_string(), "B".to_string()); + role_map.insert("B".to_string(), "A".to_string()); + + let role = StructRole::from_name_with_role_map("A", &role_map); + // Should terminate (as Other) rather than loop forever + assert!(matches!(role, StructRole::Other(_))); + } + + #[test] + fn test_flat_struct_element() { + let tree = StructTree { + children: vec![StructElement { + role: StructRole::Document, + alt_text: None, + actual_text: None, + lang: None, + content_refs: Vec::new(), + children: vec![ + StructElement { + role: StructRole::H1, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![MarkedContentRef { + mcid: 0, + page_id: Some((1, 0)), + }], + children: Vec::new(), + }, + StructElement { + role: StructRole::P, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![MarkedContentRef { + mcid: 1, + page_id: Some((1, 0)), + }], + children: Vec::new(), + }, + ], + }], + }; + + let flat = tree.flatten(); + assert_eq!(flat.len(), 3); + assert_eq!(flat[0].role, StructRole::Document); + assert_eq!(flat[0].depth, 0); + assert_eq!(flat[1].role, StructRole::H1); + assert_eq!(flat[1].depth, 1); + assert_eq!(flat[2].role, StructRole::P); + assert_eq!(flat[2].depth, 1); + } + + #[test] + fn test_mcid_count() { + let tree = StructTree { + children: vec![StructElement { + role: StructRole::Document, + alt_text: None, + actual_text: None, + lang: None, + content_refs: Vec::new(), + children: vec![ + StructElement { + role: StructRole::H1, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![ + MarkedContentRef { + mcid: 0, + page_id: Some((1, 0)), + }, + MarkedContentRef { + mcid: 1, + page_id: Some((1, 0)), + }, + ], + children: Vec::new(), + }, + StructElement { + role: StructRole::P, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![MarkedContentRef { + mcid: 2, + page_id: Some((1, 0)), + }], + children: Vec::new(), + }, + ], + }], + }; + + assert_eq!(tree.mcid_count(), 3); + } + + #[test] + fn test_mcid_to_roles() { + use std::collections::BTreeMap; + + let page_id: ObjectId = (5, 0); + let mut page_ids = BTreeMap::new(); + page_ids.insert(1u32, page_id); + + let tree = StructTree { + children: vec![StructElement { + role: StructRole::Document, + alt_text: None, + actual_text: None, + lang: None, + content_refs: Vec::new(), + children: vec![ + StructElement { + role: StructRole::H1, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![MarkedContentRef { + mcid: 0, + page_id: Some(page_id), + }], + children: Vec::new(), + }, + StructElement { + role: StructRole::P, + alt_text: None, + actual_text: None, + lang: None, + content_refs: vec![MarkedContentRef { + mcid: 1, + page_id: Some(page_id), + }], + children: Vec::new(), + }, + ], + }], + }; + + let roles = tree.mcid_to_roles(&page_ids); + let page1 = roles.get(&1).unwrap(); + assert_eq!(page1.get(&0), Some(&StructRole::H1)); + assert_eq!(page1.get(&1), Some(&StructRole::P)); + } + + #[test] + fn test_fix_bare_struct_names() { + // Verify the byte-level pre-processor fixes bare names. + // All inputs include /StructTreeRoot to pass the early-return guard. + let input = b"/StructTreeRoot /S Code\n/Type /StructElem"; + let fixed = fix_bare_struct_names(input); + assert!( + fixed.windows(b"/S /Code".len()).any(|w| w == b"/S /Code"), + "Should fix bare Code: {:?}", + String::from_utf8_lossy(&fixed) + ); + + // Already correct — should return borrowed + let input = b"/StructTreeRoot /S /Code\n/Type /StructElem"; + let fixed = fix_bare_struct_names(input); + assert!(matches!(fixed, std::borrow::Cow::Borrowed(_))); + + // Multiple bare names + let input = b"/StructTreeRoot /S H1\n/foo\n/S P\n/bar"; + let fixed = fix_bare_struct_names(input); + let s = String::from_utf8_lossy(&fixed); + assert!(s.contains("/S /H1"), "Should fix H1: {s}"); + assert!(s.contains("/S /P"), "Should fix P: {s}"); + + // Unknown name should not be touched + let input = b"/StructTreeRoot /S FooBar\n"; + let fixed = fix_bare_struct_names(input); + let s = String::from_utf8_lossy(&fixed); + assert!(s.contains("/S FooBar"), "Should not fix unknown: {s}"); + + // No StructTreeRoot — skip entirely + let input = b"/S Code\nno struct tree"; + let fixed = fix_bare_struct_names(input); + assert!(matches!(fixed, std::borrow::Cow::Borrowed(_))); + } + + #[test] + fn test_bare_name_struct_types() { + // Some PDF generators (e.g. fpdf2) write /S Code instead of /S /Code. + // lopdf silently drops objects with invalid tokens. Our pre-processor + // fixes these before loading. + let raw = std::fs::read("tests/fixtures/bare_name_struct.pdf").unwrap(); + let fixed = fix_bare_struct_names(&raw); + let doc = Document::load_mem(fixed.as_ref()).unwrap(); + + let tree = StructTree::from_doc(&doc); + assert!(tree.is_some(), "Should parse bare-name struct tree"); + let tree = tree.unwrap(); + + let flat = tree.flatten(); + let roles: Vec<&StructRole> = flat.iter().map(|e| &e.role).collect(); + + assert!( + roles.iter().any(|r| matches!(r, StructRole::H1)), + "Should find H1 from bare name: {:?}", + roles + ); + assert!( + roles.iter().any(|r| matches!(r, StructRole::Code)), + "Should find Code from bare name: {:?}", + roles + ); + } + + #[test] + fn test_parse_real_tagged_pdf() { + let doc = Document::load("tests/fixtures/2013-app2.pdf").unwrap(); + let tree = StructTree::from_doc(&doc); + assert!(tree.is_some(), "2013-app2.pdf should have a structure tree"); + let tree = tree.unwrap(); + + // Should have a non-trivial structure + assert!(!tree.children.is_empty()); + assert!( + tree.mcid_count() > 0, + "Should have marked content references" + ); + + // Flatten and verify we get heading/paragraph/table elements + let flat = tree.flatten(); + let roles: Vec<&StructRole> = flat.iter().map(|e| &e.role).collect(); + assert!( + roles.iter().any(|r| matches!(r, StructRole::P)), + "Should contain paragraph elements" + ); + + // Verify mcid_to_roles produces a populated map + let page_ids = doc.get_pages(); + let role_map = tree.mcid_to_roles(&page_ids); + assert!(!role_map.is_empty(), "Should have MCID→role mappings"); + } +} diff --git a/src/tables/detect_heuristic.rs b/src/tables/detect_heuristic.rs index 873caa0..4576d9a 100644 --- a/src/tables/detect_heuristic.rs +++ b/src/tables/detect_heuristic.rs @@ -105,6 +105,7 @@ pub(crate) fn merge_adjacent_items(items: &[TextItem]) -> (Vec, Vec Option> is_bold: item.is_bold, is_italic: item.is_italic, item_type: item.item_type.clone(), + mcid: item.mcid, }); } Some(sub_items) diff --git a/src/tables/grid.rs b/src/tables/grid.rs index ceb9e69..e3361ef 100644 --- a/src/tables/grid.rs +++ b/src/tables/grid.rs @@ -381,6 +381,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, } } @@ -727,6 +728,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, page: 1, }, )); @@ -762,6 +764,7 @@ mod tests { is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, page: 1, }, )); diff --git a/src/tables/mod.rs b/src/tables/mod.rs index 043dacb..1e9e0b8 100644 --- a/src/tables/mod.rs +++ b/src/tables/mod.rs @@ -231,6 +231,7 @@ fn split_merged_numbers(item: &TextItem, col_boundaries: &[f32]) -> Vec Vec, } /// A line of text (grouped text items) diff --git a/tests/fixtures/bare_name_struct.pdf b/tests/fixtures/bare_name_struct.pdf new file mode 100644 index 0000000..bb7af19 Binary files /dev/null and b/tests/fixtures/bare_name_struct.pdf differ diff --git a/tests/fixtures/firecrawl_docs_tagged.pdf b/tests/fixtures/firecrawl_docs_tagged.pdf new file mode 100644 index 0000000..e4b3a78 Binary files /dev/null and b/tests/fixtures/firecrawl_docs_tagged.pdf differ diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index 1d2cd85..bd599c9 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -23,6 +23,7 @@ fn make_text_item(text: &str, x: f32, y: f32, font_size: f32, page: u32) -> Text is_bold: false, is_italic: false, item_type: ItemType::Text, + mcid: None, } } @@ -47,6 +48,7 @@ fn make_text_item_with_font( is_bold: is_bold_font(font), is_italic: is_italic_font(font), item_type: ItemType::Text, + mcid: None, } } @@ -999,3 +1001,39 @@ startxref ); } } + +#[test] +fn test_firecrawl_tagged_pdf_struct_tree() { + use lopdf::Document; + use pdf_inspector::structure_tree::{StructRole, StructTree}; + + let doc = Document::load("tests/fixtures/firecrawl_docs_tagged.pdf").unwrap(); + let tree = StructTree::from_doc(&doc).expect("Should have a structure tree"); + + // Verify structure tree contains expected roles + let page_ids = doc.get_pages(); + let roles = tree.mcid_to_roles(&page_ids); + assert!(!roles.is_empty(), "Should have MCID roles across pages"); + + let flat = tree.flatten(); + let has_code = flat.iter().any(|e| matches!(e.role, StructRole::Code)); + let has_h1 = flat.iter().any(|e| matches!(e.role, StructRole::H1)); + let has_li = flat.iter().any(|e| matches!(e.role, StructRole::LI)); + let has_caption = flat.iter().any(|e| matches!(e.role, StructRole::Caption)); + assert!(has_code, "Should have Code elements"); + assert!(has_h1, "Should have H1 elements"); + assert!(has_li, "Should have LI elements"); + assert!(has_caption, "Should have Caption elements"); + + // Full conversion: code fences should be generated from Code struct elements + let buf = std::fs::read("tests/fixtures/firecrawl_docs_tagged.pdf").unwrap(); + let result = pdf_inspector::process_pdf_mem(&buf).unwrap(); + let md = result.markdown.unwrap(); + let fence_count = md.matches("```").count(); + assert!( + fence_count > 0, + "Should produce code fences from tagged Code elements" + ); + // Fences come in open/close pairs + assert_eq!(fence_count % 2, 0, "Code fences should be balanced"); +}