use crate::TreeOptions;
use dashmap::DashMap;
use parking_lot::Mutex;
use std::ops::{Deref, Index, IndexMut};
use std::{sync::Arc, time::SystemTime};
#[derive(Clone, Debug)]
pub struct Root {
pub(crate) inner: Arc<Mutex<Item>>,
}
impl Root {
pub fn new() -> Root {
TreeOptions::default().into()
}
pub fn messages_capacity(&self) -> usize {
self.inner.lock().messages.lock().buf.capacity()
}
pub fn num_tasks(&self) -> usize {
self.inner.lock().tree.len()
}
pub fn add_child(&self, name: impl Into<String>) -> Item {
self.inner.lock().add_child(name)
}
pub fn sorted_snapshot(&self, out: &mut Vec<(Key, Value)>) {
out.clear();
out.extend(
self.inner
.lock()
.tree
.iter()
.map(|r| (r.key().clone(), r.value().clone())),
);
out.sort_by_key(|t| t.0);
}
pub fn copy_messages(&self, out: &mut Vec<Message>) {
self.inner.lock().messages.lock().copy_into(out);
}
pub fn deep_clone(&self) -> Root {
Root {
inner: Arc::new(Mutex::new(self.inner.lock().deep_clone())),
}
}
pub fn deep_eq(&self, other: &Root) -> bool {
self.inner.lock().deep_eq(other.inner.lock().deref())
}
}
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub enum MessageLevel {
Info,
Failure,
Success,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Message {
pub time: SystemTime,
pub level: MessageLevel,
pub origin: String,
pub message: String,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub(crate) struct MessageRingBuffer {
buf: Vec<Message>,
cursor: usize,
}
impl MessageRingBuffer {
pub fn with_capacity(capacity: usize) -> MessageRingBuffer {
MessageRingBuffer {
buf: Vec::with_capacity(capacity),
cursor: 0,
}
}
fn has_capacity(&self) -> bool {
self.buf.len() < self.buf.capacity()
}
pub fn push_overwrite(
&mut self,
level: MessageLevel,
origin: String,
message: impl Into<String>,
) {
let msg = Message {
time: SystemTime::now(),
level,
origin,
message: message.into(),
};
if self.has_capacity() {
self.buf.push(msg)
} else {
self.buf[self.cursor] = msg;
self.cursor = (self.cursor + 1) % self.buf.len();
}
}
pub fn copy_into(&self, out: &mut Vec<Message>) {
out.clear();
if self.has_capacity() {
out.extend_from_slice(self.buf.as_slice());
} else {
out.extend_from_slice(&self.buf[(self.cursor + 1) % self.buf.len()..]);
if self.cursor + 1 != self.buf.len() {
out.extend_from_slice(&self.buf[..self.cursor]);
}
}
}
}
#[derive(Debug)]
pub struct Item {
pub(crate) key: Key,
pub(crate) highest_child_id: ItemId,
pub(crate) tree: Arc<DashMap<Key, Value>>,
pub(crate) messages: Arc<Mutex<MessageRingBuffer>>,
}
impl Drop for Item {
fn drop(&mut self) {
self.tree.remove(&self.key);
}
}
impl Item {
pub fn init(&mut self, max: Option<ProgressStep>, unit: Option<&'static str>) {
self.tree.get_mut(&self.key).map(|mut r| {
r.value_mut().progress = Some(Progress {
done_at: max,
unit,
..Default::default()
})
});
}
fn alter_progress(&mut self, f: impl FnMut(&mut Progress)) {
self.tree.get_mut(&self.key).map(|mut r| {
r.value_mut().progress.as_mut().map(f);
});
}
pub fn set_name(&mut self, name: impl Into<String>) {
self.tree.get_mut(&self.key).map(|mut r| {
r.value_mut().name = name.into();
});
}
pub fn name(&self) -> Option<String> {
self.tree.get(&self.key).map(|r| r.value().name.to_owned())
}
pub fn set(&mut self, step: ProgressStep) {
self.alter_progress(|p| {
p.step = step;
p.state = ProgressState::Running;
});
}
pub fn blocked(&mut self, reason: &'static str, eta: Option<SystemTime>) {
self.alter_progress(|p| p.state = ProgressState::Blocked(reason, eta));
}
pub fn halted(&mut self, reason: &'static str, eta: Option<SystemTime>) {
self.alter_progress(|p| p.state = ProgressState::Halted(reason, eta));
}
pub fn add_child(&mut self, name: impl Into<String>) -> Item {
let child_key = self.key.add_child(self.highest_child_id);
self.tree.insert(
child_key,
Value {
name: name.into(),
progress: None,
},
);
self.highest_child_id = self.highest_child_id.wrapping_add(1);
Item {
highest_child_id: 0,
key: child_key,
tree: self.tree.clone(),
messages: self.messages.clone(),
}
}
pub fn message(&mut self, level: MessageLevel, message: impl Into<String>) {
let message: String = message.into();
self.messages.lock().push_overwrite(
level,
{
let name = self
.tree
.get(&self.key)
.map(|v| v.name.to_owned())
.unwrap_or_default();
#[cfg(feature = "log-renderer")]
match level {
MessageLevel::Failure => crate::warn!("{} → {}", name, message),
MessageLevel::Info | MessageLevel::Success => {
crate::info!("{} → {}", name, message)
}
};
name
},
message,
)
}
pub fn done(&mut self, message: impl Into<String>) {
self.message(MessageLevel::Success, message)
}
pub fn fail(&mut self, message: impl Into<String>) {
self.message(MessageLevel::Failure, message)
}
pub fn info(&mut self, message: impl Into<String>) {
self.message(MessageLevel::Info, message)
}
fn deep_clone(&self) -> Item {
Item {
key: self.key.clone(),
highest_child_id: self.highest_child_id,
tree: Arc::new(self.tree.deref().clone()),
messages: Arc::new(Mutex::new(self.messages.lock().clone())),
}
}
fn deep_eq(&self, other: &Item) -> bool {
if !(*self.messages.lock() == *other.messages.lock()) {
return false;
}
for (lhs, rhs) in self.tree.iter().zip(other.tree.iter()) {
if lhs.key() != rhs.key() || lhs.value() != rhs.value() {
return false;
}
}
true
}
}
type ItemId = u16;
pub type Level = u8;
pub type ProgressStep = u32;
#[derive(Copy, Clone, Default, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Key(
Option<ItemId>,
Option<ItemId>,
Option<ItemId>,
Option<ItemId>,
);
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub(crate) enum SiblingLocation {
Above,
Below,
AboveAndBelow,
NotFound,
}
impl SiblingLocation {
fn merge(&mut self, other: SiblingLocation) {
use SiblingLocation::*;
*self = match (*self, other) {
(any, NotFound) => any,
(NotFound, any) => any,
(Above, Below) => AboveAndBelow,
(Below, Above) => AboveAndBelow,
(AboveAndBelow, _) => AboveAndBelow,
(_, AboveAndBelow) => AboveAndBelow,
(Above, Above) => Above,
(Below, Below) => Below,
};
}
}
impl Default for SiblingLocation {
fn default() -> Self {
SiblingLocation::NotFound
}
}
#[derive(Copy, Clone, Default, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub(crate) struct Adjacency(
pub SiblingLocation,
pub SiblingLocation,
pub SiblingLocation,
pub SiblingLocation,
);
impl Adjacency {
pub(crate) fn level(&self) -> Level {
use SiblingLocation::*;
match self {
Adjacency(NotFound, NotFound, NotFound, NotFound) => 0,
Adjacency(_a, NotFound, NotFound, NotFound) => 1,
Adjacency(_a, _b, NotFound, NotFound) => 2,
Adjacency(_a, _b, _c, NotFound) => 3,
Adjacency(_a, _b, _c, _d) => 4,
}
}
pub fn get(&self, level: Level) -> Option<&SiblingLocation> {
Some(match level {
1 => &self.0,
2 => &self.1,
3 => &self.2,
4 => &self.3,
_ => return None,
})
}
pub fn get_mut(&mut self, level: Level) -> Option<&mut SiblingLocation> {
Some(match level {
1 => &mut self.0,
2 => &mut self.1,
3 => &mut self.2,
4 => &mut self.3,
_ => return None,
})
}
}
impl Index<Level> for Adjacency {
type Output = SiblingLocation;
fn index(&self, index: Level) -> &Self::Output {
self.get(index).expect("adjacency index in bound")
}
}
impl IndexMut<Level> for Adjacency {
fn index_mut(&mut self, index: Level) -> &mut Self::Output {
self.get_mut(index).expect("adjacency index in bound")
}
}
impl Key {
pub(crate) fn add_child(self, child_id: ItemId) -> Key {
match self {
Key(None, None, None, None) => Key(Some(child_id), None, None, None),
Key(a, None, None, None) => Key(a, Some(child_id), None, None),
Key(a, b, None, None) => Key(a, b, Some(child_id), None),
Key(a, b, c, None) => Key(a, b, c, Some(child_id)),
Key(a, b, c, _d) => {
crate::warn!("Maximum nesting level reached. Adding tasks to current parent");
Key(a, b, c, Some(child_id))
}
}
}
pub fn level(&self) -> Level {
match self {
Key(None, None, None, None) => 0,
Key(Some(_), None, None, None) => 1,
Key(Some(_), Some(_), None, None) => 2,
Key(Some(_), Some(_), Some(_), None) => 3,
Key(Some(_), Some(_), Some(_), Some(_)) => 4,
_ => unreachable!("This is a bug - Keys follow a certain pattern"),
}
}
fn get(&self, level: Level) -> Option<&ItemId> {
match level {
1 => self.0.as_ref(),
2 => self.1.as_ref(),
3 => self.2.as_ref(),
4 => self.3.as_ref(),
_ => return None,
}
}
pub(crate) fn shares_parent_with(&self, other: &Key, parent_level: Level) -> bool {
if parent_level < 1 {
return true;
}
for level in 1..=parent_level {
if let (Some(lhs), Some(rhs)) = (self.get(level), other.get(level)) {
if lhs != rhs {
return false;
}
} else {
return false;
}
}
return true;
}
pub(crate) fn adjacency(sorted: &[(Key, Value)], index: usize) -> Adjacency {
use SiblingLocation::*;
let key = &sorted[index].0;
let key_level = key.level();
let mut adjecency = Adjacency::default();
if key_level == 0 {
return adjecency;
}
fn search<'a>(
iter: impl Iterator<Item = &'a (Key, Value)>,
key: &Key,
key_level: Level,
current_level: Level,
_id_at_level: ItemId,
) -> Option<usize> {
iter.map(|(k, _)| k)
.take_while(|other| key.shares_parent_with(other, current_level.saturating_sub(1)))
.enumerate()
.find(|(_idx, k)| {
if current_level == key_level {
k.level() == key_level || k.level() + 1 == key_level
} else {
k.level() == current_level
}
})
.map(|(idx, _)| idx)
};
let upward_iter = |from: usize, key: &Key, level: Level, id_at_level: ItemId| {
search(
sorted[..from].iter().rev(),
key,
key_level,
level,
id_at_level,
)
};
let downward_iter = |from: usize, key: &Key, level: Level, id_at_level: ItemId| {
sorted
.get(from + 1..)
.and_then(|s| search(s.iter(), key, key_level, level, id_at_level))
};
{
let mut cursor = index;
for level in (1..=key_level).rev() {
if level == 1 {
adjecency[level].merge(Above);
continue;
}
if let Some(key_offset) = upward_iter(cursor, &key, level, key[level]) {
cursor = index.saturating_sub(key_offset);
adjecency[level].merge(Above);
}
}
}
{
let mut cursor = index;
for level in (1..=key_level).rev() {
if let Some(key_offset) = downward_iter(cursor, &key, level, key[level]) {
cursor = index + key_offset;
adjecency[level].merge(Below);
}
}
}
for level in 1..key_level {
if key_level == 1 && index + 1 == sorted.len() {
continue;
}
adjecency[level] = match adjecency[level] {
Above | Below | NotFound => NotFound,
AboveAndBelow => AboveAndBelow,
};
}
adjecency
}
pub const fn max_level() -> Level {
4
}
}
impl Index<Level> for Key {
type Output = ItemId;
fn index(&self, index: Level) -> &Self::Output {
self.get(index).expect("key index in bound")
}
}
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum ProgressState {
Blocked(&'static str, Option<SystemTime>),
Halted(&'static str, Option<SystemTime>),
Running,
}
impl Default for ProgressState {
fn default() -> Self {
ProgressState::Running
}
}
#[derive(Copy, Clone, Default, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Progress {
pub step: ProgressStep,
pub done_at: Option<ProgressStep>,
pub unit: Option<&'static str>,
pub state: ProgressState,
}
impl Progress {
pub fn fraction(&self) -> Option<f32> {
self.done_at
.map(|done_at| self.step as f32 / done_at as f32)
}
}
#[derive(Clone, Default, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Value {
pub name: String,
pub progress: Option<Progress>,
}