1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use crate::internal::*;
use crate::model::*;
use std::ops::{Deref, DerefMut};

/// A change to apply to a model.
///
/// Actually structured around a model that represent the new nodes to be
/// inserted, plus information about how to connect these new nodes to the
/// pre-existing graph.
#[derive(Clone, Debug)]
pub struct ModelPatch<TI: TensorInfo> {
    /// the model-like 'patch' of nodes to add to the model
    pub model: Model<TI>,
    incoming: HashMap<OutletId, OutletId>,
    shunt_outlet_by: HashMap<OutletId, OutletId>,
}

impl<TI: TensorInfo> Default for ModelPatch<TI> {
    fn default() -> ModelPatch<TI> {
        ModelPatch {
            model: Model::default(),
            incoming: HashMap::new(),
            shunt_outlet_by: HashMap::new(),
        }
    }
}

impl<TI: TensorInfo> Deref for ModelPatch<TI> {
    type Target = Model<TI>;
    fn deref(&self) -> &Model<TI> {
        &self.model
    }
}

impl<TI: TensorInfo> DerefMut for ModelPatch<TI> {
    fn deref_mut(&mut self) -> &mut Model<TI> {
        &mut self.model
    }
}

impl<TI: TensorInfo> ModelPatch<TI> {
    /// Draw a tap from a preexisting node.
    ///
    /// returns an OutletId usable in the little "patch" model
    pub fn tap_model(&mut self, model: &Model<TI>, outlet: OutletId) -> TractResult<OutletId> {
        let fact = model.outlet_fact(outlet)?;
        let node_id = self
            .add_source(format!("incoming-{}/{}", outlet.node, outlet.slot), objekt::clone(fact))?;
        let inside = OutletId::new(node_id, 0);
        self.incoming.insert(inside, outlet);
        Ok(inside)
    }

    /// Replace an Outlet in the target model by one from the patch.
    pub fn shunt_outside(&mut self, outlet: OutletId, by: OutletId) -> TractResult<()> {
        self.shunt_outlet_by.insert(outlet, by);
        Ok(())
    }

    /// Convenience method creating a patch that replace a single operation.
    pub fn replace_single_op<O: Into<Box<Op>>>(
        patched_model: &Model<TI>,
        node: &Node<TI>,
        inputs: &[OutletId],
        new_op: O,
    ) -> TractResult<ModelPatch<TI>> {
        let mut patch = ModelPatch::default();
        let new_op = new_op.into();
        let outputs = node.outputs.iter().map(|o| objekt::clone(&o.fact)).collect();
        let by = patch.add_node(&*node.name, new_op, outputs)?;
        for (ix, i) in inputs.iter().enumerate() {
            let o = patch.tap_model(&patched_model, *i)?;
            patch.add_edge(o, InletId::new(by, ix))?;
        }
        for ix in 0..node.outputs.len() {
            patch.shunt_outside(OutletId::new(node.id, ix), OutletId::new(by, ix))?;
        }
        Ok(patch)
    }

    /// Convenience method creating a patch that replace a single unary operation.
    pub fn single_unary_op<O: Into<Box<Op>>>(
        patched_model: &Model<TI>,
        node: &Node<TI>,
        new_op: O,
    ) -> TractResult<ModelPatch<TI>> {
        Self::replace_single_op(patched_model, node, &[node.inputs[0]], new_op)
    }

    /// Apply all changes in the patch to the target model.
    pub fn apply(self, target: &mut Model<TI>) -> TractResult<()> {
        let ModelPatch { model: patch, incoming: mut mapping, shunt_outlet_by } = self;
        for node in patch.nodes {
            if node.op_is::<crate::ops::source::Source>() {
                continue;
            }
            let Node { id, name, inputs, op, outputs } = node;
            let n_outputs = outputs.len();
            let facts = outputs.into_iter().map(|of| of.fact).collect();
            let added_node_id = target.add_node_disable_output_guess(name, op, facts, true)?;
            for (ix, input) in inputs.into_iter().enumerate() {
                target.add_edge(mapping[&input], InletId::new(added_node_id, ix))?;
            }
            for ix in 0..n_outputs {
                mapping.insert(OutletId::new(id, ix), OutletId::new(added_node_id, ix));
            }
        }
        for (outlet, by) in shunt_outlet_by {
            let fixed_by = mapping[&by];
            let succs = target.nodes()[outlet.node].outputs[outlet.slot].successors.clone();
            for succ in succs {
                target.add_edge(fixed_by, succ)?;
            }
            for o in target.outputs.iter_mut() {
                if *o == outlet {
                    *o = fixed_by;
                }
            }
        }
        Ok(())
    }
}