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

tract_core/model/
translator.rs
