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 records the result here — an arena of items in
15//! creation order, each with its parent and children — and the JSON export
16//! ([`DoclingDocument::export_to_json`](crate::DoclingDocument::export_to_json))
17//! serializes this tree instead of deriving one from the nodes. Every other
18//! serializer keeps reading the flat nodes, so their output is unaffected.
19
20use crate::{ContentLayer, FieldItem, PictureImage, Script, Table};
21
22/// docling-core's `Formatting`: the inline styles an item carries in JSON.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
24pub struct Formatting {
25 pub bold: bool,
26 pub italic: bool,
27 pub underline: bool,
28 pub strikethrough: bool,
29 pub script: Script,
30}
31
32/// A `list_item`'s docling fields.
33#[derive(Debug, Clone, PartialEq, Eq, Default)]
34pub struct ListMeta {
35 pub enumerated: bool,
36 /// docling's `marker` — the HTML backend writes `""` unless an ordered
37 /// list carries an explicit `start`, then `"{n}."`.
38 pub marker: String,
39}
40
41/// What an item in the tree is. Mirrors the docling-core item classes the
42/// JSON `texts` / `groups` / `tables` / `pictures` / `field_regions` buckets
43/// hold.
44#[derive(Debug, Clone, PartialEq)]
45pub enum TreeKind {
46 /// A `TextItem` / `TitleItem` / `SectionHeaderItem` / `ListItem`, told
47 /// apart by `label` (`text`, `title`, `section_header`, `list_item`,
48 /// `caption`, `checkbox_selected`, `checkbox_unselected`, …).
49 Text {
50 label: String,
51 text: String,
52 /// docling's `orig` when it differs from `text` (the heading text
53 /// before unicode cleanup, say); `None` = same as `text`.
54 orig: Option<String>,
55 formatting: Option<Formatting>,
56 hyperlink: Option<String>,
57 /// `section_header` only: docling's heading level.
58 level: Option<u8>,
59 /// `list_item` only.
60 list: Option<ListMeta>,
61 },
62 /// A `CodeItem`.
63 Code {
64 text: String,
65 orig: Option<String>,
66 /// The language hint (a highlighter class token such as `python`),
67 /// mapped onto docling's `CodeLanguageLabel` at export; `None` →
68 /// `unknown`.
69 language: Option<String>,
70 formatting: Option<Formatting>,
71 hyperlink: Option<String>,
72 },
73 /// A `GroupItem`: `label` is docling's `GroupLabel` value (`inline`,
74 /// `list`, `section`, `unspecified`, …), `name` its name (`group`, `list`,
75 /// `ordered list`, `header-2`, `rich_cell_group_1_0_3`, …).
76 Group { label: String, name: String },
77 /// A `TableItem`. `rich_cells` marks the cells docling serialized as a
78 /// `RichTableCell`: `(row, col)` grid anchor → the group item (a child of
79 /// the table) that holds the cell's content. `captions` are caption text
80 /// items in the tree.
81 Table {
82 table: Table,
83 rich_cells: Vec<(usize, usize, usize)>,
84 captions: Vec<usize>,
85 },
86 /// A `PictureItem`, its caption text items and optional payload.
87 /// `classification` is a `PictureClassificationLabel` value written as
88 /// the picture's `meta.classification` (an HTML `<stamp>` / `<signature>`).
89 Picture {
90 captions: Vec<usize>,
91 image: Option<PictureImage>,
92 classification: Option<String>,
93 },
94 /// A form key-value region (`field_regions` / `field_items`).
95 FieldRegion { items: Vec<FieldItem> },
96}
97
98/// One item of an [`ItemTree`].
99#[derive(Debug, Clone, PartialEq)]
100pub struct TreeItem {
101 /// The parent item's index; `None` = the document body.
102 pub parent: Option<usize>,
103 /// Child item indices, in docling's `children` order.
104 pub children: Vec<usize>,
105 /// The content layer; `None` = `body`.
106 pub layer: Option<ContentLayer>,
107 pub kind: TreeKind,
108}
109
110/// docling's item tree in creation order (see the [module docs](self)).
111#[derive(Debug, Clone, PartialEq, Default)]
112pub struct ItemTree {
113 /// Every item, indexed by creation order — which is how docling numbers
114 /// `#/texts/N`, `#/groups/N`, … within each bucket.
115 pub items: Vec<TreeItem>,
116 /// The body's `children`, as item indices.
117 pub body: Vec<usize>,
118}
119
120impl ItemTree {
121 /// Append an item under `parent` (`None` = body) on `layer`, registering
122 /// it as its parent's last child — docling's `add_*` calls do exactly that.
123 pub fn add(
124 &mut self,
125 parent: Option<usize>,
126 layer: Option<ContentLayer>,
127 kind: TreeKind,
128 ) -> usize {
129 let id = self.items.len();
130 self.items.push(TreeItem {
131 parent,
132 children: Vec::new(),
133 layer,
134 kind,
135 });
136 match parent {
137 Some(p) => self.items[p].children.push(id),
138 None => self.body.push(id),
139 }
140 id
141 }
142
143 /// Move `id` under `new_parent`, dropping it from its current parent's
144 /// children and appending it to the new one's — docling's
145 /// `group_cell_elements` re-parenting of a rich cell's items.
146 pub fn reparent(&mut self, id: usize, new_parent: Option<usize>) {
147 let old = self.items[id].parent;
148 let siblings = match old {
149 Some(p) => &mut self.items[p].children,
150 None => &mut self.body,
151 };
152 siblings.retain(|&c| c != id);
153 self.items[id].parent = new_parent;
154 match new_parent {
155 Some(p) => self.items[p].children.push(id),
156 None => self.body.push(id),
157 }
158 }
159
160 /// How many items of a bucket precede `id` — its `#/{bucket}/N` index.
161 pub fn bucket_index(&self, id: usize) -> usize {
162 let same = |k: &TreeKind| {
163 std::mem::discriminant(k) == std::mem::discriminant(&self.items[id].kind)
164 || matches!(
165 (k, &self.items[id].kind),
166 (TreeKind::Text { .. }, TreeKind::Code { .. })
167 | (TreeKind::Code { .. }, TreeKind::Text { .. })
168 )
169 };
170 self.items[..id].iter().filter(|it| same(&it.kind)).count()
171 }
172
173 /// The number of tables created so far (docling's `len(doc.tables)`).
174 pub fn table_count(&self) -> usize {
175 self.items
176 .iter()
177 .filter(|it| matches!(it.kind, TreeKind::Table { .. }))
178 .count()
179 }
180}
181
182#[cfg(test)]
183mod tests {
184 use super::*;
185
186 fn text(t: &str) -> TreeKind {
187 TreeKind::Text {
188 label: "text".into(),
189 text: t.into(),
190 orig: None,
191 formatting: None,
192 hyperlink: None,
193 level: None,
194 list: None,
195 }
196 }
197
198 /// `add` registers the item as its parent's (or the body's) last child;
199 /// `reparent` moves it — a rich cell's items leave the heading they were
200 /// created under for the table's group.
201 #[test]
202 fn add_and_reparent_keep_docling_children_order() {
203 let mut t = ItemTree::default();
204 let title = t.add(None, None, text("Title"));
205 let a = t.add(Some(title), None, text("a"));
206 let b = t.add(Some(title), None, text("b"));
207 let table = t.add(
208 Some(title),
209 None,
210 TreeKind::Table {
211 table: Table::default(),
212 rich_cells: Vec::new(),
213 captions: Vec::new(),
214 },
215 );
216 let group = t.add(
217 Some(table),
218 None,
219 TreeKind::Group {
220 label: "unspecified".into(),
221 name: "rich_cell_group_1_0_0".into(),
222 },
223 );
224 assert_eq!(t.body, vec![title]);
225 assert_eq!(t.items[title].children, vec![a, b, table]);
226 t.reparent(a, Some(group));
227 assert_eq!(t.items[title].children, vec![b, table]);
228 assert_eq!(t.items[group].children, vec![a]);
229 assert_eq!(t.items[a].parent, Some(group));
230 assert_eq!(t.table_count(), 1);
231 // Text and code share the `texts` bucket.
232 let code = t.add(
233 None,
234 None,
235 TreeKind::Code {
236 text: "x".into(),
237 orig: None,
238 language: None,
239 formatting: None,
240 hyperlink: None,
241 },
242 );
243 assert_eq!(t.bucket_index(code), 3, "title, a, b precede it in `texts`");
244 assert_eq!(t.bucket_index(group), 0);
245 assert_eq!(t.body, vec![title, code]);
246 }
247}