Skip to main content

oxidize_pdf/pipeline/
element.rs

1use std::cmp::Ordering;
2
3#[cfg(feature = "semantic")]
4use serde::{Deserialize, Serialize};
5
6/// A typed document element extracted from a PDF page.
7///
8/// Each variant carries its specific data plus shared [`ElementMetadata`]
9/// for page number, bounding box, confidence, and optional font info.
10#[derive(Debug, Clone)]
11#[cfg_attr(
12    feature = "semantic",
13    derive(Serialize, Deserialize),
14    serde(tag = "type", rename_all = "snake_case")
15)]
16pub enum Element {
17    Title(ElementData),
18    Paragraph(ElementData),
19    Table(TableElementData),
20    Header(ElementData),
21    Footer(ElementData),
22    ListItem(ElementData),
23    Image(ImageElementData),
24    CodeBlock(ElementData),
25    KeyValue(KeyValueElementData),
26}
27
28impl Element {
29    /// Returns the primary text content of this element.
30    pub fn text(&self) -> &str {
31        match self {
32            Self::Title(d)
33            | Self::Paragraph(d)
34            | Self::Header(d)
35            | Self::Footer(d)
36            | Self::ListItem(d)
37            | Self::CodeBlock(d) => &d.text,
38            Self::Table(t) => {
39                // Tables don't have a single text — return empty.
40                // Use row_count()/cell() for structured access.
41                let _ = t;
42                ""
43            }
44            Self::Image(img) => img.alt_text.as_deref().unwrap_or(""),
45            Self::KeyValue(kv) => &kv.value,
46        }
47    }
48
49    /// Returns a human-readable text representation of this element.
50    ///
51    /// Unlike [`text()`](Self::text) which returns raw content (empty for tables,
52    /// value-only for KV pairs), this method produces a complete display form:
53    /// - Tables: pipe-separated rows
54    /// - Key-Value: "key: value"
55    /// - All others: same as `text()`
56    pub fn display_text(&self) -> String {
57        match self {
58            Self::Table(t) => t
59                .rows
60                .iter()
61                .map(|row| row.join(" | "))
62                .collect::<Vec<_>>()
63                .join("\n"),
64            Self::Image(img) => img.alt_text.clone().unwrap_or_default(),
65            Self::KeyValue(kv) => format!("{}: {}", kv.key, kv.value),
66            _ => self.text().to_string(),
67        }
68    }
69
70    /// Returns the page number (0-indexed) where this element appears.
71    pub fn page(&self) -> u32 {
72        self.metadata().page
73    }
74
75    /// Returns the bounding box of this element on the page.
76    pub fn bbox(&self) -> &ElementBBox {
77        &self.metadata().bbox
78    }
79
80    /// Returns the full metadata for this element.
81    pub fn metadata(&self) -> &ElementMetadata {
82        match self {
83            Self::Title(d)
84            | Self::Paragraph(d)
85            | Self::Header(d)
86            | Self::Footer(d)
87            | Self::ListItem(d)
88            | Self::CodeBlock(d) => &d.metadata,
89            Self::Table(t) => &t.metadata,
90            Self::Image(img) => &img.metadata,
91            Self::KeyValue(kv) => &kv.metadata,
92        }
93    }
94
95    /// Returns a mutable reference to the metadata of this element.
96    pub fn metadata_mut(&mut self) -> &mut ElementMetadata {
97        match self {
98            Self::Title(d)
99            | Self::Paragraph(d)
100            | Self::Header(d)
101            | Self::Footer(d)
102            | Self::ListItem(d)
103            | Self::CodeBlock(d) => &mut d.metadata,
104            Self::Table(t) => &mut t.metadata,
105            Self::Image(img) => &mut img.metadata,
106            Self::KeyValue(kv) => &mut kv.metadata,
107        }
108    }
109
110    /// Returns the snake_case type name of this element variant.
111    ///
112    /// Useful for logging, serialization, and metadata tagging.
113    /// Returns one of: `"title"`, `"paragraph"`, `"table"`, `"header"`,
114    /// `"footer"`, `"list_item"`, `"image"`, `"code_block"`, `"key_value"`.
115    pub fn type_name(&self) -> &'static str {
116        match self {
117            Self::Title(_) => "title",
118            Self::Paragraph(_) => "paragraph",
119            Self::Table(_) => "table",
120            Self::Header(_) => "header",
121            Self::Footer(_) => "footer",
122            Self::ListItem(_) => "list_item",
123            Self::Image(_) => "image",
124            Self::CodeBlock(_) => "code_block",
125            Self::KeyValue(_) => "key_value",
126        }
127    }
128
129    /// Set the parent heading for this element.
130    pub fn set_parent_heading(&mut self, heading: Option<String>) {
131        self.metadata_mut().parent_heading = heading;
132    }
133
134    /// Set the full heading breadcrumb for this element.
135    pub fn set_heading_path(&mut self, path: Vec<String>) {
136        self.metadata_mut().heading_path = path;
137    }
138
139    /// Returns the number of rows if this is a Table element.
140    pub fn row_count(&self) -> Option<usize> {
141        match self {
142            Self::Table(t) => Some(t.rows.len()),
143            _ => None,
144        }
145    }
146
147    /// Returns the number of columns if this is a Table element.
148    pub fn column_count(&self) -> Option<usize> {
149        match self {
150            Self::Table(t) => t.rows.first().map(|r| r.len()),
151            _ => None,
152        }
153    }
154
155    /// Returns the cell text at (row, col) if this is a Table element.
156    pub fn cell(&self, row: usize, col: usize) -> Option<&str> {
157        match self {
158            Self::Table(t) => t.rows.get(row).and_then(|r| r.get(col)).map(|s| s.as_str()),
159            _ => None,
160        }
161    }
162}
163
164/// Comparator for sorting elements in natural reading order:
165/// page ASC, then Y DESC (top-to-bottom in PDF coordinates), then X ASC (left-to-right).
166///
167/// Use with `elements.sort_by(element_reading_order)` instead of `elements.sort()`.
168pub fn element_reading_order(a: &Element, b: &Element) -> Ordering {
169    let page_cmp = a.page().cmp(&b.page());
170    if page_cmp != Ordering::Equal {
171        return page_cmp;
172    }
173    // Higher Y = higher on page in PDF coords → should come first
174    let y_cmp = b.bbox().y.total_cmp(&a.bbox().y);
175    if y_cmp != Ordering::Equal {
176        return y_cmp;
177    }
178    a.bbox().x.total_cmp(&b.bbox().x)
179}
180
181impl PartialEq for Element {
182    fn eq(&self, other: &Self) -> bool {
183        std::mem::discriminant(self) == std::mem::discriminant(other) && self.content_eq(other)
184    }
185}
186
187impl Eq for Element {}
188
189impl Element {
190    /// Content-based equality: compares text/data content, ignoring metadata (position, font, etc.).
191    fn content_eq(&self, other: &Self) -> bool {
192        match (self, other) {
193            (Self::Title(a), Self::Title(b))
194            | (Self::Paragraph(a), Self::Paragraph(b))
195            | (Self::Header(a), Self::Header(b))
196            | (Self::Footer(a), Self::Footer(b))
197            | (Self::ListItem(a), Self::ListItem(b))
198            | (Self::CodeBlock(a), Self::CodeBlock(b)) => a.text == b.text,
199            (Self::Table(a), Self::Table(b)) => a.rows == b.rows,
200            (Self::Image(a), Self::Image(b)) => a.alt_text == b.alt_text,
201            (Self::KeyValue(a), Self::KeyValue(b)) => a.key == b.key && a.value == b.value,
202            _ => false,
203        }
204    }
205}
206
207/// Shared data for text-based element variants.
208#[derive(Debug, Clone)]
209#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
210pub struct ElementData {
211    pub text: String,
212    pub metadata: ElementMetadata,
213}
214
215/// One cell of a table's rich structure. `row`/`col` are the cell's top-left
216/// position in the base grid; `row_span`/`col_span` are >= 1.
217#[derive(Debug, Clone)]
218#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
219pub struct RichCell {
220    pub row: usize,
221    pub col: usize,
222    pub row_span: usize,
223    pub col_span: usize,
224    pub text: String,
225    pub is_header: bool,
226}
227
228/// Rich table structure: merged cells and header rows. Present only when a hard
229/// signal (drawn grid / structure tags) revealed it; borderless tables leave it
230/// `None` and use the flat `rows` view only.
231#[derive(Debug, Clone)]
232#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
233pub struct TableStructure {
234    pub cells: Vec<RichCell>,
235    pub num_rows: usize,
236    pub num_cols: usize,
237    /// Number of leading rows that are headers (0 = none, 1 = single header row,
238    /// >1 = multi-level header expressed as header rows + merged cells).
239    pub header_rows: usize,
240}
241
242/// Data specific to table elements.
243#[derive(Debug, Clone)]
244#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
245#[non_exhaustive]
246pub struct TableElementData {
247    /// Row-major flat cell data. Each inner Vec is one row. When `structure` is
248    /// present this is its expanded (span-repeated) view; otherwise it is the
249    /// primary representation.
250    pub rows: Vec<Vec<String>>,
251    /// Rich structure (merged cells / header rows) when a hard signal revealed it.
252    pub structure: Option<TableStructure>,
253    pub metadata: ElementMetadata,
254}
255
256impl TableElementData {
257    /// Build a plain (non-rich) table from flat row-major cells. `structure` is
258    /// `None`; use `from_structure` when merged cells / header rows are known.
259    pub fn new(rows: Vec<Vec<String>>, metadata: ElementMetadata) -> Self {
260        Self {
261            rows,
262            structure: None,
263            metadata,
264        }
265    }
266
267    /// Build from rich structure, deriving the flat `rows` view by expanding each
268    /// spanning cell's text across every covered (row, col).
269    pub fn from_structure(structure: TableStructure, metadata: ElementMetadata) -> Self {
270        let mut rows = vec![vec![String::new(); structure.num_cols]; structure.num_rows];
271        for cell in &structure.cells {
272            for r in cell.row..(cell.row + cell.row_span).min(structure.num_rows) {
273                for c in cell.col..(cell.col + cell.col_span).min(structure.num_cols) {
274                    rows[r][c] = cell.text.clone();
275                }
276            }
277        }
278        Self {
279            rows,
280            structure: Some(structure),
281            metadata,
282        }
283    }
284}
285
286/// Data specific to image elements.
287#[derive(Debug, Clone)]
288#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
289pub struct ImageElementData {
290    pub alt_text: Option<String>,
291    pub metadata: ElementMetadata,
292}
293
294/// Data specific to key-value pair elements.
295#[derive(Debug, Clone)]
296#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
297pub struct KeyValueElementData {
298    pub key: String,
299    pub value: String,
300    pub metadata: ElementMetadata,
301}
302
303/// Metadata common to all element types.
304#[derive(Debug, Clone)]
305#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
306pub struct ElementMetadata {
307    /// Page number (0-indexed).
308    pub page: u32,
309    /// Bounding box on the page.
310    pub bbox: ElementBBox,
311    /// Classification confidence (0.0–1.0).
312    pub confidence: f64,
313    /// Font name if detected.
314    pub font_name: Option<String>,
315    /// Font size in points if detected.
316    pub font_size: Option<f64>,
317    /// Whether the text is bold.
318    pub is_bold: bool,
319    /// Whether the text is italic.
320    pub is_italic: bool,
321    /// The text of the nearest preceding Title element, if any.
322    pub parent_heading: Option<String>,
323    /// Full ancestor heading breadcrumb, root→leaf. Empty if outside any heading.
324    pub heading_path: Vec<String>,
325    /// Open class label assigned by a custom
326    /// [`ElementClassifier`](crate::pipeline::spi::ElementClassifier) before
327    /// chunking (e.g. `"clause"`, `"definition"`). `None` unless a classifier
328    /// set it. A chunking strategy may read it to make boundary decisions.
329    #[cfg(feature = "unstable-spi")]
330    #[cfg_attr(
331        feature = "semantic",
332        serde(default, skip_serializing_if = "Option::is_none")
333    )]
334    pub class_label: Option<String>,
335}
336
337impl Default for ElementMetadata {
338    fn default() -> Self {
339        Self {
340            page: 0,
341            bbox: ElementBBox::ZERO,
342            confidence: 1.0,
343            font_name: None,
344            font_size: None,
345            is_bold: false,
346            is_italic: false,
347            parent_heading: None,
348            heading_path: Vec::new(),
349            #[cfg(feature = "unstable-spi")]
350            class_label: None,
351        }
352    }
353}
354
355/// Axis-aligned bounding box for an element on a PDF page.
356#[derive(Debug, Clone, Copy, PartialEq)]
357#[cfg_attr(feature = "semantic", derive(Serialize, Deserialize))]
358pub struct ElementBBox {
359    /// Left edge X coordinate.
360    pub x: f64,
361    /// Bottom edge Y coordinate (PDF coordinate system).
362    pub y: f64,
363    /// Width of the bounding box.
364    pub width: f64,
365    /// Height of the bounding box.
366    pub height: f64,
367}
368
369impl ElementBBox {
370    /// A zero-sized bounding box at the origin.
371    pub const ZERO: Self = Self {
372        x: 0.0,
373        y: 0.0,
374        width: 0.0,
375        height: 0.0,
376    };
377
378    /// Creates a new bounding box.
379    pub fn new(x: f64, y: f64, width: f64, height: f64) -> Self {
380        Self {
381            x,
382            y,
383            width,
384            height,
385        }
386    }
387
388    /// Right edge X coordinate (x + width).
389    pub fn right(&self) -> f64 {
390        self.x + self.width
391    }
392
393    /// Top edge Y coordinate (y + height) in PDF coordinate system.
394    pub fn top(&self) -> f64 {
395        self.y + self.height
396    }
397}