use crate::{ContentLayer, FieldItem, PictureImage, Script, Table};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Formatting {
pub bold: bool,
pub italic: bool,
pub underline: bool,
pub strikethrough: bool,
pub script: Script,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ListMeta {
pub enumerated: bool,
pub marker: String,
}
#[derive(Debug, Clone, PartialEq)]
pub enum TreeKind {
Text {
label: String,
text: String,
orig: Option<String>,
formatting: Option<Formatting>,
hyperlink: Option<String>,
level: Option<u8>,
list: Option<ListMeta>,
},
Code {
text: String,
orig: Option<String>,
language: Option<String>,
formatting: Option<Formatting>,
hyperlink: Option<String>,
},
Group { label: String, name: String },
Table {
table: Table,
rich_cells: Vec<(usize, usize, usize)>,
captions: Vec<usize>,
},
Picture {
captions: Vec<usize>,
image: Option<PictureImage>,
classification: Option<String>,
confidence: Option<f64>,
chart: Option<Table>,
dpi: Option<u32>,
},
FieldRegion { items: Vec<FieldItem> },
KeyValueGraph {
cells: Vec<crate::GraphCell>,
links: Vec<crate::GraphLink>,
},
}
#[derive(Debug, Clone, PartialEq)]
pub struct TreeProv {
pub page_no: usize,
pub bbox: [f64; 4],
pub bottom_left: bool,
pub charspan: [usize; 2],
}
#[derive(Debug, Clone, PartialEq)]
pub struct TreeTrack {
pub start_time: f64,
pub end_time: f64,
pub identifier: Option<String>,
pub voice: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct TreeItem {
pub parent: Option<usize>,
pub children: Vec<usize>,
pub layer: Option<ContentLayer>,
pub kind: TreeKind,
pub prov: Option<TreeProv>,
pub comments: Vec<usize>,
pub source: Option<TreeTrack>,
pub deleted: bool,
pub notes: Vec<TreeNote>,
pub note_body: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TreeNote {
pub offset: usize,
pub text: String,
}
impl ItemTree {
pub fn place_note_calls(
&self,
first_new: usize,
full_text: &str,
offsets: &[usize],
) -> Vec<Option<(usize, usize)>> {
let texts: Vec<(usize, &str)> = (first_new..self.items.len())
.filter(|&i| !self.items[i].deleted)
.filter_map(|i| match &self.items[i].kind {
TreeKind::Text { text, .. } | TreeKind::Code { text, .. } => {
Some((i, text.as_str()))
}
_ => None,
})
.collect();
let mut spans: Vec<(usize, usize, usize)> = Vec::new();
let mut cursor = 0usize; for &(item, text) in &texts {
if text.is_empty() {
continue;
}
if let Some(pos) = full_text[cursor..].find(text) {
let start_b = cursor + pos;
let start = full_text[..start_b].chars().count();
spans.push((item, start, start + text.chars().count()));
cursor = start_b + text.len();
}
}
offsets
.iter()
.map(|&offset| {
spans
.iter()
.find(|&&(_, start, end)| offset >= start && offset <= end)
.map(|&(item, start, _)| (item, offset - start))
.or_else(
|| match spans.iter().rev().find(|&&(_, _, end)| end <= offset) {
Some(&(item, start, end)) => Some((item, end - start)),
None => spans.first().map(|&(item, _, _)| (item, 0)),
},
)
.or_else(|| {
texts
.last()
.map(|&(item, text)| (item, text.chars().count()))
})
})
.collect()
}
}
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ItemTree {
pub items: Vec<TreeItem>,
pub body: Vec<usize>,
}
impl ItemTree {
pub fn add(
&mut self,
parent: Option<usize>,
layer: Option<ContentLayer>,
kind: TreeKind,
) -> usize {
let id = self.items.len();
self.items.push(TreeItem {
parent,
children: Vec::new(),
layer,
kind,
prov: None,
comments: Vec::new(),
source: None,
deleted: false,
notes: Vec::new(),
note_body: false,
});
match parent {
Some(p) => self.items[p].children.push(id),
None => self.body.push(id),
}
id
}
pub fn add_with_prov(
&mut self,
parent: Option<usize>,
layer: Option<ContentLayer>,
kind: TreeKind,
prov: TreeProv,
) -> usize {
let id = self.add(parent, layer, kind);
self.items[id].prov = Some(prov);
id
}
pub fn append(&mut self, other: ItemTree) {
let off = self.items.len();
let shift = |i: usize| i + off;
for mut item in other.items {
item.parent = item.parent.map(shift);
for c in item.children.iter_mut().chain(item.comments.iter_mut()) {
*c = shift(*c);
}
match &mut item.kind {
TreeKind::Table {
rich_cells,
captions,
..
} => {
for (_, _, g) in rich_cells.iter_mut() {
*g = shift(*g);
}
for c in captions.iter_mut() {
*c = shift(*c);
}
}
TreeKind::Picture { captions, .. } => {
for c in captions.iter_mut() {
*c = shift(*c);
}
}
_ => {}
}
self.items.push(item);
}
self.body.extend(other.body.into_iter().map(shift));
}
pub fn reparent(&mut self, id: usize, new_parent: Option<usize>) {
let old = self.items[id].parent;
let siblings = match old {
Some(p) => &mut self.items[p].children,
None => &mut self.body,
};
siblings.retain(|&c| c != id);
self.items[id].parent = new_parent;
match new_parent {
Some(p) => self.items[p].children.push(id),
None => self.body.push(id),
}
}
pub fn renumber_in_traversal_order(&mut self) -> Vec<Option<usize>> {
let mut order = Vec::with_capacity(self.items.len());
let mut stack: Vec<usize> = self.body.iter().rev().copied().collect();
while let Some(id) = stack.pop() {
if self.items[id].deleted {
continue;
}
order.push(id);
stack.extend(self.items[id].children.iter().rev());
}
let mut new_of: Vec<Option<usize>> = vec![None; self.items.len()];
for (new, &old) in order.iter().enumerate() {
new_of[old] = Some(new);
}
let remap = |ids: &mut Vec<usize>| {
*ids = ids.iter().filter_map(|&i| new_of[i]).collect();
};
let mut old_items: Vec<Option<TreeItem>> = std::mem::take(&mut self.items)
.into_iter()
.map(Some)
.collect();
for &old in &order {
let mut item = old_items[old].take().expect("each item visited once");
item.parent = item.parent.and_then(|p| new_of[p]);
remap(&mut item.children);
remap(&mut item.comments);
match &mut item.kind {
TreeKind::Table {
rich_cells,
captions,
..
} => {
rich_cells.retain_mut(|(_, _, g)| match new_of[*g] {
Some(n) => {
*g = n;
true
}
None => false,
});
remap(captions);
}
TreeKind::Picture { captions, .. } => remap(captions),
_ => {}
}
self.items.push(item);
}
remap(&mut self.body);
new_of
}
pub fn delete(&mut self, id: usize) {
match self.items[id].parent {
Some(p) => self.items[p].children.retain(|&c| c != id),
None => self.body.retain(|&c| c != id),
}
self.items[id].deleted = true;
}
fn delete_subtree(&mut self, id: usize) {
self.delete(id);
let mut stack = self.items[id].children.clone();
while let Some(c) = stack.pop() {
self.items[c].deleted = true;
stack.extend(self.items[c].children.iter().copied());
}
}
pub fn wrap_misplaced_list_items(&mut self) {
let is_list_item = |t: &Self, id: usize| matches!(&t.items[id].kind, TreeKind::Text { label, .. } if label == "list_item");
let in_list_group = |t: &Self, id: usize| {
t.items[id].parent.is_some_and(
|p| matches!(&t.items[p].kind, TreeKind::Group { label, .. } if label == "list"),
)
};
let mut runs: Vec<Vec<usize>> = Vec::new();
let mut prev: Option<usize> = None;
let mut stack: Vec<usize> = self.body.iter().rev().copied().collect();
while let Some(id) = stack.pop() {
if self.items[id].deleted {
continue;
}
if is_list_item(self, id) && !in_list_group(self, id) {
let continues =
prev.is_some_and(|p| is_list_item(self, p) && self.items[p].parent.is_none());
match runs.last_mut() {
Some(run) if continues => run.push(id),
_ => runs.push(vec![id]),
}
}
prev = Some(id);
stack.extend(self.items[id].children.iter().rev().copied());
}
for run in runs.into_iter().rev() {
let parent = self.items[run[0]].parent;
let group = self.add(
parent,
None,
TreeKind::Group {
label: "list".into(),
name: "group".into(),
},
);
let siblings = match parent {
Some(p) => &mut self.items[p].children,
None => &mut self.body,
};
siblings.pop();
let at = siblings
.iter()
.position(|&c| c == run[0])
.unwrap_or(siblings.len());
siblings.insert(at, group);
for &li in &run {
self.delete_subtree(li);
}
for &li in &run {
let copy = TreeItem {
parent: Some(group),
children: Vec::new(),
comments: Vec::new(),
source: None,
deleted: false,
..self.items[li].clone()
};
let id = self.items.len();
self.items.push(copy);
self.items[group].children.push(id);
}
}
}
pub fn last_text(&self) -> Option<usize> {
self.items.iter().rposition(|it| {
!it.deleted && matches!(it.kind, TreeKind::Text { .. } | TreeKind::Code { .. })
})
}
pub fn bucket_index(&self, id: usize) -> usize {
let same = |k: &TreeKind| {
std::mem::discriminant(k) == std::mem::discriminant(&self.items[id].kind)
|| matches!(
(k, &self.items[id].kind),
(TreeKind::Text { .. }, TreeKind::Code { .. })
| (TreeKind::Code { .. }, TreeKind::Text { .. })
)
};
self.items[..id]
.iter()
.filter(|it| !it.deleted && same(&it.kind))
.count()
}
pub fn table_count(&self) -> usize {
self.items
.iter()
.filter(|it| !it.deleted && matches!(it.kind, TreeKind::Table { .. }))
.count()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn text(t: &str) -> TreeKind {
TreeKind::Text {
label: "text".into(),
text: t.into(),
orig: None,
formatting: None,
hyperlink: None,
level: None,
list: None,
}
}
#[test]
fn add_and_reparent_keep_docling_children_order() {
let mut t = ItemTree::default();
let title = t.add(None, None, text("Title"));
let a = t.add(Some(title), None, text("a"));
let b = t.add(Some(title), None, text("b"));
let table = t.add(
Some(title),
None,
TreeKind::Table {
table: Table::default(),
rich_cells: Vec::new(),
captions: Vec::new(),
},
);
let group = t.add(
Some(table),
None,
TreeKind::Group {
label: "unspecified".into(),
name: "rich_cell_group_1_0_0".into(),
},
);
assert_eq!(t.body, vec![title]);
assert_eq!(t.items[title].children, vec![a, b, table]);
t.reparent(a, Some(group));
assert_eq!(t.items[title].children, vec![b, table]);
assert_eq!(t.items[group].children, vec![a]);
assert_eq!(t.items[a].parent, Some(group));
assert_eq!(t.table_count(), 1);
let code = t.add(
None,
None,
TreeKind::Code {
text: "x".into(),
orig: None,
language: None,
formatting: None,
hyperlink: None,
},
);
assert_eq!(t.bucket_index(code), 3, "title, a, b precede it in `texts`");
assert_eq!(t.bucket_index(group), 0);
assert_eq!(t.body, vec![title, code]);
}
#[test]
fn append_renumbers_a_fragment_into_creation_order() {
let mut whole = ItemTree::default();
let slide0 = whole.add(
None,
None,
TreeKind::Group {
label: "chapter".into(),
name: "slide-0".into(),
},
);
whole.add(Some(slide0), None, text("first"));
let mut frag = ItemTree::default();
let slide1 = frag.add(
None,
None,
TreeKind::Group {
label: "chapter".into(),
name: "slide-1".into(),
},
);
let cap = frag.add_with_prov(
Some(slide1),
None,
TreeKind::Text {
label: "caption".into(),
text: "Title".into(),
orig: None,
formatting: None,
hyperlink: None,
level: None,
list: None,
},
TreeProv {
page_no: 2,
bbox: [1.0, 2.0, 3.0, 4.0],
bottom_left: true,
charspan: [0, 5],
},
);
let pic = frag.add(
Some(slide1),
None,
TreeKind::Picture {
captions: vec![cap],
image: None,
classification: Some("bar_chart".into()),
confidence: None,
chart: None,
dpi: None,
},
);
let note = frag.add(
None,
Some(ContentLayer::Notes),
TreeKind::Group {
label: "comment_section".into(),
name: "comment-slide2-1".into(),
},
);
frag.items[pic].comments.push(note);
whole.append(frag);
assert_eq!(whole.body, vec![slide0, 2, 5]);
assert_eq!(whole.items[2].children, vec![3, 4]);
assert_eq!(whole.items[3].parent, Some(2));
assert_eq!(whole.items[3].prov.as_ref().map(|p| p.page_no), Some(2));
assert!(
matches!(&whole.items[4].kind, TreeKind::Picture { captions, .. } if captions == &[3])
);
assert_eq!(whole.items[4].comments, vec![5]);
assert_eq!(whole.items[5].parent, None);
assert_eq!(
whole.bucket_index(4),
0,
"the fragment's picture is #/pictures/0"
);
assert_eq!(
whole.bucket_index(5),
2,
"slide-0, slide-1 precede it in `groups`"
);
}
#[test]
fn renumbering_follows_the_traversal() {
let mut t = ItemTree::default();
let a = t.add(None, None, text("a"));
let table = t.add(
None,
None,
TreeKind::Table {
table: Table::default(),
rich_cells: Vec::new(),
captions: Vec::new(),
},
);
let cell_text = t.add(None, None, text("cell"));
let group = t.add(
Some(table),
None,
TreeKind::Group {
label: "unspecified".into(),
name: "rich_cell_group_1_0_0".into(),
},
);
t.reparent(cell_text, Some(group));
if let TreeKind::Table { rich_cells, .. } = &mut t.items[table].kind {
rich_cells.push((0, 0, group));
}
let gone = t.add(None, None, text("gone"));
t.delete(gone);
let z = t.add(None, None, text("z"));
let new_of = t.renumber_in_traversal_order();
assert_eq!(
new_of,
vec![Some(0), Some(1), Some(3), Some(2), None, Some(4)]
);
assert_eq!(t.items.len(), 5);
assert_eq!(t.body, vec![0, 1, 4]);
assert_eq!(t.items[1].children, vec![2], "the group follows its table");
assert_eq!(t.items[2].parent, Some(1));
assert_eq!(
t.items[3].parent,
Some(2),
"the cell text follows its group"
);
assert!(matches!(&t.items[3].kind, TreeKind::Text { text, .. } if text == "cell"));
assert!(
matches!(&t.items[1].kind, TreeKind::Table { rich_cells, .. } if rich_cells == &[(0, 0, 2)])
);
let _ = (a, z);
}
}