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 mm: bool,
34}
35
36impl DisplayParams {
37    pub fn filter(
38        &self,
39        model: &dyn Model,
40        scope: &[(usize, String)],
41        node_id: usize,
42    ) -> TractResult<bool> {
43        if let Some(nodes) = self.node_ids.as_ref() {
44            return Ok(nodes.iter().any(|n| {
45                n.len() == scope.len() + 1
46                    && &n[0..scope.len()] == scope
47                    && n.last().unwrap().0 == node_id
48            }));
49        }
50        if let Some(node_name) = self.node_name.as_ref() {
51            return Ok(model.node_name(node_id).starts_with(node_name));
52        }
53        if let Some(op_name) = self.op_name.as_ref() {
54            return Ok(model.node_op_name(node_id).starts_with(op_name));
55        }
56        /*
57        if let Some(successor) = self.successors {
58        return Ok(model.node_inputs(node_id).iter().any(|i| i.node == successor));
59        }
60        */
61        Ok(!model.node_const(node_id) || self.konst)
62    }
63
64    pub fn should_draw(&self) -> bool {
65        !self.natural_order
66    }
67
68    pub fn order(&self, model: &dyn Model) -> TractResult<Vec<usize>> {
69        if self.natural_order {
70            Ok((0..model.nodes_len()).collect())
71        } else if self.opt_ram_order {
72            model.eval_order_opt_ram()
73        } else {
74            model.eval_order()
75        }
76    }
77}