Skip to main content

tract_core/model/
graph.rs

1use super::*;
2use crate::internal::*;
3use crate::ops::Op;
4use crate::plan::PlanOptions;
5use crate::prelude::*;
6
7use std::fmt;
8use tract_data::internal::*;
9use tract_itertools::Itertools;
10
11pub trait SpecialOps<F, O> {
12    fn create_dummy(&self) -> O;
13    fn create_source(&self, fact: F) -> O;
14    fn is_source(op: &O) -> bool;
15    fn wire_node(
16        &mut self,
17        name: impl Into<String>,
18        op: impl Into<O>,
19        inputs: &[OutletId],
20    ) -> TractResult<TVec<OutletId>>;
21    fn add_const(
22        &mut self,
23        name: impl Into<String>,
24        v: impl IntoArcTensor,
25    ) -> TractResult<OutletId>;
26}
27
28/// Main model class
29///
30/// Parameterized by a Fact class.
31#[derive(Clone, Debug)]
32pub struct Graph<F, O>
33where
34    F: Fact + Clone + 'static,
35    O: fmt::Debug + fmt::Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
36{
37    /// all nodes in the model
38    pub nodes: Vec<Node<F, O>>,
39    /// model inputs
40    pub inputs: Vec<OutletId>,
41    /// model outputs
42    pub outputs: Vec<OutletId>,
43    /// outlet labels
44    pub outlet_labels: HashMap<OutletId, String>,
45    /// model properties
46    pub properties: HashMap<String, Arc<Tensor>>,
47    /// symbol scope, including table
48    pub symbols: SymbolScope,
49}
50
51impl<F, O> Default for Graph<F, O>
52where
53    F: Fact + Clone + 'static,
54    O: fmt::Debug + fmt::Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
55{
56    fn default() -> Graph<F, O> {
57        Graph {
58            nodes: vec![],
59            inputs: vec![],
60            outputs: vec![],
61            outlet_labels: HashMap::new(),
62            properties: HashMap::new(),
63            symbols: Default::default(),
64        }
65    }
66}
67
68impl<F, O> Graph<F, O>
69where
70    F: Fact + Clone + 'static,
71    O: fmt::Debug + fmt::Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
72    Graph<F, O>: SpecialOps<F, O>,
73{
74    pub fn add_source(&mut self, name: impl Into<String>, fact: F) -> TractResult<OutletId> {
75        let source = self.create_source(fact.clone());
76        let id = self.add_node(name, source, tvec!(fact))?;
77        let id = OutletId::new(id, 0);
78        self.inputs.push(id);
79        Ok(id)
80    }
81}
82
83impl<F, O> Graph<F, O>
84where
85    F: Fact + Clone + 'static,
86    O: fmt::Debug + fmt::Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
87{
88    pub fn add_node(
89        &mut self,
90        name: impl Into<String>,
91        op: impl Into<O>,
92        output_facts: TVec<F>,
93    ) -> TractResult<usize> {
94        let op = op.into();
95        let name = name.into();
96        let id = self.nodes.len();
97        let outputs =
98            output_facts.into_iter().map(|fact| Outlet { fact, successors: tvec!() }).collect();
99        let node = Node { id, name, op, inputs: vec![], outputs };
100        self.nodes.push(node);
101        Ok(id)
102    }
103
104    /// Connect a node outlet to a node inlet.
105    pub fn add_edge(&mut self, outlet: OutletId, inlet: InletId) -> TractResult<()> {
106        if let Some(previous) = self.nodes[inlet.node].inputs.get(inlet.slot).cloned() {
107            self.nodes[previous.node].outputs[previous.slot]
108                .successors
109                .retain(|&mut succ| succ != inlet);
110        }
111        {
112            let prec = &mut self.nodes[outlet.node];
113            prec.outputs[outlet.slot].successors.push(inlet);
114        }
115        let succ = &mut self.nodes[inlet.node];
116        #[allow(clippy::comparison_chain)]
117        if inlet.slot == succ.inputs.len() {
118            succ.inputs.push(outlet);
119        } else if inlet.slot < succ.inputs.len() {
120            succ.inputs[inlet.slot] = outlet;
121        } else {
122            bail!(
123                "Edges must be added in order and consecutive. Trying to connect input {:?} of node {:?} ",
124                inlet.slot,
125                succ
126            )
127        }
128        Ok(())
129    }
130
131    // Inputs
132
133    /// Get model inputs.
134    pub fn input_outlets(&self) -> TractResult<&[OutletId]> {
135        Ok(&self.inputs)
136    }
137
138    /// Change model inputs.
139    pub fn set_input_outlets(&mut self, inputs: &[OutletId]) -> TractResult<()> {
140        self.inputs = inputs.to_vec();
141        Ok(())
142    }
143
144    /// Change model inputs and return `self`.
145    pub fn with_input_outlets(mut self, inputs: &[OutletId]) -> TractResult<Self> {
146        self.set_input_outlets(inputs)?;
147        Ok(self)
148    }
149
150    /// Set model inputs by the node name.
151    pub fn set_input_names(
152        &mut self,
153        inputs: impl IntoIterator<Item = impl AsRef<str>>,
154    ) -> TractResult<()> {
155        let mut ids = vec![];
156        for i in inputs.into_iter() {
157            let node = self.node_by_name(&i)?;
158            for o in 0..node.outputs.len() {
159                ids.push(OutletId::new(node.id, o))
160            }
161        }
162        self.inputs = ids;
163        Ok(())
164    }
165
166    /// Set model inputs by the node name and return `self`.
167    pub fn with_input_names(
168        mut self,
169        inputs: impl IntoIterator<Item = impl AsRef<str>>,
170    ) -> TractResult<Self> {
171        self.set_input_names(inputs)?;
172        Ok(self)
173    }
174
175    /// Get the `ix`-th input tensor type information.
176    pub fn input_fact(&self, ix: usize) -> TractResult<&F> {
177        let input = self.input_outlets()?[ix];
178        self.outlet_fact(input)
179    }
180
181    /// Get the `ix`-th input tensor type information, mutably.
182    pub fn input_fact_mut(&mut self, ix: usize) -> TractResult<&mut F> {
183        let input = self.input_outlets()?[ix];
184        self.outlet_fact_mut(input)
185    }
186
187    /// Set the `ix`-th input tensor type information.
188    pub fn set_input_fact(&mut self, input: usize, fact: F) -> TractResult<()> {
189        let outlet = self.inputs[input];
190        self.set_outlet_fact(outlet, fact)
191    }
192
193    /// Set the `ix`-th input tensor type information and return `self`.
194    pub fn with_input_fact(mut self, input: usize, fact: F) -> TractResult<Self> {
195        self.set_input_fact(input, fact)?;
196        Ok(self)
197    }
198
199    // Outputs
200    /// Get model outputs.
201    pub fn output_outlets(&self) -> TractResult<&[OutletId]> {
202        Ok(&self.outputs)
203    }
204
205    /// Guess outputs from the topology: node or nodes with no successors.
206    pub fn auto_outputs(&mut self) -> TractResult<()> {
207        let outputs = self
208            .nodes
209            .iter()
210            .flat_map(|n| {
211                let id = n.id;
212                n.outputs.iter().enumerate().map(move |(ix, output_fact)| {
213                    (OutletId::new(id, ix), output_fact.successors.len())
214                })
215            })
216            .filter(|(_f, succs)| *succs == 0)
217            .map(|(f, _)| f)
218            .collect();
219        self.outputs = outputs;
220        Ok(())
221    }
222
223    /// Change model outputs.
224    pub fn set_output_outlets(&mut self, outputs: &[OutletId]) -> TractResult<()> {
225        self.outputs = outputs.to_vec();
226        Ok(())
227    }
228
229    /// Change model outputs and return `self`.
230    pub fn with_output_outlets(mut self, outputs: &[OutletId]) -> TractResult<Self> {
231        self.set_output_outlets(outputs)?;
232        Ok(self)
233    }
234
235    /// Set model outputs by node names.
236    pub fn set_output_names(
237        &mut self,
238        outputs: impl IntoIterator<Item = impl AsRef<str>>,
239    ) -> TractResult<()> {
240        let mut labels: HashMap<StaticName, OutletId> =
241            self.outlet_labels.iter().map(|(o, s)| (Cow::Owned((*s).to_string()), *o)).collect();
242        for n in self.nodes() {
243            for ix in 0..n.outputs.len() {
244                labels.insert(Cow::Owned(format!("{}:{}", &n.name, ix)), OutletId::new(n.id, ix));
245            }
246        }
247        let ids: Vec<OutletId> = outputs
248            .into_iter()
249            .map(|s| {
250                let s = s.as_ref();
251                labels
252                    .get(s)
253                    .cloned()
254                    .or_else(|| self.nodes.iter().find(|n| n.name == s).map(|n| n.id.into()))
255                    .ok_or_else(|| format_err!("Node {} not found", s))
256            })
257            .collect::<TractResult<_>>()?;
258        self.outputs = ids;
259        Ok(())
260    }
261
262    /// Set model outputs by node names and return `self`.
263    pub fn with_output_names(
264        mut self,
265        outputs: impl IntoIterator<Item = impl AsRef<str>>,
266    ) -> TractResult<Self> {
267        self.set_output_names(outputs)?;
268        Ok(self)
269    }
270
271    /// Get the `ix`-th input tensor type information.
272    pub fn output_fact(&self, ix: usize) -> TractResult<&F> {
273        let output = self.output_outlets()?[ix];
274        self.outlet_fact(output)
275    }
276
277    /// Get the `ix`-th input tensor type information, mutably.
278    pub fn output_fact_mut(&mut self, ix: usize) -> TractResult<&mut F> {
279        let output = self.output_outlets()?[ix];
280        self.outlet_fact_mut(output)
281    }
282
283    /// Set the `ix`-th output tensor type information.
284    pub fn set_output_fact(&mut self, output: usize, fact: F) -> TractResult<()> {
285        let outlet = self.outputs[output];
286        self.set_outlet_fact(outlet, fact)
287    }
288
289    /// Set the `ix`-th output tensor type information and return `self`.
290    pub fn with_output_fact(mut self, output: usize, fact: F) -> TractResult<Self> {
291        self.set_output_fact(output, fact)?;
292        Ok(self)
293    }
294
295    // nodes and their facts
296
297    /// Iterate over all node names.
298    pub fn node_names(&self) -> impl Iterator<Item = &str> {
299        self.nodes.iter().map(|s| &*s.name)
300    }
301
302    pub fn node_id_by_name(&self, name: &str) -> TractResult<usize> {
303        self.nodes
304            .iter()
305            .find(|n| n.name == name)
306            .map(|n| n.id)
307            .with_context(|| format!("No node found for name: \"{name}\""))
308    }
309
310    /// Find a node by its name.
311    pub fn node_by_name(&self, name: impl AsRef<str>) -> TractResult<&Node<F, O>> {
312        let id: usize = self.node_id_by_name(name.as_ref())?;
313        Ok(&self.nodes[id])
314    }
315
316    /// Borrow mutably a node by its name.
317    pub fn node_by_name_mut(&mut self, name: impl AsRef<str>) -> TractResult<&mut Node<F, O>> {
318        let id: usize = self.node_id_by_name(name.as_ref())?;
319        Ok(&mut self.nodes[id])
320    }
321
322    pub fn rename_node(&mut self, id: usize, name: &str) -> TractResult<()> {
323        self.node_mut(id).name = name.to_string();
324        Ok(())
325    }
326
327    /// Find a node by its id.
328    pub fn node(&self, id: usize) -> &Node<F, O> {
329        &self.nodes[id]
330    }
331
332    /// Find a node by its id.
333    pub fn node_mut(&mut self, id: usize) -> &mut Node<F, O> {
334        &mut self.nodes[id]
335    }
336
337    /// Access the nodes table.
338    pub fn nodes(&self) -> &[Node<F, O>] {
339        &self.nodes
340    }
341
342    /// Access the nodes table.
343    pub fn nodes_mut(&mut self) -> &mut [Node<F, O>] {
344        &mut self.nodes
345    }
346
347    /// Get input and output tensor information for a node.
348    pub fn node_facts(&self, id: usize) -> TractResult<(TVec<&F>, TVec<&F>)> {
349        Ok((self.node_input_facts(id)?, self.node_output_facts(id)?))
350    }
351
352    /// Get input tensor information for a node.
353    pub fn node_input_facts(&self, node_id: usize) -> TractResult<TVec<&F>> {
354        self.nodes[node_id].inputs.iter().map(|o| self.outlet_fact(*o)).collect()
355    }
356
357    /// Get output tensor information for a node.
358    pub fn node_output_facts(&self, node_id: usize) -> TractResult<TVec<&F>> {
359        Ok(self.nodes[node_id].outputs.iter().map(|o| &o.fact).collect())
360    }
361
362    // outlets
363
364    /// Get tensor information for a single outlet.
365    pub fn outlet_fact(&self, outlet: OutletId) -> TractResult<&F> {
366        ensure!(outlet.node < self.nodes.len(), "Invalid outlet for graph");
367        let outlets = &self.nodes[outlet.node].outputs;
368        outlets
369            .get(outlet.slot)
370            .map(|o| &o.fact)
371            .with_context(|| format!("Invalid outlet reference: {outlet:?}"))
372    }
373
374    /// Get tensor information for a single outlet.
375    pub fn outlet_fact_mut(&mut self, outlet: OutletId) -> TractResult<&mut F> {
376        let outlets = &mut self.nodes[outlet.node].outputs;
377        outlets
378            .get_mut(outlet.slot)
379            .map(|o| &mut o.fact)
380            .with_context(|| format!("Invalid outlet reference: {outlet:?}"))
381    }
382
383    /// Get multiple mutable tensor information for outlets.
384    pub fn outlets_fact_mut(&mut self, outlets: &[OutletId]) -> TractResult<TVec<&mut F>> {
385        assert!(outlets.iter().tuple_combinations().all(|(a, b)| a != b));
386        unsafe {
387            outlets
388                .iter()
389                .map(|o| Ok((self.outlet_fact(*o)? as *const F as *mut F).as_mut().unwrap()))
390                .collect()
391        }
392    }
393
394    /// Set tensor information for a single outlet.
395    pub fn set_outlet_fact(&mut self, outlet: OutletId, fact: F) -> TractResult<()> {
396        let outlets = &mut self.nodes[outlet.node].outputs;
397        if outlets.len() <= outlet.slot {
398            bail!("Invalid outlet refererence: {:?}", outlet)
399        }
400        outlets[outlet.slot].fact = fact;
401        Ok(())
402    }
403
404    /// Set tensor information for a single outlet and return `self`.
405    pub fn with_outlet_fact(mut self, outlet: OutletId, fact: F) -> TractResult<Self> {
406        self.set_outlet_fact(outlet, fact)?;
407        Ok(self)
408    }
409
410    // outlet labels
411
412    /// Get label for an outlet.
413    pub fn outlet_label(&self, outlet: OutletId) -> Option<&str> {
414        self.outlet_labels.get(&outlet).map(|s| &**s)
415    }
416
417    /// Set label for an outlet.
418    pub fn set_outlet_label(&mut self, outlet: OutletId, label: String) -> TractResult<()> {
419        self.outlet_labels.insert(outlet, label);
420        Ok(())
421    }
422
423    /// Set label for an outlet and return `self`.
424    pub fn with_outlet_label(mut self, outlet: OutletId, label: String) -> TractResult<Self> {
425        self.set_outlet_label(outlet, label)?;
426        Ok(self)
427    }
428
429    /// Find outlet by label.
430    pub fn find_outlet_label(&self, label: &str) -> Option<OutletId> {
431        self.outlet_labels.iter().find(|(_k, v)| **v == label).map(|(k, _v)| *k)
432    }
433
434    // misc
435
436    /// Computes an evalutation order for the graph inputs and outputs
437    pub fn eval_order(&self) -> TractResult<Vec<usize>> {
438        super::order::eval_order(self)
439    }
440
441    /// Computes an evalutation order for the graph inputs and outputs. This order will minimize
442    /// temporary buffers.
443    pub fn eval_order_opt_ram(&self) -> TractResult<Vec<usize>> {
444        super::order::eval_order_opt_ram(self)
445    }
446
447    #[cfg(not(all(debug_assertions, feature = "paranoid_assertions")))]
448    #[inline]
449    pub fn check_edges(&self) -> TractResult<()> {
450        Ok(())
451    }
452
453    /// Performs a sanity check on network connections.
454    #[cfg(all(debug_assertions, feature = "paranoid_assertions"))]
455    pub fn check_edges(&self) -> TractResult<()> {
456        for node_id in self.eval_order()? {
457            let node = &self.nodes[node_id];
458            for (ix, input) in node.inputs.iter().enumerate() {
459                let prec = &self.nodes[input.node];
460                if !prec.outputs[input.slot].successors.contains(&InletId::new(node.id, ix)) {
461                    bail!(
462                        "Mismatched oncoming edge, node:{} input:{} to {:?} not reciprocated",
463                        node.id,
464                        ix,
465                        prec
466                    )
467                }
468            }
469            for (ix, output) in node.outputs.iter().enumerate() {
470                for succ in &output.successors {
471                    if self.nodes[succ.node].inputs[succ.slot] != OutletId::new(node.id, ix) {
472                        bail!(
473                            "Mismatched outgoing edge, node:{} output:{} to {:?} not reciprocated",
474                            node.id,
475                            ix,
476                            succ
477                        )
478                    }
479                }
480            }
481        }
482        Ok(())
483    }
484
485    /// Evaluate temporary memory usage with its related node at each step of the given order.
486    pub fn eval_tmp_memory_usage<Flushable>(
487        &self,
488        order: &[usize],
489        flushable: Flushable,
490    ) -> TractResult<TVec<(usize, TDim)>>
491    where
492        Flushable: Fn(&Node<F, O>) -> bool,
493    {
494        super::memory::eval_tmp_memory_usage(self, order, flushable)
495    }
496
497    #[cfg(not(all(debug_assertions, feature = "paranoid_assertions")))]
498    #[inline]
499    pub fn check_names(&self) -> TractResult<()> {
500        Ok(())
501    }
502
503    /// Performs a sanity check on network connections.
504    #[cfg(all(debug_assertions, feature = "paranoid_assertions"))]
505    pub fn check_names(&self) -> TractResult<()> {
506        let dups =
507            self.eval_order()?.iter().map(|n| &self.nodes[*n].name).duplicates().collect_vec();
508        ensure!(dups.len() == 0, "Duplicate node name(s) : {:?}\n{}", dups, &self);
509        Ok(())
510    }
511
512    /// Converts the model into a `RunnableModel` to actually process user data.
513    pub fn into_runnable(self) -> TractResult<RunnableModel<F, O, Self>> {
514        crate::plan::SimplePlan::new_with_options(self, &PlanOptions::default())
515    }
516
517    /// Converts the model into a `RunnableModel` to actually process user data. This variant
518    /// accepts options.
519    pub fn into_runnable_with_options(
520        self,
521        options: &PlanOptions,
522    ) -> TractResult<RunnableModel<F, O, Self>> {
523        crate::plan::SimplePlan::new_with_options(self, options)
524    }
525
526    pub fn linear_prec(&self, id: usize) -> TractResult<Option<&Node<F, O>>> {
527        let node = &self.nodes()[id];
528        if node.inputs.len() != 1 {
529            return Ok(None);
530        }
531        let prec = &self.nodes()[node.inputs[0].node];
532        if prec.outputs.iter().map(|of| of.successors.len()).sum::<usize>() != 1 {
533            return Ok(None);
534        }
535        Ok(Some(prec))
536    }
537
538    pub fn single_prec(&self, id: usize) -> TractResult<Option<&Node<F, O>>> {
539        let node = &self.nodes()[id];
540        if node.inputs.len() != 1 {
541            return Ok(None);
542        }
543        let prec = &self.nodes()[node.inputs[0].node];
544        Ok(Some(prec))
545    }
546
547    pub fn all_prec(&self, id: usize) -> TractResult<Option<TVec<&Node<F, O>>>> {
548        let node = &self.nodes()[id];
549        if node.inputs.is_empty() {
550            return Ok(None);
551        };
552        Ok(Some(node.inputs.iter().map(|n| &self.nodes()[n.node]).collect()))
553    }
554
555    pub fn single_prec_at(&self, id: usize, count: usize) -> TractResult<Option<&Node<F, O>>> {
556        let mut node = self.node(id);
557        for _ in 0..count {
558            if let Some(next) = self.linear_prec(node.id)? {
559                node = next
560            } else {
561                return Ok(None);
562            }
563        }
564        Ok(Some(node))
565    }
566
567    pub fn single_succ_at(&self, id: usize, count: usize) -> TractResult<Option<&Node<F, O>>> {
568        let mut node = self.node(id);
569        for _ in 0..count {
570            if let Some(next) = self.linear_succ(node.id)? {
571                node = next
572            } else {
573                return Ok(None);
574            }
575        }
576        Ok(Some(node))
577    }
578
579    /// linear_succ is only intended for optimisation of simple operators
580    /// with 1 output, and only 1 output successors (successor with only 1 input)
581    pub fn linear_succ(&self, id: usize) -> TractResult<Option<&Node<F, O>>> {
582        let node = &self.nodes()[id];
583
584        if node.outputs.len() != 1 || node.outputs[0].successors.len() != 1 {
585            return Ok(None);
586        }
587        let succ = node.outputs[0].successors[0];
588        let succ = &self.nodes()[succ.node];
589        if succ.inputs.len() != 1 {
590            return Ok(None);
591        }
592        Ok(Some(succ))
593    }
594
595    pub fn single_succ(&self, id: usize) -> TractResult<Option<&Node<F, O>>> {
596        let node = &self.nodes()[id];
597
598        if node.outputs.len() != 1 || node.outputs[0].successors.len() != 1 {
599            return Ok(None);
600        }
601        let succ = node.outputs[0].successors[0];
602        Ok(Some(&self.nodes()[succ.node]))
603    }
604
605    pub fn all_succ(&self, id: usize) -> TractResult<Option<TVec<&Node<F, O>>>> {
606        let node = &self.nodes()[id];
607        if node.outputs.is_empty() {
608            return Ok(None);
609        };
610
611        Ok(Some(
612            node.outputs
613                .iter()
614                .flat_map(|o| {
615                    o.successors.iter().map(|succ| &self.nodes()[succ.node]).collect::<Vec<_>>()
616                })
617                .collect(),
618        ))
619    }
620
621    pub fn outlet_successors(&self, outlet: OutletId) -> &[InletId] {
622        &self.nodes[outlet.node].outputs[outlet.slot].successors
623    }
624
625    /// retrieve of create a symbol
626    pub fn sym(&self, s: &str) -> Symbol {
627        self.symbols.sym(s)
628    }
629
630    /// create a new symbol with the prefix
631    pub fn new_sym_with_prefix(&self, prefix: &str) -> Symbol {
632        self.symbols.new_with_prefix(prefix)
633    }
634
635    /// generates a name for a new node in the model that will not conflict (by suffixing with a
636    /// dot and number)
637    pub fn unique_name<'n>(&self, prefix: impl Into<Cow<'n, str>>) -> Cow<'n, str> {
638        let prefix = prefix.into();
639        if self.nodes.iter().all(|n| n.name != *prefix) {
640            return prefix;
641        }
642        for i in 1.. {
643            let s = format!("{prefix}.{i}");
644            if self.nodes.iter().all(|n| n.name != s) {
645                return Cow::Owned(s);
646            }
647        }
648        unreachable!();
649    }
650}
651
652impl<F, O> fmt::Display for Graph<F, O>
653where
654    F: Fact + Clone + 'static,
655    O: fmt::Debug + fmt::Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
656{
657    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
658        for i in 0..self.nodes.len() {
659            let input_1 =
660                self.nodes[i].inputs.first().map(|o| format!("{o:?}")).unwrap_or_default();
661            let input_2 = self.nodes[i].inputs.get(1).map(|o| format!("{o:?}")).unwrap_or_default();
662            let successors = self.nodes[i]
663                .outputs
664                .first()
665                .iter()
666                .flat_map(|o| o.successors.iter())
667                .collect_vec();
668            let output_1 = successors.first().map(|o| format!("{o:?}")).unwrap_or_default();
669            let output_2 = successors.get(1).map(|o| format!("{o:?}")).unwrap_or_default();
670            writeln!(
671                fmt,
672                "{:5} | {:8} {:8} -> {:8} {:8} | {:25} {:50} {} => {}",
673                i,
674                input_1,
675                input_2,
676                output_1,
677                output_2,
678                self.nodes[i].op().name(),
679                self.nodes[i].name,
680                self.node_input_facts(i).unwrap().iter().map(|f| format!("{f:?}")).join(" ; "),
681                self.node_output_facts(i).unwrap().iter().map(|f| format!("{f:?}")).join(" ; "),
682            )?;
683            if self.nodes[i].inputs.len() > 2 {
684                writeln!(
685                    fmt,
686                    "                                               |   * inputs: {}",
687                    self.nodes[i].inputs.iter().map(|s| format!("{s:?}")).join(", ")
688                )?;
689            }
690            if self.nodes[i].outputs.len() > 1
691                || successors.len() > 2
692                || (self.outlet_label(i.into()).is_some()
693                    && self.outlet_label(i.into()).unwrap() != self.nodes[i].name)
694            {
695                for o in 0..self.nodes[i].outputs.len() {
696                    if self.outlet_successors((i, o).into()).len() > 0 {
697                        writeln!(
698                            fmt,
699                            "                                               |   * output #{}: {} {}",
700                            o,
701                            self.outlet_label((i, o).into()).unwrap_or(""),
702                            self.outlet_successors((i, o).into())
703                                .iter()
704                                .map(|s| format!("{s:?}"))
705                                .join(", "),
706                        )?;
707                    }
708                }
709            }
710        }
711        writeln!(fmt, "outputs: {}", self.outputs.iter().map(|o| format!("{o:?}")).join(", "))?;
712        Ok(())
713    }
714}
715
716impl<F, O> Graph<F, O>
717where
718    F: Fact + Clone + 'static + for<'a> std::convert::From<&'a F>,
719    O: std::fmt::Display
720        + std::fmt::Debug
721        + Clone
722        + AsRef<dyn Op>
723        + AsMut<dyn Op>
724        + Clone
725        + 'static
726        + for<'a> std::convert::From<&'a O>,
727    Graph<F, O>: SpecialOps<F, O>,
728{
729    #[cfg(debug_assertions)]
730    pub fn check_compact(&self) -> TractResult<()> {
731        let order = self.eval_order()?;
732        let useless_sources = self
733            .input_outlets()?
734            .iter()
735            .filter(|io| {
736                self.outlet_successors(**io).len() == 0
737                    && !self.output_outlets().unwrap().contains(io)
738            })
739            .count();
740        if order.len() + useless_sources != self.nodes.len() {
741            bail!(
742                "Eval order is {} long, nodes are {}, including {} unused sources",
743                order.len(),
744                self.nodes.len(),
745                useless_sources
746            );
747        }
748        if (0..order.len()).any(|ix| order[ix] != ix) {
749            bail!("eval order is not trivial");
750        }
751        let mut seen = std::collections::HashSet::new();
752        for (ix, n) in self.nodes.iter().enumerate() {
753            if ix != n.id {
754                bail!("Invalid node id: position is {}, node is {}", ix, n);
755            }
756            if seen.contains(&n.name) {
757                bail!("duplicate name for node {n}");
758            }
759            seen.insert(&n.name);
760        }
761        Ok(())
762    }
763
764    pub fn compact(&mut self) -> TractResult<()> {
765        let mut order = self.eval_order()?;
766        if order.len() == self.nodes.len() && order.iter().enumerate().all(|(a, b)| a == *b) {
767            return Ok(());
768        }
769        for i in &self.inputs {
770            if !order.contains(&i.node) {
771                order.push(i.node);
772            }
773        }
774        let mut old_to_new = vec![usize::MAX; self.nodes.len()];
775        let mut new_nodes = vec![
776            Node {
777                id: self.nodes.len(),
778                name: "".to_string(),
779                inputs: vec![],
780                op: self.create_dummy(),
781                outputs: tvec!(),
782            };
783            order.len()
784        ];
785        for (ix, id) in order.iter().enumerate() {
786            old_to_new[*id] = ix;
787            std::mem::swap(&mut new_nodes[ix], &mut self.nodes[*id]);
788        }
789        for node in &mut new_nodes {
790            if self.inputs.iter().any(|n| n.node == node.id) && !Self::is_source(&node.op) {
791                node.inputs.clear();
792                node.op = self.create_source(node.outputs[0].fact.clone());
793            }
794            node.id = old_to_new[node.id];
795            for input in &mut node.inputs {
796                assert!(old_to_new[input.node] < order.len());
797                input.node = old_to_new[input.node];
798            }
799            for output in &mut node.outputs {
800                for succ in &mut output.successors {
801                    succ.node = old_to_new[succ.node];
802                }
803                output.successors.retain(|s| s.node < order.len());
804                output.successors.sort();
805            }
806        }
807        self.nodes = new_nodes;
808        for input in &mut self.inputs {
809            assert!(old_to_new[input.node] < order.len());
810            input.node = old_to_new[input.node];
811        }
812        for output in &mut self.outputs {
813            assert!(old_to_new[output.node] < order.len());
814            output.node = old_to_new[output.node];
815        }
816        self.outlet_labels = std::mem::take(&mut self.outlet_labels)
817            .into_iter()
818            .map(|(k, v)| (OutletId::new(old_to_new[k.node], k.slot), v))
819            .filter(|(k, _)| k.node < order.len())
820            .collect();
821        ensure!(self.nodes.iter().enumerate().all(|(ix, n)| n.id == ix));
822        #[cfg(debug_assertions)]
823        {
824            self.check_compact().context("after graph compaction")?;
825        }
826        Ok(())
827    }
828
829    pub fn into_compact(mut self) -> TractResult<Self> {
830        self.compact()?;
831        Ok(self)
832    }
833}