tract_core/model/
translator.rs

1use crate::internal::*;
2use crate::model::{Fact, Graph, OutletId};
3use std::collections::HashMap;
4use std::convert::*;
5use std::fmt;
6
7pub trait Translate<TI1, O1, TI2, O2>: fmt::Debug
8where
9    TI1: Fact + Clone + 'static,
10    TI2: Fact + Clone + 'static,
11    O1: fmt::Display + fmt::Debug + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
12    O2: fmt::Display + fmt::Debug + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
13{
14    fn translate_node(
15        &self,
16        source: &Graph<TI1, O1>,
17        node: &Node<TI1, O1>,
18        target: &mut Graph<TI2, O2>,
19        mapping: &HashMap<OutletId, OutletId>,
20    ) -> TractResult<TVec<OutletId>>;
21
22    fn translate_model(&self, source: &Graph<TI1, O1>) -> TractResult<Graph<TI2, O2>> {
23        Ok(self.translate_model_with_mappings(source)?.0)
24    }
25
26    fn translate_model_with_mappings(
27        &self,
28        source: &Graph<TI1, O1>,
29    ) -> TractResult<(Graph<TI2, O2>, HashMap<OutletId, OutletId>)> {
30        let _proof_session = source.symbols.proof_cache_session();
31        let mut target = Graph { symbols: source.symbols.clone(), ..Graph::default() };
32        let mut mapping = HashMap::new();
33        for old_id in source.eval_order()? {
34            let node = source.node(old_id);
35            let outlets = self
36                .translate_node(source, node, &mut target, &mapping)
37                .with_context(|| format!("Translating node {node} {self:?}"))?;
38            for (ix, outlet) in outlets.into_iter().enumerate() {
39                mapping.insert(OutletId::new(node.id, ix), outlet);
40                if let Some(label) = source.outlet_label(OutletId::new(node.id, ix)) {
41                    target.set_outlet_label(outlet, label.to_string())?;
42                }
43            }
44        }
45        // do not drop inputs, even if they are useless, to maintain interface
46        for i in source.input_outlets()? {
47            if !mapping.contains_key(i) {
48                let node = source.node(i.node);
49                trace!("Translate useless source {node}");
50                let outlets = self
51                    .translate_node(source, node, &mut target, &mapping)
52                    .with_context(|| format!("Translating input {node} {self:?}"))?;
53                mapping.insert(*i, outlets[0]);
54            }
55        }
56        // maintaining order of i/o interface
57        target.inputs = source.input_outlets()?.iter().map(|i| mapping[i]).collect();
58        target.outputs = source.output_outlets()?.iter().map(|o| mapping[o]).collect();
59        target.properties.clone_from(&source.properties);
60        Ok((target, mapping))
61    }
62}
63
64#[derive(Debug)]
65pub struct IntoTranslator;
66impl<TI1, O1, TI2, O2, EO, ETI> Translate<TI1, O1, TI2, O2> for IntoTranslator
67where
68    TractError: From<EO> + From<ETI>,
69    TI1: Fact + Clone + 'static,
70    TI2: Fact + for<'a> TryFrom<&'a TI1, Error = EO> + Clone + 'static,
71    O1: fmt::Display + fmt::Debug + Clone + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
72    O2: fmt::Display
73        + for<'a> TryFrom<&'a O1, Error = ETI>
74        + fmt::Debug
75        + AsRef<dyn Op>
76        + AsMut<dyn Op>
77        + Clone
78        + 'static,
79    Graph<TI2, O2>: SpecialOps<TI2, O2>,
80{
81    fn translate_node(
82        &self,
83        source: &Graph<TI1, O1>,
84        node: &Node<TI1, O1>,
85        target: &mut Graph<TI2, O2>,
86        mapping: &HashMap<OutletId, OutletId>,
87    ) -> TractResult<TVec<OutletId>> {
88        let node_is_input =
89            (0..node.outputs.len()).all(|o| source.inputs.contains(&(node.id, o).into()));
90        if node_is_input {
91            (0..node.outputs.len())
92                .map(|i| {
93                    target.add_source(
94                        if node.outputs.len() > 1 {
95                            format!("{}-{}", node.name, i)
96                        } else {
97                            node.name.to_string()
98                        },
99                        TI2::try_from(&node.outputs[i].fact)?,
100                    )
101                })
102                .collect()
103        } else {
104            let new_op = O2::try_from(&node.op)?;
105            let facts = node
106                .outputs
107                .iter()
108                .map(|of| Ok(TI2::try_from(&of.fact)?))
109                .collect::<TractResult<TVec<_>>>()?;
110            let new_id = target.add_node(node.name.clone(), new_op, facts)?;
111            for (ix, o) in node.inputs.iter().enumerate() {
112                target.add_edge(mapping[o], InletId::new(new_id, ix))?
113            }
114            Ok(node.outputs.iter().enumerate().map(|(ix, _)| OutletId::new(new_id, ix)).collect())
115        }
116    }
117}
tract_core/model/translator.rs

tract_core/model/
translator.rs
