use color_eyre::{Result, eyre::eyre};
use std::{
cell::RefCell,
collections::{HashMap, HashSet, VecDeque, hash_map::Entry},
ops::Deref,
rc::Rc,
};
pub trait Treeable: std::fmt::Debug + std::cmp::Ord {
type ID: std::cmp::Eq + std::hash::Hash;
fn id(&self) -> Self::ID;
fn parent_id(&self) -> Option<Self::ID>;
fn reset_parent(&mut self);
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Tree<T: Treeable> {
pub item: T,
pub subitems: Vec<Tree<T>>,
pub depth: usize,
}
impl<T: Treeable> Deref for Tree<T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.item
}
}
impl<T: Treeable> Ord for Tree<T> {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.item.cmp(&other.item)
}
}
impl<T: Treeable + std::cmp::PartialEq> PartialOrd for Tree<T> {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
#[derive(Debug)]
struct TreeBuilder<Treeable> {
item: Treeable,
parent: Option<()>,
subitems: Vec<Rc<RefCell<TreeBuilder<Treeable>>>>,
}
impl<T: Treeable> TreeBuilder<T> {
fn finalize(self, depth: usize) -> Tree<T> {
let subitems: Vec<Tree<T>> = self
.subitems
.into_iter()
.map(|c| {
Rc::try_unwrap(c)
.expect("should consume single Rc")
.into_inner()
.finalize(depth + 1)
})
.collect();
Tree {
item: self.item,
subitems,
depth,
}
}
}
impl<T: Treeable + std::cmp::Eq> Tree<T> {
pub fn new(item: T) -> Self {
Self {
item,
subitems: vec![],
depth: 0,
}
}
pub fn from_items(items: Vec<T>) -> Result<Vec<Tree<T>>> {
let ids = items.iter().map(|t| t.id()).collect::<HashSet<_>>();
let (top_level_items, mut subitems): (VecDeque<_>, VecDeque<_>) = items
.into_iter()
.map(|mut item| {
if let Some(parent) = item.parent_id()
&& !ids.contains(&parent)
{
item.reset_parent();
}
Rc::new(RefCell::new(TreeBuilder {
item,
parent: None,
subitems: vec![],
}))
})
.partition(|item| item.borrow().item.parent_id().is_none());
let mut items: HashMap<_, Rc<RefCell<TreeBuilder<T>>>> = top_level_items
.into_iter()
.map(|item| (item.borrow().item.id(), item.clone()))
.collect();
while !subitems.is_empty() {
let subitem = subitems.pop_front().unwrap();
let parent = items.entry(
subitem
.borrow()
.item
.parent_id()
.ok_or_else(|| eyre!("Subitem has bad parent assigned"))?,
);
if let Entry::Vacant(_) = parent {
subitems.push_back(subitem);
continue;
}
parent.and_modify(|entry| {
subitem.borrow_mut().parent = Some(());
entry.borrow_mut().subitems.push(subitem.clone())
});
items.insert(subitem.borrow().item.id(), subitem.clone());
}
let items: Result<Vec<_>> = items
.into_iter()
.filter(|(_, c)| c.borrow().parent.is_none())
.collect::<Vec<_>>()
.into_iter()
.map(|(_, c)| {
Ok(Rc::try_unwrap(c)
.map_err(|_| eyre!("Expected single item reference"))?
.into_inner()
.finalize(0))
})
.collect();
items.map(|mut f| {
f.sort();
f
})
}
pub fn flatten(&self) -> Vec<&Tree<T>> {
let mut items = vec![self];
for item in &self.subitems {
items.extend(item.flatten())
}
items
}
pub fn find(&self, id: &<T as Treeable>::ID) -> Option<&Tree<T>> {
if self.item.id() == *id {
return Some(self);
}
for item in &self.subitems {
if let Some(tree) = item.find(id) {
return Some(tree);
}
}
None
}
pub fn find_mut(&mut self, id: &<T as Treeable>::ID) -> Option<&mut Tree<T>> {
if self.item.id() == *id {
return Some(self);
}
for item in &mut self.subitems {
if let Some(tree) = item.find_mut(id) {
return Some(tree);
}
}
None
}
}
pub trait TreeFlattenExt<T: Treeable> {
fn flat_tree(&self) -> Vec<&Tree<T>>;
fn find(&self, id: &T::ID) -> Option<&Tree<T>>;
fn find_mut(&mut self, id: &T::ID) -> Option<&mut Tree<T>>;
fn keep_trees(self, filter_items: &[T::ID]) -> Self
where
Self: Sized;
}
impl<T: Treeable> TreeFlattenExt<T> for Vec<Tree<T>> {
fn flat_tree(&self) -> Vec<&Tree<T>> {
self.iter().flat_map(Tree::flatten).collect()
}
fn find(&self, id: &T::ID) -> Option<&Tree<T>> {
for item in self {
if let Some(item) = item.find(id) {
return Some(item);
}
}
None
}
fn find_mut(&mut self, id: &T::ID) -> Option<&mut Tree<T>> {
for item in self {
if let Some(item) = item.find_mut(id) {
return Some(item);
}
}
None
}
fn keep_trees(self, filter_items: &[T::ID]) -> Self {
fn find_tree<T: Treeable>(tree: &Tree<T>, id: &T::ID, ids: &mut HashSet<T::ID>) {
if tree.find(id).is_none() {
return;
}
ids.insert(tree.id());
for item in &tree.subitems {
find_tree(item, id, ids);
}
}
let mut hs = HashSet::new();
for filter_item in filter_items {
for item in &self {
find_tree(item, filter_item, &mut hs);
}
}
let mut tree = self;
tree.retain(|t| hs.contains(&t.id()));
fn filter<T: Treeable>(tree: &mut Tree<T>, hs: &HashSet<T::ID>) {
tree.subitems.retain(|p| hs.contains(&p.id()));
for item in tree.subitems.iter_mut() {
filter(item, hs);
}
}
for t in tree.iter_mut() {
filter(t, &hs);
}
tree
}
}
#[cfg(test)]
mod tests {
use crate::api::rest::Task;
use super::*;
#[test]
fn test_tree_no_subitems() {
let tasks = vec![
Task::new("1", "one"),
Task::new("2", "two"),
Task::new("3", "three"),
];
let trees = Tree::from_items(tasks).unwrap();
assert_eq!(trees.len(), 3);
}
#[test]
fn test_tree_some_subtasks() {
let tasks = vec![
Task::new("1", "one"),
Task::new("2", "two"),
Task::new("3", "three"),
Task {
parent_id: Some("1".to_string()),
..Task::new("4", "four")
},
];
let trees = Tree::from_items(tasks).unwrap();
assert_eq!(trees.len(), 3);
let task = trees
.iter()
.filter(|t| t.item.id == "1")
.collect::<Vec<_>>();
assert_eq!(task.len(), 1);
let task = task[0];
assert_eq!(task.subitems.len(), 1);
assert_eq!(task.subitems[0].item.id, "4");
for task in trees.into_iter().filter(|t| t.item.id != "1") {
assert_eq!(task.subitems.len(), 0);
}
}
#[test]
fn task_tree_complex_subtasks() {
let tasks = vec![
Task::new("1", "one"),
Task {
parent_id: Some("1".to_string()),
..Task::new("2", "two")
},
Task {
parent_id: Some("2".to_string()),
..Task::new("3", "three")
},
Task {
parent_id: Some("3".to_string()),
..Task::new("4", "four")
},
];
let trees = Tree::from_items(tasks).unwrap();
assert_eq!(trees.len(), 1);
assert_eq!(trees[0].item.id, "1");
assert_eq!(trees[0].depth, 0);
assert_eq!(trees[0].subitems[0].item.id, "2");
assert_eq!(trees[0].subitems[0].depth, 1);
assert_eq!(trees[0].subitems[0].subitems[0].item.id, "3");
assert_eq!(trees[0].subitems[0].subitems[0].depth, 2);
assert_eq!(trees[0].subitems[0].subitems[0].subitems[0].item.id, "4");
assert_eq!(trees[0].subitems[0].subitems[0].subitems[0].depth, 3);
}
#[test]
fn task_tree_no_parent() {
let tasks = vec![
Task {
parent_id: Some("1".to_string()),
..Task::new("2", "two")
},
Task {
parent_id: Some("2".to_string()),
..Task::new("3", "three")
},
];
let trees = Tree::from_items(tasks).unwrap();
assert_eq!(trees.len(), 1);
assert_eq!(trees[0].item.parent_id, None);
assert_eq!(trees[0].subitems[0].item.id, "3");
}
#[test]
fn keep_trees() {
let tasks = vec![
Task {
parent_id: None,
..Task::new("1", "one")
},
Task {
parent_id: Some("1".to_string()),
..Task::new("2", "two")
},
Task {
parent_id: Some("2".to_string()),
..Task::new("3", "three")
},
Task {
parent_id: None,
..Task::new("4", "four")
},
Task {
parent_id: None,
..Task::new("5", "five")
},
];
let trees = Tree::from_items(tasks).unwrap();
let trees = trees.keep_trees(&["3".to_string(), "4".to_string()]);
assert_eq!(trees.len(), 2);
assert_eq!(trees[0].item.id, "1");
assert_eq!(trees[0].item.parent_id, None);
assert_eq!(trees[0].subitems[0].item.id, "2");
assert_eq!(trees[0].subitems[0].subitems[0].item.id, "3");
assert_eq!(trees[1].item.id, "4");
}
}