rpdfium_core/

fx_stream.rs

//! PDF source data abstraction.
//!
//! The [`PdfSource`] trait abstracts over the source of PDF data, allowing
//! the same parsing code to work with heap-owned data (`Vec<u8>`, `Arc<[u8]>`)
//! and future backends like memory-mapped files.
//!
//! On WASM, PDF data is provided as `Vec<u8>` from a JavaScript `ArrayBuffer`.
//! On native platforms, `Arc<[u8]>` is the default. A future `MmapSource`
//! behind the `mmap` feature flag will enable memory-mapped file access.

use std::sync::Arc;

/// Abstraction over PDF source data.
///
/// Implementations must provide contiguous byte access via [`AsRef<[u8]>`].
/// All implementations must be `Send + Sync` to support concurrent object
/// resolution across threads.
pub trait PdfSource: AsRef<[u8]> + Send + Sync {
    /// Total size in bytes.
    fn len(&self) -> usize {
        self.as_ref().len()
    }

    /// Returns `true` if the source contains no bytes.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Heap-owned byte slice (safe, works everywhere including WASM).
impl PdfSource for Arc<[u8]> {}

/// Owned byte vector (safe, works everywhere including WASM).
impl PdfSource for Vec<u8> {}

/// Trait for writing a block of bytes.
///
/// Interface-only port of `IFX_WriteStream` (fx_stream.h).
pub trait WriteStream {
    /// Write `data` to the stream. Empty slice is a no-op; returns `true`.
    fn write_block(&mut self, data: &[u8]) -> bool;
}

/// Trait for a write stream that exposes its current byte offset.
///
/// Interface-only port of `IFX_ArchiveStream` (fx_stream.h).
pub trait ArchiveStream: WriteStream {
    fn current_offset(&self) -> i64;
}

/// Trait for a seekable writable stream.
///
/// Interface-only port of `IFX_SeekableWriteStream` (fx_stream.h).
pub trait SeekableWriteStream: WriteStream {
    /// Returns the total size of the stream in bytes.
    ///
    /// Corresponds to `IFX_StreamWithSize::GetSize()` in PDFium upstream.
    fn size(&self) -> i64;

    /// Upstream-aligned alias for [`size`](Self::size).
    ///
    /// Corresponds to `IFX_StreamWithSize::GetSize()` in PDFium upstream.
    #[inline]
    fn get_size(&self) -> i64 {
        self.size()
    }

    fn flush(&mut self) -> bool;
}

/// Trait combining seekable read and write access.
///
/// Interface-only port of `IFX_SeekableStream` (fx_stream.h).
pub trait SeekableStream: SeekableWriteStream {}

/// Concrete `WriteStream` implementation for `Vec<u8>`.
impl WriteStream for Vec<u8> {
    fn write_block(&mut self, data: &[u8]) -> bool {
        self.extend_from_slice(data);
        true
    }
}

/// Trait for writing PDF bytes with position tracking.
///
/// Port of `IFX_ArchiveStream` (fx_stream.h): write + offset tracking.
/// Uses `io::Result` instead of `bool` return (idiomatic Rust error handling).
pub trait PdfWrite {
    /// Write raw bytes to the output.
    fn write_bytes(&mut self, data: &[u8]) -> std::io::Result<()>;
    /// Current byte offset in the output.
    fn position(&self) -> u64;
}

/// `PdfWrite` implementation for `Vec<u8>`.
impl PdfWrite for Vec<u8> {
    fn write_bytes(&mut self, data: &[u8]) -> std::io::Result<()> {
        self.extend_from_slice(data);
        Ok(())
    }

    fn position(&self) -> u64 {
        self.len() as u64
    }
}

/// Position-tracking writer wrapper.
///
/// Wraps any [`std::io::Write`] implementor and tracks the byte offset written.
pub struct CountingWriter<W: std::io::Write> {
    inner: W,
    pos: u64,
}

impl<W: std::io::Write> CountingWriter<W> {
    /// Create a new counting writer starting at offset 0.
    pub fn new(inner: W) -> Self {
        Self { inner, pos: 0 }
    }

    /// Create a counting writer with an initial offset (e.g. for incremental appends).
    pub fn with_offset(inner: W, offset: u64) -> Self {
        Self { inner, pos: offset }
    }

    /// Consume the writer and return the inner writer.
    pub fn into_inner(self) -> W {
        self.inner
    }

    /// Flush the underlying writer.
    pub fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}

impl<W: std::io::Write> PdfWrite for CountingWriter<W> {
    fn write_bytes(&mut self, data: &[u8]) -> std::io::Result<()> {
        self.inner.write_all(data)?;
        self.pos += data.len() as u64;
        Ok(())
    }

    fn position(&self) -> u64 {
        self.pos
    }
}

/// In-memory stream supporting both read and write access.
///
/// Port of `CFX_MemoryStream` (cfx_memorystream.h/cpp).
/// Implements both [`PdfSource`] (read) and [`PdfWrite`] (write).
pub struct MemoryStream {
    data: Vec<u8>,
}

impl MemoryStream {
    /// Create an empty `MemoryStream`.
    pub fn new() -> Self {
        Self { data: Vec::new() }
    }

    /// Create a `MemoryStream` from existing bytes.
    pub fn from_bytes(data: Vec<u8>) -> Self {
        Self { data }
    }

    /// Consume the stream and return the underlying bytes.
    pub fn into_bytes(self) -> Vec<u8> {
        self.data
    }

    /// Return a slice of the stream's bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.data
    }
}

impl Default for MemoryStream {
    fn default() -> Self {
        Self::new()
    }
}

impl AsRef<[u8]> for MemoryStream {
    fn as_ref(&self) -> &[u8] {
        &self.data
    }
}

// Vec<u8> is Send + Sync, so MemoryStream { data: Vec<u8> } is too.
impl PdfSource for MemoryStream {}

impl PdfWrite for MemoryStream {
    fn write_bytes(&mut self, data: &[u8]) -> std::io::Result<()> {
        self.data.extend_from_slice(data);
        Ok(())
    }

    fn position(&self) -> u64 {
        self.data.len() as u64
    }
}

impl WriteStream for MemoryStream {
    fn write_block(&mut self, data: &[u8]) -> bool {
        self.data.extend_from_slice(data);
        true
    }
}

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

    #[test]
    fn test_arc_source_len() {
        let data: Arc<[u8]> = Arc::from(b"hello" as &[u8]);
        assert_eq!(PdfSource::len(&data), 5);
        assert!(!PdfSource::is_empty(&data));
    }

    #[test]
    fn test_arc_source_empty() {
        let data: Arc<[u8]> = Arc::from(b"" as &[u8]);
        assert_eq!(PdfSource::len(&data), 0);
        assert!(PdfSource::is_empty(&data));
    }

    #[test]
    fn test_vec_source_len() {
        let data: Vec<u8> = b"hello world".to_vec();
        assert_eq!(PdfSource::len(&data), 11);
        assert!(!PdfSource::is_empty(&data));
    }

    #[test]
    fn test_vec_source_empty() {
        let data: Vec<u8> = Vec::new();
        assert_eq!(PdfSource::len(&data), 0);
        assert!(PdfSource::is_empty(&data));
    }

    #[test]
    fn test_arc_source_as_ref() {
        let data: Arc<[u8]> = Arc::from(b"PDF" as &[u8]);
        assert_eq!(data.as_ref(), b"PDF");
    }

    #[test]
    fn test_vec_source_as_ref() {
        let data: Vec<u8> = b"PDF".to_vec();
        assert_eq!(<Vec<u8> as AsRef<[u8]>>::as_ref(&data), b"PDF");
    }

    #[test]
    fn test_pdf_source_is_send_sync() {
        fn assert_send_sync<T: PdfSource>() {}
        assert_send_sync::<Arc<[u8]>>();
        assert_send_sync::<Vec<u8>>();
    }

    #[test]
    fn test_pdf_source_object_safe() {
        // Verify that PdfSource can be used as a trait object
        fn use_source(source: &dyn PdfSource) -> usize {
            source.len()
        }
        let data: Vec<u8> = b"test".to_vec();
        assert_eq!(use_source(&data), 4);
    }

    #[test]
    fn test_vec_write_stream_write_block() {
        let mut buf: Vec<u8> = Vec::new();
        assert!(buf.write_block(b"hello"));
        assert!(buf.write_block(b" world"));
        assert_eq!(&buf, b"hello world");
    }

    #[test]
    fn test_vec_write_stream_empty_slice() {
        let mut buf: Vec<u8> = Vec::new();
        assert!(buf.write_block(b""));
        assert!(buf.is_empty());
    }

    #[test]
    fn test_vec_pdf_write_tracks_position() {
        let mut buf: Vec<u8> = Vec::new();
        assert_eq!(buf.position(), 0);
        buf.write_bytes(b"abc").unwrap();
        assert_eq!(buf.position(), 3);
        buf.write_bytes(b"de").unwrap();
        assert_eq!(buf.position(), 5);
    }

    #[test]
    fn test_counting_writer_tracks_position() {
        let mut writer = CountingWriter::new(Vec::new());
        assert_eq!(writer.position(), 0);
        writer.write_bytes(b"hello").unwrap();
        assert_eq!(writer.position(), 5);
        writer.write_bytes(b" world").unwrap();
        assert_eq!(writer.position(), 11);
    }

    #[test]
    fn test_counting_writer_with_offset() {
        let mut writer = CountingWriter::with_offset(Vec::new(), 100);
        assert_eq!(writer.position(), 100);
        writer.write_bytes(b"data").unwrap();
        assert_eq!(writer.position(), 104);
    }

    #[test]
    fn test_counting_writer_into_inner() {
        let mut writer = CountingWriter::new(Vec::new());
        writer.write_bytes(b"test").unwrap();
        let inner = writer.into_inner();
        assert_eq!(&inner, b"test");
    }

    #[test]
    fn test_memory_stream_default_is_empty() {
        let ms = MemoryStream::default();
        assert!(ms.as_bytes().is_empty());
        assert_eq!(ms.position(), 0);
    }

    #[test]
    fn test_memory_stream_new_is_empty() {
        let ms = MemoryStream::new();
        assert!(ms.as_bytes().is_empty());
    }

    #[test]
    fn test_memory_stream_from_bytes() {
        let ms = MemoryStream::from_bytes(b"hello".to_vec());
        assert_eq!(ms.as_bytes(), b"hello");
        assert_eq!(ms.position(), 5);
    }

    #[test]
    fn test_memory_stream_write_and_read_back() {
        let mut ms = MemoryStream::new();
        ms.write_bytes(b"PDF").unwrap();
        ms.write_bytes(b"-1.7").unwrap();
        assert_eq!(ms.as_bytes(), b"PDF-1.7");
        assert_eq!(ms.position(), 7);
    }

    #[test]
    fn test_memory_stream_into_bytes() {
        let mut ms = MemoryStream::new();
        ms.write_bytes(b"data").unwrap();
        let bytes = ms.into_bytes();
        assert_eq!(&bytes, b"data");
    }

    #[test]
    fn test_memory_stream_write_block() {
        let mut ms = MemoryStream::new();
        assert!(ms.write_block(b"block"));
        assert_eq!(ms.as_bytes(), b"block");
    }

    #[test]
    fn test_memory_stream_as_pdf_source() {
        let ms = MemoryStream::from_bytes(b"hello".to_vec());
        assert_eq!(PdfSource::len(&ms), 5);
        assert!(!PdfSource::is_empty(&ms));
    }

    #[test]
    fn test_memory_stream_is_send_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<MemoryStream>();
    }
}