use itertools::Itertools;
use liwe::{
graph::{Graph, GraphContext},
model::{node::NodePointer, tree::Tree, Key, NodeId},
};
use std::collections::{HashMap, HashSet};
#[derive(Default, PartialEq, Debug, Clone)]
pub struct GraphData {
pub sections: HashMap<NodeId, Section>,
pub documents: HashMap<NodeId, Document>,
pub section_to_section: Vec<(NodeId, NodeId)>,
pub section_to_document: Vec<(NodeId, NodeId)>,
pub document_to_document: Vec<(NodeId, NodeId)>,
}
impl GraphData {
pub fn merge(&mut self, other: GraphData) {
self.sections.extend(other.sections);
self.documents.extend(other.documents);
self.section_to_section.extend(other.section_to_section);
self.section_to_document.extend(other.section_to_document);
self.document_to_document.extend(other.document_to_document);
}
}
#[derive(PartialEq, Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Section {
pub id: NodeId,
pub title: String,
pub key: String,
pub depth: u8,
}
#[derive(PartialEq, Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Document {
pub id: NodeId,
pub title: String,
pub key: String,
pub depth: u8,
}
#[derive(Default, Debug, PartialEq)]
struct GraphCache {
sections: HashMap<NodeId, Section>,
documents: HashMap<NodeId, Document>,
section_to_section: HashSet<(NodeId, NodeId)>,
section_to_document: HashSet<(NodeId, Key)>,
document_to_document: HashSet<(Key, Key)>,
backlinks: HashSet<(Key, Key)>,
}
pub fn graph_data(key_filter: Vec<Key>, depth: u8, graph: &Graph) -> GraphData {
let keys = filter_keys(graph, key_filter, depth);
keys.iter()
.map(|pair| build_graph_data(graph, pair.0, *pair.1))
.fold(
GraphData {
sections: HashMap::new(),
documents: HashMap::new(),
section_to_section: Vec::new(),
section_to_document: Vec::new(),
document_to_document: Vec::new(),
},
|mut acc, data| {
acc.merge(data);
acc
},
)
}
fn build_graph_data(graph: &Graph, key: &Key, key_depth: u8) -> GraphData {
let tree = graph.collect(key);
let mut cache = GraphCache {
documents: HashMap::new(),
sections: HashMap::new(),
section_to_section: HashSet::new(),
section_to_document: HashSet::new(),
document_to_document: HashSet::new(),
backlinks: HashSet::new(),
};
build_sections(&key.to_string(), &mut cache, key_depth, 0, 100, &tree);
GraphData {
sections: cache.sections.clone(),
documents: cache.documents.clone(),
section_to_section: cache.section_to_section.into_iter().collect_vec(),
section_to_document: cache
.section_to_document
.into_iter()
.map(|r| (r.0, resolve_key(graph, &r.1)))
.collect_vec(),
document_to_document: cache
.document_to_document
.into_iter()
.map(|r| (resolve_key(graph, &r.0), resolve_key(graph, &r.1)))
.collect_vec(),
}
}
fn build_sections(
key: &str,
cache: &mut GraphCache,
key_depth: u8,
depth: u8,
_max_depth: u8,
tree: &Tree,
) {
tree.children.iter().for_each(|child| {
let Some(child_id) = child.id else {
return;
};
if child.is_list() {
} else if depth == 0 && child.is_section() {
cache.documents.insert(
child_id,
Document {
depth: key_depth,
id: child_id,
title: child.node.plain_text(),
key: key.to_string(),
},
);
build_sections(key, cache, key_depth, depth + 1, _max_depth, child);
} else if child.is_section() {
cache.sections.insert(
child_id,
Section {
depth,
id: child_id,
title: child.node.plain_text(),
key: key.to_string(),
},
);
if tree.is_section() {
if let Some(tree_id) = tree.id {
cache.section_to_section.insert((tree_id, child_id));
}
}
build_sections(key, cache, key_depth, depth + 1, _max_depth, child);
} else if child.is_reference() {
if let Some(ref_key) = child.node.reference_key() {
if let Some(tree_id) = tree.id {
cache.section_to_document.insert((tree_id, ref_key.clone()));
}
cache.document_to_document.insert((key.into(), ref_key));
}
}
})
}
fn filter_keys(graph: &Graph, key_filter: Vec<Key>, depth_limit: u8) -> HashMap<Key, u8> {
if key_filter.is_empty() {
top_level_keys(graph, depth_limit)
} else {
let mut result = HashMap::new();
for key in key_filter {
for (k, d) in filter_key(graph, key, depth_limit) {
result.entry(k).or_insert(d);
}
}
result
}
}
fn top_level_keys(graph: &Graph, depth_limit: u8) -> HashMap<Key, u8> {
let top_level = graph
.paths()
.into_iter()
.filter(|n| n.ids().len() <= 1_usize)
.filter_map(|n| {
let first_id = n.first_id()?;
let key = (&graph).node(first_id).node_key();
if graph.maybe_key(&key).is_some() {
Some((key, 0))
} else {
None
}
});
let mut keys = HashMap::new();
top_level.for_each(|k| get_keys_for_depth(graph, &k.0, depth_limit, &mut keys));
keys
}
fn filter_key(graph: &Graph, key: Key, depth_limit: u8) -> HashMap<Key, u8> {
if graph.maybe_key(&key).is_none() {
return HashMap::new();
}
let mut keys = HashMap::new();
get_keys_for_depth(graph, &key, depth_limit, &mut keys);
for key in keys.keys().cloned().collect::<Vec<_>>() {
let references = graph.get_inclusion_edges_to(&key);
for node_id in &references {
let ref_key = graph.node(*node_id).node_key();
keys.entry(ref_key).or_insert(0);
}
}
keys
}
fn get_keys_for_depth(graph: &Graph, key: &Key, depth: u8, collected_keys: &mut HashMap<Key, u8>) {
collected_keys.insert(key.clone(), depth);
if depth == 0 {
return;
}
let keys = graph.collect(key).get_all_inclusion_edge_keys();
for k in keys {
if !collected_keys.contains_key(&k) && graph.maybe_key(&k).is_some() {
get_keys_for_depth(graph, &k, depth - 1, collected_keys);
}
}
}
fn resolve_key(graph: &Graph, key: &Key) -> NodeId {
graph
.get_node_id(&key.clone())
.map(|doc_id| graph.node(doc_id).child_id().unwrap_or_default())
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use liwe::model::node::{Node, Reference, ReferenceType};
use pretty_assertions::assert_eq;
use super::*;
#[test]
fn build_sections_one_doc() {
let doc1 = Tree::new(
Some(1),
Node::Document(Key::name("1"), None),
vec![Tree::new(
Some(2),
Node::Section(vec!["title".to_string().into()]),
vec![
Tree::new(
Some(3),
Node::Section(vec!["1.1".to_string().into()]),
vec![Tree::new(
Some(4),
Node::Reference(Reference {
key: Key::name("2"),
text: "".into(),
reference_type: ReferenceType::Regular,
url: String::new(),
display_url: None,
}),
vec![],
)],
),
Tree::new(
Some(5),
Node::Section(vec!["1.2".to_string().into()]),
vec![Tree::new(
Some(6),
Node::Reference(Reference {
key: Key::name("3"),
text: "".into(),
reference_type: ReferenceType::Regular,
url: String::new(),
display_url: None,
}),
vec![],
)],
),
],
)],
);
let mut cache = GraphCache::default();
build_sections("1", &mut cache, 2, 0, 100, &doc1);
assert_eq!(
GraphCache {
backlinks: HashSet::new(),
sections: vec![
(
3,
Section {
id: 3,
title: "1.1".into(),
key: "1".into(),
depth: 1,
}
),
(
5,
Section {
id: 5,
title: "1.2".into(),
key: "1".into(),
depth: 1,
}
)
]
.into_iter()
.collect(),
documents: vec![(
2,
Document {
id: 2,
title: "title".into(),
key: "1".into(),
depth: 2,
}
)]
.into_iter()
.collect(),
section_to_section: vec![(2, 5), (2, 3)].into_iter().collect(),
section_to_document: vec![(3, Key::name("2")), (5, Key::name("3"))]
.into_iter()
.collect(),
document_to_document: vec![
(Key::name("1"), Key::name("2")),
(Key::name("1"), Key::name("3"))
]
.into_iter()
.collect(),
},
cache
);
}
}