tract_libcli/
display_params.rs

1use crate::model::Model;
2use tract_core::prelude::*;
3
4#[derive(Debug, Clone, PartialEq, Eq, Default)]
5pub enum Io {
6    None,
7    #[default]
8    Short,
9    Long,
10}
11
12#[derive(Debug, Clone, Default, PartialEq, Eq)]
13pub struct DisplayParams {
14    pub konst: bool,
15    pub invariants: bool,
16    pub quiet: bool,
17    pub natural_order: bool,
18    pub opt_ram_order: bool,
19    pub debug_op: bool,
20    pub cost: bool,
21    pub tmp_mem_usage: bool,
22    pub profile: bool,
23    pub folded: bool,
24    pub node_ids: Option<Vec<TVec<(usize, String)>>>,
25    pub op_name: Option<String>,
26    pub node_name: Option<String>,
27    pub expect_core: bool,
28    pub outlet_labels: bool,
29    pub io: Io,
30    pub json: bool,
31    pub info: bool,
32    pub left_column_width: usize,
33    pub has_accelerator: bool,
34    pub mm: bool,
35}
36
37impl DisplayParams {
38    pub fn filter(
39        &self,
40        model: &dyn Model,
41        scope: &[(usize, String)],
42        node_id: usize,
43    ) -> TractResult<bool> {
44        if let Some(nodes) = self.node_ids.as_ref() {
45            return Ok(nodes.iter().any(|n| {
46                n.len() == scope.len() + 1
47                    && &n[0..scope.len()] == scope
48                    && n.last().unwrap().0 == node_id
49            }));
50        }
51        if let Some(node_name) = self.node_name.as_ref() {
52            return Ok(model.node_name(node_id).starts_with(node_name));
53        }
54        if let Some(op_name) = self.op_name.as_ref() {
55            return Ok(model.node_op_name(node_id).starts_with(op_name));
56        }
57        /*
58        if let Some(successor) = self.successors {
59        return Ok(model.node_inputs(node_id).iter().any(|i| i.node == successor));
60        }
61        */
62        Ok(!model.node_const(node_id) || self.konst)
63    }
64
65    pub fn should_draw(&self) -> bool {
66        !self.natural_order
67    }
68
69    pub fn order(&self, model: &dyn Model) -> TractResult<Vec<usize>> {
70        if self.natural_order {
71            Ok((0..model.nodes_len()).collect())
72        } else if self.opt_ram_order {
73            model.eval_order_opt_ram()
74        } else {
75            model.eval_order()
76        }
77    }
78}