Skip to main content

docling_core/
tree.rs

1//! docling's **item tree**, for a backend that knows the exact shape upstream
2//! gives a document and wants the JSON export to reproduce it.
3//!
4//! [`DoclingDocument::nodes`](crate::DoclingDocument::nodes) is a flat,
5//! reading-order stream tuned for Markdown / DocLang / LaTeX; the JSON export
6//! rebuilds docling's parent/child structure from it with generic rules (runs
7//! of list items become list groups, a heading is a flat sibling of the text
8//! that follows it). Upstream's backends do not all agree on that structure:
9//! the HTML backend nests everything after a heading *under* the heading,
10//! splits a paragraph of mixed formatting into an `inline` group of one text
11//! item per formatting run, parents a rich table cell's content to a group
12//! under the table, keeps site chrome on the `furniture` layer… and numbers
13//! every item in the order it *creates* them. A backend that ports those
14//! rules call-for-call (HTML's `html_tree.rs`, DOCX's `docx_tree.rs`) records
15//! the result here — an arena of items in
16//! creation order, each with its parent and children — and the JSON export
17//! ([`DoclingDocument::export_to_json`](crate::DoclingDocument::export_to_json))
18//! serializes this tree instead of deriving one from the nodes. Every other
19//! serializer keeps reading the flat nodes, so their output is unaffected.
20
21use crate::{ContentLayer, FieldItem, PictureImage, Script, Table};
22
23/// docling-core's `Formatting`: the inline styles an item carries in JSON.
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
25pub struct Formatting {
26    pub bold: bool,
27    pub italic: bool,
28    pub underline: bool,
29    pub strikethrough: bool,
30    pub script: Script,
31}
32
33/// A `list_item`'s docling fields.
34#[derive(Debug, Clone, PartialEq, Eq, Default)]
35pub struct ListMeta {
36    pub enumerated: bool,
37    /// docling's `marker` — the HTML backend writes `""` unless an ordered
38    /// list carries an explicit `start`, then `"{n}."`.
39    pub marker: String,
40}
41
42/// What an item in the tree is. Mirrors the docling-core item classes the
43/// JSON `texts` / `groups` / `tables` / `pictures` / `field_regions` buckets
44/// hold.
45#[derive(Debug, Clone, PartialEq)]
46pub enum TreeKind {
47    /// A `TextItem` / `TitleItem` / `SectionHeaderItem` / `ListItem`, told
48    /// apart by `label` (`text`, `title`, `section_header`, `list_item`,
49    /// `caption`, `checkbox_selected`, `checkbox_unselected`, …).
50    Text {
51        label: String,
52        text: String,
53        /// docling's `orig` when it differs from `text` (the heading text
54        /// before unicode cleanup, say); `None` = same as `text`.
55        orig: Option<String>,
56        formatting: Option<Formatting>,
57        hyperlink: Option<String>,
58        /// `section_header` only: docling's heading level.
59        level: Option<u8>,
60        /// `list_item` only.
61        list: Option<ListMeta>,
62    },
63    /// A `CodeItem`.
64    Code {
65        text: String,
66        orig: Option<String>,
67        /// The language hint (a highlighter class token such as `python`),
68        /// mapped onto docling's `CodeLanguageLabel` at export; `None` →
69        /// `unknown`.
70        language: Option<String>,
71        formatting: Option<Formatting>,
72        hyperlink: Option<String>,
73    },
74    /// A `GroupItem`: `label` is docling's `GroupLabel` value (`inline`,
75    /// `list`, `section`, `unspecified`, …), `name` its name (`group`, `list`,
76    /// `ordered list`, `header-2`, `rich_cell_group_1_0_3`, …).
77    Group { label: String, name: String },
78    /// A `TableItem`. `rich_cells` marks the cells docling serialized as a
79    /// `RichTableCell`: `(row, col)` grid anchor → the group item (a child of
80    /// the table) that holds the cell's content. `captions` are caption text
81    /// items in the tree.
82    Table {
83        table: Table,
84        rich_cells: Vec<(usize, usize, usize)>,
85        captions: Vec<usize>,
86    },
87    /// A `PictureItem`, its caption text items and optional payload.
88    /// `classification` is a `PictureClassificationLabel` value written as
89    /// the picture's `meta.classification` (an HTML `<stamp>` / `<signature>`).
90    Picture {
91        captions: Vec<usize>,
92        image: Option<PictureImage>,
93        classification: Option<String>,
94        /// The prediction's `confidence`, when the backend writes one: the
95        /// DocLang deserializer stamps `1.0`; the office and HTML backends
96        /// leave it out (`None`).
97        confidence: Option<f64>,
98        /// A native chart's data grid (docling's `meta.tabular_chart.chart_data`,
99        /// the series reconstructed as a `TableData`), for a DOCX chart drawing.
100        chart: Option<Table>,
101        /// The `ImageRef.dpi` docling writes for `image` when the backend
102        /// read one from the file (python-pptx's `Image.dpi`: PIL's `dpi`
103        /// info, rounded, 72 when absent or out of 1–2048); `None` → 72,
104        /// which is what upstream's other office backends pass.
105        dpi: Option<u32>,
106    },
107    /// A form key-value region (`field_regions` / `field_items`).
108    FieldRegion { items: Vec<FieldItem> },
109    /// A `KeyValueItem` (`key_value_items`): docling's `GraphData` of key and
110    /// value cells and their links, written verbatim.
111    KeyValueGraph {
112        cells: Vec<crate::GraphCell>,
113        links: Vec<crate::GraphLink>,
114    },
115}
116
117/// docling's `ProvenanceItem` for a tree item, written verbatim: the
118/// backend's own geometry in the page's units — a PPTX shape's EMU box,
119/// whose `pages` entry is the slide size in EMU — rather than the 0–511
120/// DocLang grid the flat [`Node::Located`](crate::Node::Located) carries
121/// (which cannot round-trip those integers).
122#[derive(Debug, Clone, PartialEq)]
123pub struct TreeProv {
124    /// 1-based page (slide) number.
125    pub page_no: usize,
126    /// `[l, t, r, b]`, exactly as docling computed them.
127    pub bbox: [f64; 4],
128    /// docling's `coord_origin` tag. The office backends tag their
129    /// top-left-based boxes `TOPLEFT` (the PPTX backend since docling#4294 —
130    /// it used to tag them `BOTTOMLEFT`, which read the tuple as `(l, b, r,
131    /// t)` and swapped the vertical edges); a backend that really works in a
132    /// bottom-left space sets this.
133    pub bottom_left: bool,
134    /// `[0, len(text)]` in characters for a text item, `[0, 0]` for a table
135    /// or picture.
136    pub charspan: [usize; 2],
137}
138
139/// docling's `TrackSource` — where in a time-based track (a WebVTT cue) a
140/// text item came from. Written as the item's `source: [{"kind": "track", …}]`.
141#[derive(Debug, Clone, PartialEq)]
142pub struct TreeTrack {
143    /// The cue's start offset in seconds (docling's `WebVTTTimestamp.seconds`:
144    /// `h*3600 + m*60 + s + millis/1000.0`, so the float is bit-identical).
145    pub start_time: f64,
146    /// The cue's end offset in seconds.
147    pub end_time: f64,
148    /// The cue identifier line, when the cue has one.
149    pub identifier: Option<String>,
150    /// The `<v …>` voice annotation the text sits in, when any.
151    pub voice: Option<String>,
152}
153
154/// One item of an [`ItemTree`].
155#[derive(Debug, Clone, PartialEq)]
156pub struct TreeItem {
157    /// The parent item's index; `None` = the document body.
158    pub parent: Option<usize>,
159    /// Child item indices, in docling's `children` order.
160    pub children: Vec<usize>,
161    /// The content layer; `None` = `body`.
162    pub layer: Option<ContentLayer>,
163    pub kind: TreeKind,
164    /// The item's `prov` entry, when the backend has page geometry for it
165    /// (`None` → `prov: []`, what the HTML and DOCX backends write).
166    pub prov: Option<TreeProv>,
167    /// docling's `DocItem.comments`: the `comment_section` groups (or note
168    /// text items) annotating this item, as item indices — written after
169    /// `prov` when non-empty.
170    pub comments: Vec<usize>,
171    /// docling's `DocItem.source`: the track segment a text item was taken
172    /// from (WebVTT cues) — written after `prov` when set.
173    pub source: Option<TreeTrack>,
174    /// Removed by [`ItemTree::delete`] (docling's `delete_items`): the slot
175    /// stays so every other index keeps its meaning, but the item is not
176    /// numbered or written.
177    pub deleted: bool,
178}
179
180/// docling's item tree in creation order (see the [module docs](self)).
181#[derive(Debug, Clone, PartialEq, Default)]
182pub struct ItemTree {
183    /// Every item, indexed by creation order — which is how docling numbers
184    /// `#/texts/N`, `#/groups/N`, … within each bucket.
185    pub items: Vec<TreeItem>,
186    /// The body's `children`, as item indices.
187    pub body: Vec<usize>,
188}
189
190impl ItemTree {
191    /// Append an item under `parent` (`None` = body) on `layer`, registering
192    /// it as its parent's last child — docling's `add_*` calls do exactly that.
193    pub fn add(
194        &mut self,
195        parent: Option<usize>,
196        layer: Option<ContentLayer>,
197        kind: TreeKind,
198    ) -> usize {
199        let id = self.items.len();
200        self.items.push(TreeItem {
201            parent,
202            children: Vec::new(),
203            layer,
204            kind,
205            prov: None,
206            comments: Vec::new(),
207            source: None,
208            deleted: false,
209        });
210        match parent {
211            Some(p) => self.items[p].children.push(id),
212            None => self.body.push(id),
213        }
214        id
215    }
216
217    /// [`add`](Self::add) with the item's provenance — docling's
218    /// `add_text(…, prov=prov)`.
219    pub fn add_with_prov(
220        &mut self,
221        parent: Option<usize>,
222        layer: Option<ContentLayer>,
223        kind: TreeKind,
224        prov: TreeProv,
225    ) -> usize {
226        let id = self.add(parent, layer, kind);
227        self.items[id].prov = Some(prov);
228        id
229    }
230
231    /// Append every item of `other` after this tree's, renumbering its
232    /// indices (parents, children, comments, table/picture caption and
233    /// rich-cell refs) and adding its body children to this body — so a
234    /// backend can build independent fragments in parallel (one per PPTX
235    /// slide) and still hand the export one tree in creation order, exactly
236    /// as if it had been built sequentially.
237    pub fn append(&mut self, other: ItemTree) {
238        let off = self.items.len();
239        let shift = |i: usize| i + off;
240        for mut item in other.items {
241            item.parent = item.parent.map(shift);
242            for c in item.children.iter_mut().chain(item.comments.iter_mut()) {
243                *c = shift(*c);
244            }
245            match &mut item.kind {
246                TreeKind::Table {
247                    rich_cells,
248                    captions,
249                    ..
250                } => {
251                    for (_, _, g) in rich_cells.iter_mut() {
252                        *g = shift(*g);
253                    }
254                    for c in captions.iter_mut() {
255                        *c = shift(*c);
256                    }
257                }
258                TreeKind::Picture { captions, .. } => {
259                    for c in captions.iter_mut() {
260                        *c = shift(*c);
261                    }
262                }
263                _ => {}
264            }
265            self.items.push(item);
266        }
267        self.body.extend(other.body.into_iter().map(shift));
268    }
269
270    /// Move `id` under `new_parent`, dropping it from its current parent's
271    /// children and appending it to the new one's — docling's
272    /// `group_cell_elements` re-parenting of a rich cell's items.
273    pub fn reparent(&mut self, id: usize, new_parent: Option<usize>) {
274        let old = self.items[id].parent;
275        let siblings = match old {
276            Some(p) => &mut self.items[p].children,
277            None => &mut self.body,
278        };
279        siblings.retain(|&c| c != id);
280        self.items[id].parent = new_parent;
281        match new_parent {
282            Some(p) => self.items[p].children.push(id),
283            None => self.body.push(id),
284        }
285    }
286
287    /// Re-number the items in traversal order — a pre-order walk of the body
288    /// through every layer, which is how docling's
289    /// `DoclingDocument.concatenate` (and `_normalize_references`) re-creates
290    /// a document's items: a group created after the content it was later
291    /// wrapped around comes before that content afterwards. Items the walk
292    /// does not reach (deleted ones) are dropped. Returns each old id's new
293    /// id.
294    pub fn renumber_in_traversal_order(&mut self) -> Vec<Option<usize>> {
295        let mut order = Vec::with_capacity(self.items.len());
296        let mut stack: Vec<usize> = self.body.iter().rev().copied().collect();
297        while let Some(id) = stack.pop() {
298            if self.items[id].deleted {
299                continue;
300            }
301            order.push(id);
302            stack.extend(self.items[id].children.iter().rev());
303        }
304        let mut new_of: Vec<Option<usize>> = vec![None; self.items.len()];
305        for (new, &old) in order.iter().enumerate() {
306            new_of[old] = Some(new);
307        }
308        let remap = |ids: &mut Vec<usize>| {
309            *ids = ids.iter().filter_map(|&i| new_of[i]).collect();
310        };
311        let mut old_items: Vec<Option<TreeItem>> = std::mem::take(&mut self.items)
312            .into_iter()
313            .map(Some)
314            .collect();
315        for &old in &order {
316            let mut item = old_items[old].take().expect("each item visited once");
317            item.parent = item.parent.and_then(|p| new_of[p]);
318            remap(&mut item.children);
319            remap(&mut item.comments);
320            match &mut item.kind {
321                TreeKind::Table {
322                    rich_cells,
323                    captions,
324                    ..
325                } => {
326                    rich_cells.retain_mut(|(_, _, g)| match new_of[*g] {
327                        Some(n) => {
328                            *g = n;
329                            true
330                        }
331                        None => false,
332                    });
333                    remap(captions);
334                }
335                TreeKind::Picture { captions, .. } => remap(captions),
336                _ => {}
337            }
338            self.items.push(item);
339        }
340        remap(&mut self.body);
341        new_of
342    }
343
344    /// Remove `id` from the tree — docling's `delete_items`, which the DOCX
345    /// backend uses to drop the empty text item a blank spacer paragraph left
346    /// between two items of a resumed list. The item leaves its parent's
347    /// children and is neither numbered nor written; its slot stays so the
348    /// indices held elsewhere stay valid.
349    pub fn delete(&mut self, id: usize) {
350        match self.items[id].parent {
351            Some(p) => self.items[p].children.retain(|&c| c != id),
352            None => self.body.retain(|&c| c != id),
353        }
354        self.items[id].deleted = true;
355    }
356
357    /// The last live text-bucket item (docling's `doc.texts[-1]`).
358    pub fn last_text(&self) -> Option<usize> {
359        self.items.iter().rposition(|it| {
360            !it.deleted && matches!(it.kind, TreeKind::Text { .. } | TreeKind::Code { .. })
361        })
362    }
363
364    /// How many items of a bucket precede `id` — its `#/{bucket}/N` index.
365    pub fn bucket_index(&self, id: usize) -> usize {
366        let same = |k: &TreeKind| {
367            std::mem::discriminant(k) == std::mem::discriminant(&self.items[id].kind)
368                || matches!(
369                    (k, &self.items[id].kind),
370                    (TreeKind::Text { .. }, TreeKind::Code { .. })
371                        | (TreeKind::Code { .. }, TreeKind::Text { .. })
372                )
373        };
374        self.items[..id]
375            .iter()
376            .filter(|it| !it.deleted && same(&it.kind))
377            .count()
378    }
379
380    /// The number of tables created so far (docling's `len(doc.tables)`).
381    pub fn table_count(&self) -> usize {
382        self.items
383            .iter()
384            .filter(|it| !it.deleted && matches!(it.kind, TreeKind::Table { .. }))
385            .count()
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392
393    fn text(t: &str) -> TreeKind {
394        TreeKind::Text {
395            label: "text".into(),
396            text: t.into(),
397            orig: None,
398            formatting: None,
399            hyperlink: None,
400            level: None,
401            list: None,
402        }
403    }
404
405    /// `add` registers the item as its parent's (or the body's) last child;
406    /// `reparent` moves it — a rich cell's items leave the heading they were
407    /// created under for the table's group.
408    #[test]
409    fn add_and_reparent_keep_docling_children_order() {
410        let mut t = ItemTree::default();
411        let title = t.add(None, None, text("Title"));
412        let a = t.add(Some(title), None, text("a"));
413        let b = t.add(Some(title), None, text("b"));
414        let table = t.add(
415            Some(title),
416            None,
417            TreeKind::Table {
418                table: Table::default(),
419                rich_cells: Vec::new(),
420                captions: Vec::new(),
421            },
422        );
423        let group = t.add(
424            Some(table),
425            None,
426            TreeKind::Group {
427                label: "unspecified".into(),
428                name: "rich_cell_group_1_0_0".into(),
429            },
430        );
431        assert_eq!(t.body, vec![title]);
432        assert_eq!(t.items[title].children, vec![a, b, table]);
433        t.reparent(a, Some(group));
434        assert_eq!(t.items[title].children, vec![b, table]);
435        assert_eq!(t.items[group].children, vec![a]);
436        assert_eq!(t.items[a].parent, Some(group));
437        assert_eq!(t.table_count(), 1);
438        // Text and code share the `texts` bucket.
439        let code = t.add(
440            None,
441            None,
442            TreeKind::Code {
443                text: "x".into(),
444                orig: None,
445                language: None,
446                formatting: None,
447                hyperlink: None,
448            },
449        );
450        assert_eq!(t.bucket_index(code), 3, "title, a, b precede it in `texts`");
451        assert_eq!(t.bucket_index(group), 0);
452        assert_eq!(t.body, vec![title, code]);
453    }
454
455    /// `append` renumbers a fragment built on its own (a slide converted in
456    /// parallel) so the merged tree reads as if built in one pass: parents,
457    /// children, comment back-refs and caption refs all shift together.
458    #[test]
459    fn append_renumbers_a_fragment_into_creation_order() {
460        let mut whole = ItemTree::default();
461        let slide0 = whole.add(
462            None,
463            None,
464            TreeKind::Group {
465                label: "chapter".into(),
466                name: "slide-0".into(),
467            },
468        );
469        whole.add(Some(slide0), None, text("first"));
470
471        let mut frag = ItemTree::default();
472        let slide1 = frag.add(
473            None,
474            None,
475            TreeKind::Group {
476                label: "chapter".into(),
477                name: "slide-1".into(),
478            },
479        );
480        let cap = frag.add_with_prov(
481            Some(slide1),
482            None,
483            TreeKind::Text {
484                label: "caption".into(),
485                text: "Title".into(),
486                orig: None,
487                formatting: None,
488                hyperlink: None,
489                level: None,
490                list: None,
491            },
492            TreeProv {
493                page_no: 2,
494                bbox: [1.0, 2.0, 3.0, 4.0],
495                bottom_left: true,
496                charspan: [0, 5],
497            },
498        );
499        let pic = frag.add(
500            Some(slide1),
501            None,
502            TreeKind::Picture {
503                captions: vec![cap],
504                image: None,
505                classification: Some("bar_chart".into()),
506                confidence: None,
507                chart: None,
508                dpi: None,
509            },
510        );
511        let note = frag.add(
512            None,
513            Some(ContentLayer::Notes),
514            TreeKind::Group {
515                label: "comment_section".into(),
516                name: "comment-slide2-1".into(),
517            },
518        );
519        frag.items[pic].comments.push(note);
520
521        whole.append(frag);
522        assert_eq!(whole.body, vec![slide0, 2, 5]);
523        assert_eq!(whole.items[2].children, vec![3, 4]);
524        assert_eq!(whole.items[3].parent, Some(2));
525        assert_eq!(whole.items[3].prov.as_ref().map(|p| p.page_no), Some(2));
526        assert!(
527            matches!(&whole.items[4].kind, TreeKind::Picture { captions, .. } if captions == &[3])
528        );
529        assert_eq!(whole.items[4].comments, vec![5]);
530        assert_eq!(whole.items[5].parent, None);
531        assert_eq!(
532            whole.bucket_index(4),
533            0,
534            "the fragment's picture is #/pictures/0"
535        );
536        assert_eq!(
537            whole.bucket_index(5),
538            2,
539            "slide-0, slide-1 precede it in `groups`"
540        );
541    }
542
543    /// `renumber_in_traversal_order`: docling's `concatenate` re-creates the
544    /// items as a pre-order walk meets them, so a group created after the
545    /// content it was wrapped around (a rich cell's group) comes first, and a
546    /// deleted item disappears; every cross-reference follows.
547    #[test]
548    fn renumbering_follows_the_traversal() {
549        let mut t = ItemTree::default();
550        let a = t.add(None, None, text("a"));
551        let table = t.add(
552            None,
553            None,
554            TreeKind::Table {
555                table: Table::default(),
556                rich_cells: Vec::new(),
557                captions: Vec::new(),
558            },
559        );
560        let cell_text = t.add(None, None, text("cell"));
561        let group = t.add(
562            Some(table),
563            None,
564            TreeKind::Group {
565                label: "unspecified".into(),
566                name: "rich_cell_group_1_0_0".into(),
567            },
568        );
569        t.reparent(cell_text, Some(group));
570        if let TreeKind::Table { rich_cells, .. } = &mut t.items[table].kind {
571            rich_cells.push((0, 0, group));
572        }
573        let gone = t.add(None, None, text("gone"));
574        t.delete(gone);
575        let z = t.add(None, None, text("z"));
576
577        let new_of = t.renumber_in_traversal_order();
578        assert_eq!(
579            new_of,
580            vec![Some(0), Some(1), Some(3), Some(2), None, Some(4)]
581        );
582        assert_eq!(t.items.len(), 5);
583        assert_eq!(t.body, vec![0, 1, 4]);
584        assert_eq!(t.items[1].children, vec![2], "the group follows its table");
585        assert_eq!(t.items[2].parent, Some(1));
586        assert_eq!(
587            t.items[3].parent,
588            Some(2),
589            "the cell text follows its group"
590        );
591        assert!(matches!(&t.items[3].kind, TreeKind::Text { text, .. } if text == "cell"));
592        assert!(
593            matches!(&t.items[1].kind, TreeKind::Table { rich_cells, .. } if rich_cells == &[(0, 0, 2)])
594        );
595        let _ = (a, z);
596    }
597}