edgeparse_core/

lib.rs

//! EdgeParse Core Library
//!
//! High-performance PDF-to-structured-data extraction engine.
//! Implements a 20-stage processing pipeline for extracting text, tables,
//! images, and semantic structure from PDF documents.

#![warn(missing_docs)]

pub mod api;
pub mod models;
pub mod output;
pub mod pdf;
pub mod pipeline;
pub mod utils;

#[cfg(feature = "hybrid")]
pub mod hybrid;

pub mod tagged;

use crate::api::config::ProcessingConfig;
use crate::models::content::ContentElement;
use crate::models::document::PdfDocument;
use crate::pdf::chunk_parser::extract_page_chunks;
use crate::pdf::page_info;
use crate::pdf::raster_table_ocr::recover_raster_table_borders;
use crate::pipeline::orchestrator::{run_pipeline, PipelineState};
use crate::tagged::struct_tree::build_mcid_map;

/// Main entry point: convert a PDF file to structured data.
///
/// # Arguments
/// * `input_path` - Path to the input PDF file
/// * `config` - Processing configuration
///
/// # Returns
/// * `Result<PdfDocument>` - The extracted structured document
///
/// # Errors
/// Returns an error if the PDF cannot be loaded or processed.
pub fn convert(
    input_path: &std::path::Path,
    config: &ProcessingConfig,
) -> Result<PdfDocument, EdgePdfError> {
    let raw_doc = pdf::loader::load_pdf(input_path, config.password.as_deref())?;

    // Extract per-page geometry (MediaBox, CropBox, rotation) for use throughout the pipeline.
    let page_info_list = page_info::extract_page_info(&raw_doc.document);

    // Extract text chunks from each page
    let pages_map = raw_doc.document.get_pages();
    let mut page_contents = Vec::with_capacity(pages_map.len());

    for (&page_num, &page_id) in &pages_map {
        let page_chunks = extract_page_chunks(&raw_doc.document, page_num, page_id)?;
        let mut recovered_tables = Vec::new();
        if let Some(page_info) = page_info_list
            .iter()
            .find(|info| info.page_number == page_num)
        {
            recovered_tables = recover_raster_table_borders(
                input_path,
                &page_info.crop_box,
                page_num,
                &page_chunks.text_chunks,
                &page_chunks.image_chunks,
            );
        }
        let mut elements: Vec<ContentElement> = page_chunks
            .text_chunks
            .into_iter()
            .map(ContentElement::TextChunk)
            .collect();

        elements.extend(
            page_chunks
                .image_chunks
                .into_iter()
                .map(ContentElement::Image),
        );
        elements.extend(
            page_chunks
                .line_chunks
                .into_iter()
                .map(ContentElement::Line),
        );
        elements.extend(
            page_chunks
                .line_art_chunks
                .into_iter()
                .map(ContentElement::LineArt),
        );
        elements.extend(
            recovered_tables
                .into_iter()
                .map(ContentElement::TableBorder),
        );

        page_contents.push(elements);
    }

    // Run the processing pipeline
    let mcid_map = build_mcid_map(&raw_doc.document);
    let mut pipeline_state = PipelineState::with_mcid_map(page_contents, config.clone(), mcid_map)
        .with_page_info(page_info_list);
    run_pipeline(&mut pipeline_state)?;

    // Build the output document
    let file_name = input_path
        .file_name()
        .and_then(|n| n.to_str())
        .unwrap_or("unknown.pdf")
        .to_string();

    let mut doc = PdfDocument::new(file_name);
    doc.number_of_pages = pages_map.len() as u32;
    doc.author = raw_doc.metadata.author;
    doc.title = raw_doc.metadata.title;
    doc.creation_date = raw_doc.metadata.creation_date;
    doc.modification_date = raw_doc.metadata.modification_date;

    // Flatten pipeline output into document kids
    for page in pipeline_state.pages {
        doc.kids.extend(page);
    }

    Ok(doc)
}

/// Top-level error type for EdgeParse operations.
#[derive(Debug, thiserror::Error)]
pub enum EdgePdfError {
    /// PDF loading error
    #[error("PDF loading error: {0}")]
    LoadError(String),

    /// Pipeline processing error
    #[error("Pipeline error at stage {stage}: {message}")]
    PipelineError {
        /// Pipeline stage number (1-20)
        stage: u32,
        /// Error description
        message: String,
    },

    /// Output generation error
    #[error("Output error: {0}")]
    OutputError(String),

    /// I/O error
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),

    /// Configuration error
    #[error("Configuration error: {0}")]
    ConfigError(String),

    /// lopdf error
    #[error("PDF parse error: {0}")]
    LopdfError(String),
}

impl From<lopdf::Error> for EdgePdfError {
    fn from(e: lopdf::Error) -> Self {
        EdgePdfError::LopdfError(e.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use lopdf::{
        content::{Content, Operation},
        dictionary, Object, Stream,
    };
    use std::io::Write;

    /// Create a synthetic PDF file for integration testing.
    fn create_test_pdf_file(path: &std::path::Path) {
        let mut doc = lopdf::Document::with_version("1.5");
        let pages_id = doc.new_object_id();

        let font_id = doc.add_object(dictionary! {
            "Type" => "Font",
            "Subtype" => "Type1",
            "BaseFont" => "Helvetica",
        });

        let resources_id = doc.add_object(dictionary! {
            "Font" => dictionary! {
                "F1" => font_id,
            },
        });

        let content = Content {
            operations: vec![
                Operation::new("BT", vec![]),
                Operation::new("Tf", vec!["F1".into(), 12.into()]),
                Operation::new("Td", vec![72.into(), 700.into()]),
                Operation::new("Tj", vec![Object::string_literal("Hello EdgeParse!")]),
                Operation::new("Td", vec![0.into(), Object::Real(-20.0)]),
                Operation::new("Tj", vec![Object::string_literal("Second line of text.")]),
                Operation::new("ET", vec![]),
            ],
        };

        let encoded = content.encode().unwrap();
        let content_id = doc.add_object(Stream::new(dictionary! {}, encoded));

        let page_id = doc.add_object(dictionary! {
            "Type" => "Page",
            "Parent" => pages_id,
            "Contents" => content_id,
            "Resources" => resources_id,
            "MediaBox" => vec![0.into(), 0.into(), 595.into(), 842.into()],
        });

        let pages = dictionary! {
            "Type" => "Pages",
            "Kids" => vec![page_id.into()],
            "Count" => 1,
        };
        doc.objects.insert(pages_id, Object::Dictionary(pages));

        let catalog_id = doc.add_object(dictionary! {
            "Type" => "Catalog",
            "Pages" => pages_id,
        });
        doc.trailer.set("Root", catalog_id);

        let mut file = std::fs::File::create(path).unwrap();
        doc.save_to(&mut file).unwrap();
        file.flush().unwrap();
    }

    #[test]
    fn test_convert_end_to_end() {
        let dir = std::env::temp_dir().join("edgeparse_test");
        std::fs::create_dir_all(&dir).unwrap();
        let pdf_path = dir.join("test_convert.pdf");

        create_test_pdf_file(&pdf_path);

        let config = ProcessingConfig::default();
        let result = convert(&pdf_path, &config);
        assert!(result.is_ok(), "convert() failed: {:?}", result.err());

        let doc = result.unwrap();
        assert_eq!(doc.number_of_pages, 1);
        assert!(
            !doc.kids.is_empty(),
            "Expected content elements in document"
        );

        // Check that we extracted content (may be TextChunks, TextLines, or TextBlocks after pipeline)
        let mut all_text = String::new();
        for element in &doc.kids {
            match element {
                models::content::ContentElement::TextChunk(tc) => {
                    all_text.push_str(&tc.value);
                    all_text.push(' ');
                }
                models::content::ContentElement::TextLine(tl) => {
                    all_text.push_str(&tl.value());
                    all_text.push(' ');
                }
                models::content::ContentElement::TextBlock(tb) => {
                    all_text.push_str(&tb.value());
                    all_text.push(' ');
                }
                models::content::ContentElement::Paragraph(p) => {
                    all_text.push_str(&p.base.value());
                    all_text.push(' ');
                }
                models::content::ContentElement::Heading(h) => {
                    all_text.push_str(&h.base.base.value());
                    all_text.push(' ');
                }
                _ => {}
            }
        }

        assert!(
            all_text.contains("Hello"),
            "Expected 'Hello' in extracted text, got: {}",
            all_text
        );
        assert!(
            all_text.contains("Second"),
            "Expected 'Second' in extracted text, got: {}",
            all_text
        );

        // Cleanup
        let _ = std::fs::remove_file(&pdf_path);
    }
}