oxidize_pdf/memory/

stream_processor.rs

//! Stream processing for memory-efficient PDF operations
//!
//! Processes PDF content incrementally without loading entire documents
//! into memory, ideal for large files or memory-constrained environments.

use crate::error::{PdfError, Result};
use crate::parser::content::{ContentOperation, ContentParser};
use crate::parser::PdfObject;
use std::io::{BufRead, BufReader, Read, Seek, Write};

/// Options for streaming operations
#[derive(Debug, Clone)]
pub struct StreamingOptions {
    /// Buffer size for reading
    pub buffer_size: usize,
    /// Maximum content stream size to process at once
    pub max_stream_size: usize,
    /// Whether to skip processing images
    pub skip_images: bool,
    /// Whether to skip processing fonts
    pub skip_fonts: bool,
}

impl Default for StreamingOptions {
    fn default() -> Self {
        Self {
            buffer_size: 64 * 1024,            // 64KB
            max_stream_size: 10 * 1024 * 1024, // 10MB
            skip_images: false,
            skip_fonts: false,
        }
    }
}

/// Stream processor for incremental PDF processing
pub struct StreamProcessor<R: Read + Seek> {
    reader: BufReader<R>,
    #[allow(dead_code)]
    options: StreamingOptions,
}

impl<R: Read + Seek> StreamProcessor<R> {
    /// Create a new stream processor
    pub fn new(reader: R, options: StreamingOptions) -> Self {
        let buf_reader = BufReader::with_capacity(options.buffer_size, reader);
        Self {
            reader: buf_reader,
            options,
        }
    }

    /// Process a PDF incrementally with a callback
    pub fn process_with<F>(&mut self, mut callback: F) -> Result<()>
    where
        F: FnMut(ProcessingEvent) -> Result<ProcessingAction>,
    {
        // Start processing
        callback(ProcessingEvent::Start)?;

        // Process header
        self.process_header(&mut callback)?;

        // Process objects incrementally
        self.process_objects(&mut callback)?;

        // End processing
        callback(ProcessingEvent::End)?;

        Ok(())
    }

    /// Process pages incrementally
    pub fn process_pages<F>(&mut self, mut page_callback: F) -> Result<()>
    where
        F: FnMut(u32, PageData) -> Result<ProcessingAction>,
    {
        let mut page_index = 0;

        self.process_with(|event| match event {
            ProcessingEvent::Page(data) => {
                let action = page_callback(page_index, data)?;
                page_index += 1;
                Ok(action)
            }
            _ => Ok(ProcessingAction::Continue),
        })
    }

    /// Extract text incrementally
    pub fn extract_text_streaming<W: Write>(&mut self, output: &mut W) -> Result<()> {
        self.process_pages(|_index, page_data| {
            if let Some(text) = page_data.text {
                output.write_all(text.as_bytes())?;
                output.write_all(b"\n")?;
            }
            Ok(ProcessingAction::Continue)
        })
    }

    fn process_header<F>(&mut self, callback: &mut F) -> Result<()>
    where
        F: FnMut(ProcessingEvent) -> Result<ProcessingAction>,
    {
        let mut header = String::new();
        self.reader.read_line(&mut header)?;

        if !header.starts_with("%PDF-") {
            return Err(PdfError::InvalidHeader);
        }

        let version = header.trim_start_matches("%PDF-").trim();
        callback(ProcessingEvent::Header {
            version: version.to_string(),
        })?;

        Ok(())
    }

    fn process_objects<F>(&mut self, callback: &mut F) -> Result<()>
    where
        F: FnMut(ProcessingEvent) -> Result<ProcessingAction>,
    {
        // In a real implementation, this would parse objects incrementally
        // For now, we'll simulate streaming behavior

        // Process some mock pages
        for i in 0..3 {
            let page_data = PageData {
                number: i,
                width: 595.0,
                height: 842.0,
                text: Some(format!("Page {} content", i + 1)),
                operations: vec![],
            };

            match callback(ProcessingEvent::Page(page_data))? {
                ProcessingAction::Continue => {}
                ProcessingAction::Skip => continue,
                ProcessingAction::Stop => break,
            }
        }

        Ok(())
    }
}

/// Events during stream processing
#[derive(Debug)]
pub enum ProcessingEvent {
    /// Processing started
    Start,
    /// PDF header found
    Header { version: String },
    /// Object encountered
    Object { id: (u32, u16), object: PdfObject },
    /// Page encountered
    Page(PageData),
    /// Resource encountered
    Resource {
        name: String,
        resource_type: ResourceType,
    },
    /// Processing ended
    End,
}

/// Page data during streaming
#[derive(Debug)]
pub struct PageData {
    /// Page number (0-indexed)
    pub number: u32,
    /// Page width in points
    pub width: f32,
    /// Page height in points
    pub height: f32,
    /// Extracted text (if any)
    pub text: Option<String>,
    /// Content operations (if requested)
    pub operations: Vec<ContentOperation>,
}

/// Resource types
#[derive(Debug, Clone)]
pub enum ResourceType {
    Font,
    Image,
    ColorSpace,
    Pattern,
    XObject,
}

/// Action to take after processing an event
#[derive(Debug, PartialEq)]
pub enum ProcessingAction {
    /// Continue processing
    Continue,
    /// Skip this item
    Skip,
    /// Stop processing
    Stop,
}

/// Stream-based content processor for individual content streams
pub struct ContentStreamProcessor {
    buffer: Vec<u8>,
    options: StreamingOptions,
}

impl ContentStreamProcessor {
    /// Create a new content stream processor
    pub fn new(options: StreamingOptions) -> Self {
        Self {
            buffer: Vec::with_capacity(options.buffer_size),
            options,
        }
    }

    /// Process a content stream incrementally
    pub fn process_stream<R: Read, F>(&mut self, mut reader: R, mut callback: F) -> Result<()>
    where
        F: FnMut(&ContentOperation) -> Result<ProcessingAction>,
    {
        self.buffer.clear();
        reader.read_to_end(&mut self.buffer)?;

        if self.buffer.len() > self.options.max_stream_size {
            return Err(PdfError::ContentStreamTooLarge(self.buffer.len()));
        }

        let operations =
            ContentParser::parse(&self.buffer).map_err(|e| PdfError::ParseError(e.to_string()))?;

        for op in operations {
            match callback(&op)? {
                ProcessingAction::Continue => {}
                ProcessingAction::Skip => continue,
                ProcessingAction::Stop => break,
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn test_streaming_options_default() {
        let options = StreamingOptions::default();
        assert_eq!(options.buffer_size, 64 * 1024);
        assert_eq!(options.max_stream_size, 10 * 1024 * 1024);
        assert!(!options.skip_images);
        assert!(!options.skip_fonts);
    }

    #[test]
    fn test_stream_processor_creation() {
        let data = b"%PDF-1.7\n";
        let cursor = Cursor::new(data);
        let options = StreamingOptions::default();
        let _processor = StreamProcessor::new(cursor, options);
    }

    #[test]
    fn test_processing_events() {
        let data = b"%PDF-1.7\n";
        let cursor = Cursor::new(data);
        let options = StreamingOptions::default();
        let mut processor = StreamProcessor::new(cursor, options);

        let mut events = Vec::new();

        processor
            .process_with(|event| {
                match &event {
                    ProcessingEvent::Start => events.push("start"),
                    ProcessingEvent::Header { version } => {
                        assert_eq!(version, "1.7");
                        events.push("header");
                    }
                    ProcessingEvent::Page(_) => events.push("page"),
                    ProcessingEvent::End => events.push("end"),
                    _ => {}
                }
                Ok(ProcessingAction::Continue)
            })
            .unwrap();

        assert!(events.contains(&"start"));
        assert!(events.contains(&"header"));
        assert!(events.contains(&"end"));
    }

    #[test]
    fn test_process_pages() {
        let data = b"%PDF-1.7\n";
        let cursor = Cursor::new(data);
        let options = StreamingOptions::default();
        let mut processor = StreamProcessor::new(cursor, options);

        let mut page_count = 0;

        processor
            .process_pages(|index, page| {
                assert_eq!(index, page_count);
                assert_eq!(page.width, 595.0);
                assert_eq!(page.height, 842.0);
                page_count += 1;
                Ok(ProcessingAction::Continue)
            })
            .unwrap();

        assert!(page_count > 0);
    }

    #[test]
    fn test_extract_text_streaming() {
        let data = b"%PDF-1.7\n";
        let cursor = Cursor::new(data);
        let options = StreamingOptions::default();
        let mut processor = StreamProcessor::new(cursor, options);

        let mut output = Vec::new();
        processor.extract_text_streaming(&mut output).unwrap();

        let text = String::from_utf8(output).unwrap();
        assert!(text.contains("Page"));
    }

    #[test]
    fn test_processing_action() {
        assert_eq!(ProcessingAction::Continue, ProcessingAction::Continue);
        assert_eq!(ProcessingAction::Skip, ProcessingAction::Skip);
        assert_eq!(ProcessingAction::Stop, ProcessingAction::Stop);
        assert_ne!(ProcessingAction::Continue, ProcessingAction::Stop);
    }

    #[test]
    fn test_content_stream_processor() {
        let options = StreamingOptions::default();
        let mut processor = ContentStreamProcessor::new(options);

        // Test with simple content
        let content = b"BT /F1 12 Tf 100 700 Td (Hello) Tj ET";
        let cursor = Cursor::new(content);

        let mut op_count = 0;
        processor
            .process_stream(cursor, |op| {
                op_count += 1;
                match op {
                    ContentOperation::BeginText => assert_eq!(op_count, 1),
                    ContentOperation::EndText => assert_eq!(op_count, 5),
                    _ => {}
                }
                Ok(ProcessingAction::Continue)
            })
            .unwrap();

        assert!(op_count > 0);
    }

    #[test]
    fn test_stop_processing() {
        let data = b"%PDF-1.7\n";
        let cursor = Cursor::new(data);
        let options = StreamingOptions::default();
        let mut processor = StreamProcessor::new(cursor, options);

        let mut page_count = 0;

        processor
            .process_pages(|_index, _page| {
                page_count += 1;
                if page_count >= 2 {
                    Ok(ProcessingAction::Stop)
                } else {
                    Ok(ProcessingAction::Continue)
                }
            })
            .unwrap();

        assert_eq!(page_count, 2);
    }
}