use eyre::{bail, ContextCompat, Result};
use nodo::{
codelet::{ScheduleId, Transition},
prelude::{ParameterDataType, ParameterValue, ParameterWithPropertiesSet},
};
use nodo_runtime::proto::nodo as nodo_pb;
use std::{cell::RefCell, cmp::Ordering, collections::HashMap, iter, rc::Rc};
#[derive(Default)]
pub struct NodoAppTree {
pub last_report_app_meta: Option<nodo_pb::AppMeta>,
pub schedules: HashMap<String, NodoAppTreeSchedule>,
pub monitors: Vec<(String, nodo_pb::MonitorMeta, nodo_pb::Monitor)>,
pub step_duration: f32,
}
impl NodoAppTree {
pub fn update_step_duration(&mut self) {
self.step_duration = 0.0;
for (_, sched) in self.schedules.iter_mut() {
sched.step_duration = 0.0;
for (_, seq) in sched.seqs.iter_mut() {
seq.step_duration = 0.0;
for (_, node) in seq.nodes.iter() {
let step_transition = &node.statistics.transitions[Transition::Step.index()];
let delta = step_transition.duration.context("TODO").unwrap().total;
seq.step_duration += delta;
}
sched.step_duration += seq.step_duration;
}
self.step_duration += sched.step_duration;
}
}
}
pub struct NodoAppTreeSchedule {
pub id: ScheduleId,
pub name: String,
pub target_period: Option<f32>,
pub status: nodo_pb::ScheduleStatus,
pub is_expanded: bool,
pub is_selected: bool,
pub tree_row_idx: RefCell<Option<usize>>,
pub seqs: HashMap<String, NodoAppTreeSequence>,
pub step_duration: f32,
}
pub struct NodoAppTreeSequence {
pub name: String,
pub is_expanded: bool,
pub is_selected: bool,
pub tree_row_idx: RefCell<Option<usize>>,
pub nodes: HashMap<String, NodoAppTreeNode>,
pub step_duration: f32,
}
pub struct NodoAppTreeNode {
pub info: nodo_pb::NodeMeta,
pub status: nodo_pb::NodeStatus,
pub statistics: nodo_pb::NodeStatistics,
pub signals: Vec<(String, nodo_pb::NodeSignal)>,
pub parameters: HashMap<String, NodoAppTreeParameter>,
pub is_selected: bool,
pub tree_row_idx: RefCell<Option<usize>>,
}
impl NodoAppTreeNode {
pub fn step_duration(&self) -> f32 {
self.statistics.transitions[Transition::Step.index()]
.duration
.unwrap()
.total
}
}
pub struct NodoAppTreeParameter {
pub name: String,
pub dtype: ParameterDataType,
pub is_mutable: bool,
pub value: ParameterValue,
}
impl NodoAppTree {
pub fn update(&mut self, report: nodo_pb::Report) -> Result<()> {
if let Some(app_meta) = report.meta {
self.last_report_app_meta = Some(app_meta);
}
if self.last_report_app_meta.is_none() {
bail!("app info not available");
}
let (
Some(nodo_pb::AppMeta {
nodes: node_meta,
monitors: monitor_meta,
schedules: schedule_meta,
}),
Some(nodo_pb::AppStatus {
nodes: status,
schedules: schedule_status,
}),
Some(nodo_pb::AppStatistics { nodes: statistics }),
Some(nodo_pb::AppSignals { nodes: signals }),
) = (
&self.last_report_app_meta,
report.status,
report.statistics,
report.signals,
)
else {
bail!("invalid report");
};
if node_meta.len() != status.len()
|| status.len() != statistics.len()
|| statistics.len() != signals.len()
{
bail!("invalid report");
}
if let Some(monitors) = report.monitors.map(|am| am.monitors) {
if monitor_meta.len() != monitors.len() {
bail!("invalid report");
}
self.monitors = monitor_meta
.iter()
.zip(monitors.into_iter())
.map(|(meta, mon)| {
let node_name = node_meta
.iter()
.find(|i| i.id == meta.node_id)
.map(|i| i.node_name.clone())
.unwrap_or_else(|| "(unknown)".into());
(node_name, meta.clone(), mon)
})
.collect();
}
Self::retain_nodes(&mut self.schedules, |nid| {
node_meta.iter().any(|n| n.id == nid)
});
self.schedules
.retain(|needle, _| schedule_meta.iter().any(|sched| &sched.name == needle));
for sched in schedule_meta.iter() {
if !self.schedules.contains_key(&sched.name) {
self.schedules.insert(
sched.name.clone(),
NodoAppTreeSchedule {
id: ScheduleId(sched.id as usize),
name: sched.name.clone(),
target_period: sched.period,
status: nodo_pb::ScheduleStatus::default(),
is_expanded: true,
is_selected: false,
tree_row_idx: RefCell::new(None),
seqs: HashMap::new(),
step_duration: 0.,
},
);
}
}
for status in schedule_status.iter() {
match self
.schedules
.values_mut()
.find(|v| v.id.0 == status.id as usize)
{
Some(sched) => {
sched.status = status.clone();
}
None => {
}
}
}
for (info, (status, (statistics, nodo_pb::NodeSignals { signals }))) in
node_meta.iter().zip(
status
.into_iter()
.zip(statistics.into_iter().zip(signals.into_iter())),
)
{
let Some(worker) = self.schedules.get_mut(&info.schedule_name) else {
continue;
};
let sequence = worker
.seqs
.entry(info.sequence_name.clone())
.or_insert_with_key(|key| NodoAppTreeSequence {
name: key.into(),
is_expanded: true,
is_selected: false,
tree_row_idx: RefCell::new(None),
nodes: HashMap::default(),
step_duration: 0.0,
});
let named_signals = info
.signal_names
.iter()
.cloned()
.zip(signals.iter().cloned())
.collect::<Vec<_>>();
let node = sequence
.nodes
.entry(info.node_name.clone())
.or_insert_with_key(|_| NodoAppTreeNode {
info: info.clone(),
status: status.clone(),
statistics: statistics.clone(),
signals: named_signals.clone(),
parameters: HashMap::default(),
is_selected: false,
tree_row_idx: RefCell::new(None),
});
node.status = status;
node.statistics = statistics;
node.signals = named_signals;
}
self.update_step_duration();
Ok(())
}
pub fn update_config_parameters(&mut self, items: ParameterWithPropertiesSet<String, String>) {
for (_, worker) in self.schedules.iter_mut() {
for (_, seq) in worker.seqs.iter_mut() {
for (node_name, node) in seq.nodes.iter_mut() {
for (key, (props, value)) in items.iter() {
if node_name == key.node() {
let tp =
node.parameters
.entry(key.param().clone())
.or_insert_with(|| NodoAppTreeParameter {
name: key.param().clone(),
dtype: props.dtype,
is_mutable: props.is_mutable,
value: value.clone(),
});
tp.dtype = value.dtype();
tp.value = value.clone();
}
}
}
}
}
}
pub fn iter_workers(&self) -> impl Iterator<Item = &NodoAppTreeSchedule> {
self.schedules.iter().map(|(_, v)| v)
}
pub fn iter_workers_mut(&mut self) -> impl Iterator<Item = &mut NodoAppTreeSchedule> {
self.schedules.iter_mut().map(|(_, v)| v)
}
pub fn iter_sequences(&self) -> impl Iterator<Item = (&str, &NodoAppTreeSequence)> {
self.schedules
.iter()
.flat_map(|(sched_name, w)| w.seqs.iter().map(|(_, seq)| (sched_name.as_str(), seq)))
}
pub fn iter_sequences_mut(&mut self) -> impl Iterator<Item = (&str, &mut NodoAppTreeSequence)> {
self.schedules.iter_mut().flat_map(|(sched_name, w)| {
w.seqs
.iter_mut()
.map(move |(_, seq)| (sched_name.as_str(), seq))
})
}
pub fn iter_nodes(
&self,
) -> impl Iterator<Item = (&NodoAppTreeSchedule, &NodoAppTreeSequence, &NodoAppTreeNode)> {
self.schedules.iter().flat_map(|(_, w)| {
w.seqs
.iter()
.flat_map(move |(_, seq)| seq.nodes.iter().map(move |(_, node)| (w, seq, node)))
})
}
pub fn iter_parameters(
&self,
) -> impl Iterator<
Item = (
&NodoAppTreeSchedule,
&NodoAppTreeSequence,
&NodoAppTreeNode,
&NodoAppTreeParameter,
),
> {
self.schedules.iter().flat_map(|(_, w)| {
w.seqs.iter().flat_map(move |(_, seq)| {
seq.nodes.iter().flat_map(move |(_, node)| {
node.parameters
.iter()
.map(move |(_, param)| (w, seq, node, param))
})
})
})
}
pub fn iter_signals(
&self,
) -> impl Iterator<
Item = (
&NodoAppTreeSchedule,
&NodoAppTreeSequence,
&NodoAppTreeNode,
&(String, nodo_pb::NodeSignal),
),
> {
self.schedules.iter().flat_map(|(_, w)| {
w.seqs.iter().flat_map(move |(_, seq)| {
seq.nodes.iter().flat_map(move |(_, node)| {
node.signals.iter().map(move |entry| (w, seq, node, entry))
})
})
})
}
pub fn iter<'a>(&'a self) -> impl Iterator<Item = AppTreeItem<'a>> {
self.schedules.iter().flat_map(|(_, w)| {
iter::once(AppTreeItem::Worker(w)).chain(w.seqs.iter().flat_map(|(_, seq)| {
iter::once(AppTreeItem::Sequence(seq))
.chain(seq.nodes.iter().map(|(_, node)| AppTreeItem::Node(node)))
}))
})
}
pub fn tree_iter<'a, F1, F2>(
&'a self,
cmp: F1,
recurse: F2,
) -> impl Iterator<Item = (TreeLocation, AppTreeItem<'a>)>
where
F1: Fn(&AppTreeItem<'a>, &AppTreeItem<'a>) -> Ordering + Copy,
F2: Fn(&AppTreeItem<'a>) -> bool + Copy,
{
let mut schedules: Vec<_> = self
.schedules
.iter()
.map(|(_, v)| AppTreeItem::Worker(v))
.collect();
schedules.sort_by(|a, b| cmp(a, b));
let sched_count = schedules.len();
let tree_row_idx = Rc::new(RefCell::new(0));
schedules
.into_iter()
.enumerate()
.flat_map(move |(group_idx, sched)| {
let AppTreeItem::Worker(sched) = sched else {
unreachable!()
};
*sched.tree_row_idx.borrow_mut() = Some(*tree_row_idx.borrow());
*tree_row_idx.borrow_mut() += 1;
let head = (
TreeLocation {
level: TreeLocationLevel::Node(0),
pos: TreeLocationPos::from_index_count(group_idx, sched_count),
is_expanded: recurse(&AppTreeItem::Worker(sched)),
},
AppTreeItem::Worker(sched),
);
let mut seqs = Vec::new();
if head.0.is_expanded {
seqs.extend(sched.seqs.iter().map(|(_, v)| AppTreeItem::Sequence(v)));
}
seqs.sort_by(|a, b| cmp(a, b));
let seq_count = seqs.len();
iter::once(head).chain(seqs.into_iter().enumerate().flat_map({
let tree_row_idx = tree_row_idx.clone();
move |(group_idx, seq)| {
let AppTreeItem::Sequence(seq) = seq else {
unreachable!()
};
*seq.tree_row_idx.borrow_mut() = Some(*tree_row_idx.borrow());
*tree_row_idx.borrow_mut() += 1;
let head = (
TreeLocation {
level: TreeLocationLevel::Node(1),
pos: TreeLocationPos::from_index_count(group_idx, seq_count),
is_expanded: recurse(&AppTreeItem::Sequence(seq)),
},
AppTreeItem::Sequence(seq),
);
let mut nodes = Vec::new();
if head.0.is_expanded {
nodes.extend(seq.nodes.iter().map(|(_, v)| AppTreeItem::Node(v)));
}
nodes.sort_by(|a, b| cmp(a, b));
let node_count = nodes.len();
iter::once(head).chain(nodes.into_iter().enumerate().map({
let tree_row_idx = tree_row_idx.clone();
move |(group_idx, node)| {
let AppTreeItem::Node(node) = node else {
unreachable!()
};
*node.tree_row_idx.borrow_mut() = Some(*tree_row_idx.borrow());
*tree_row_idx.borrow_mut() += 1;
(
TreeLocation {
level: TreeLocationLevel::Leaf,
pos: TreeLocationPos::from_index_count(
group_idx, node_count,
),
is_expanded: true,
},
AppTreeItem::Node(node),
)
}
}))
}
}))
})
}
pub fn tree_entry_select(&mut self, selection: Option<usize>) {
for (_, sched) in self.schedules.iter_mut() {
sched.is_selected = *sched.tree_row_idx.borrow() == selection;
for (_, seq) in sched.seqs.iter_mut() {
seq.is_selected = *seq.tree_row_idx.borrow() == selection;
for (_, node) in seq.nodes.iter_mut() {
node.is_selected = *node.tree_row_idx.borrow() == selection;
}
}
}
}
pub fn on_toggle_expand_selected(&mut self) {
for (_, sched) in self.schedules.iter_mut() {
if sched.is_selected {
sched.is_expanded = !sched.is_expanded;
}
for (_, seq) in sched.seqs.iter_mut() {
if seq.is_selected {
seq.is_expanded = !seq.is_expanded;
}
}
}
}
fn retain_nodes<F>(workers: &mut HashMap<String, NodoAppTreeSchedule>, f: F)
where
F: Fn(u32) -> bool,
{
for (_, wrk) in workers.iter_mut() {
for (_, seq) in wrk.seqs.iter_mut() {
seq.nodes.retain(|_, n| f(n.info.id))
}
}
}
}
pub enum AppTreeItem<'a> {
Worker(&'a NodoAppTreeSchedule),
Sequence(&'a NodoAppTreeSequence),
Node(&'a NodoAppTreeNode),
}
pub struct TreeLocation {
pub level: TreeLocationLevel,
pub pos: TreeLocationPos,
pub is_expanded: bool,
}
pub enum TreeLocationLevel {
Node(usize),
Leaf,
}
pub enum TreeLocationPos {
First,
Mid,
Last,
}
impl TreeLocationPos {
pub fn from_index_count(i: usize, n: usize) -> Self {
if i == 0 {
TreeLocationPos::First
} else if i + 1 == n {
TreeLocationPos::Last
} else {
TreeLocationPos::Mid
}
}
}