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

tract_core/model/
translator.rs
