flow_graph/
schematic.rs

1mod dot;
2pub mod iterators;
3
4use std::collections::HashMap;
5
6use self::iterators::{Connections, SchematicWalker, WalkDirection};
7use crate::connection::Connection;
8use crate::error::Error;
9use crate::node::{Node, NodeKind, NodePort};
10use crate::port::PortReference;
11use crate::{NodeReference, PortDirection};
12
13pub type ConnectionIndex = usize;
14pub type NodeIndex = usize;
15pub type PortIndex = usize;
16
17pub const SCHEMATIC_INPUT: &str = "<input>";
18pub const SCHEMATIC_INPUT_INDEX: NodeIndex = 0;
19pub const SCHEMATIC_OUTPUT: &str = "<output>";
20pub const SCHEMATIC_OUTPUT_INDEX: NodeIndex = 1;
21
22pub const NS_SCHEMATIC: &str = "__schematic__";
23
24#[derive(Debug, Clone)]
25pub struct Schematic<DATA> {
26  name: String,
27  input: NodeIndex,
28  inherent: NodeIndex,
29  output: NodeIndex,
30  nodes: Vec<Node<DATA>>,
31  node_map: HashMap<String, NodeIndex>,
32  connections: Vec<Connection<DATA>>,
33}
34
35impl<DATA> PartialEq for Schematic<DATA> {
36  fn eq(&self, other: &Self) -> bool {
37    self.name == other.name && self.input == other.input && self.output == other.output
38  }
39}
40
41impl<DATA> Schematic<DATA>
42where
43  DATA: Clone,
44{
45  pub fn new<T: Into<String>>(name: T, input_data: DATA, output_data: DATA) -> Self {
46    let nodes = vec![
47      Node::new(SCHEMATIC_INPUT, SCHEMATIC_INPUT_INDEX, NodeKind::input(), input_data),
48      Node::new(
49        SCHEMATIC_OUTPUT,
50        SCHEMATIC_OUTPUT_INDEX,
51        NodeKind::output(),
52        output_data,
53      ),
54    ];
55    let node_indices = HashMap::from([
56      (SCHEMATIC_INPUT.to_owned(), SCHEMATIC_INPUT_INDEX),
57      (SCHEMATIC_OUTPUT.to_owned(), SCHEMATIC_OUTPUT_INDEX),
58    ]);
59
60    Self {
61      name: name.into(),
62      input: 0,
63      output: 1,
64      inherent: 2,
65      nodes,
66      node_map: node_indices,
67      connections: Default::default(),
68    }
69  }
70
71  #[must_use]
72  pub fn name(&self) -> &str {
73    &self.name
74  }
75
76  pub fn input(&self) -> &Node<DATA> {
77    &self.nodes[self.input]
78  }
79
80  pub fn output(&self) -> &Node<DATA> {
81    &self.nodes[self.output]
82  }
83
84  pub fn inherent(&self) -> &Node<DATA> {
85    &self.nodes[self.inherent]
86  }
87
88  pub fn connections(&self) -> &[Connection<DATA>] {
89    &self.connections
90  }
91
92  pub fn get_port(&self, port: &PortReference) -> &NodePort {
93    match port.direction() {
94      PortDirection::In => &self.nodes()[port.node_index()].inputs()[port.port_index()],
95      PortDirection::Out => &self.nodes()[port.node_index()].outputs()[port.port_index()],
96    }
97  }
98
99  #[must_use]
100  pub fn get_ports(&self) -> Vec<PortReference> {
101    self
102      .nodes
103      .iter()
104      .flat_map(|c| {
105        let mut refs = c.input_refs();
106        refs.append(&mut c.output_refs());
107        refs
108      })
109      .collect()
110  }
111
112  #[must_use]
113  pub fn get_port_name(&self, port: &PortReference) -> &str {
114    self.get_port(port).name()
115  }
116
117  pub fn add_input<T: Into<String>>(&mut self, port: T) -> PortReference {
118    let input = self.get_mut(self.input).unwrap();
119    let port = port.into();
120    input.add_input(port.clone());
121    input.add_output(port)
122  }
123
124  pub fn add_output<T: Into<String>>(&mut self, port: T) -> PortReference {
125    let output = self.get_mut(self.output).unwrap();
126    let port = port.into();
127    output.add_output(port.clone());
128    output.add_input(port)
129  }
130
131  pub fn nodes(&self) -> &[Node<DATA>] {
132    &self.nodes
133  }
134
135  #[must_use]
136  pub fn used_nodes(&self) -> Vec<&Node<DATA>> {
137    let mut nodes_connected_to_input: HashMap<String, _> = SchematicWalker::new_from_input(self)
138      .filter_map(|hop| match hop {
139        iterators::SchematicHop::Node(n) => Some((n.name().to_owned(), n.inner())),
140        _ => None,
141      })
142      .collect();
143    let nodes_connected_to_output: HashMap<String, _> = SchematicWalker::new_from_output(self)
144      .filter_map(|hop| match hop {
145        iterators::SchematicHop::Node(n) => Some((n.name().to_owned(), n.inner())),
146        _ => None,
147      })
148      .collect();
149    nodes_connected_to_input.extend(nodes_connected_to_output);
150    nodes_connected_to_input.into_values().collect()
151  }
152
153  #[must_use]
154  pub fn get(&self, index: NodeIndex) -> Option<&Node<DATA>> {
155    self.nodes.get(index)
156  }
157
158  #[must_use]
159  pub fn get_mut(&mut self, index: NodeIndex) -> Option<&mut Node<DATA>> {
160    self.nodes.get_mut(index)
161  }
162
163  #[must_use]
164  pub fn find(&self, name: &str) -> Option<&Node<DATA>> {
165    self.node_map.get(name).map(|index| &self.nodes[*index])
166  }
167
168  pub fn find_mut(&mut self, name: &str) -> Option<&mut Node<DATA>> {
169    self.node_map.get_mut(name).map(|index| &mut self.nodes[*index])
170  }
171
172  #[must_use]
173  pub fn get_port_connections(&self, port: &NodePort) -> Vec<&Connection<DATA>> {
174    let mut connections = Vec::new();
175    for i in port.connections() {
176      connections.push(&self.connections[*i]);
177    }
178    connections
179  }
180
181  pub fn get_connections(&self) -> &[Connection<DATA>] {
182    &self.connections
183  }
184
185  #[must_use]
186  pub fn downstreams_from(&self, node: NodeIndex) -> Option<Vec<Connections<DATA>>> {
187    self.nodes.get(node).map(|node| {
188      let mut list = Vec::new();
189      for port in node.outputs() {
190        list.push(Connections::new(self, port.connections().to_owned()));
191      }
192      list
193    })
194  }
195
196  #[must_use]
197  pub fn downstream_connections<T: AsRef<PortReference>>(&self, port: T) -> Option<Connections<DATA>> {
198    let port = port.as_ref();
199    self
200      .get(port.node_index)
201      .and_then(|node| node.output_connections(port.port_index))
202      .map(|indices| Connections::new(self, indices.clone()))
203  }
204
205  #[must_use]
206  pub fn upstreams_from(&self, node: NodeIndex) -> Option<Vec<Connections<DATA>>> {
207    self.nodes.get(node).map(|node| {
208      let mut list = Vec::new();
209      for port in node.outputs() {
210        list.push(Connections::new(self, port.connections().to_owned()));
211      }
212      list
213    })
214  }
215
216  #[must_use]
217  pub fn upstream_connections<T: AsRef<PortReference>>(&self, port: T) -> Option<Connections<DATA>> {
218    let port = port.as_ref();
219    self
220      .get(port.node_index)
221      .and_then(|node| node.input_connections(port.port_index))
222      .map(|indices| Connections::new(self, indices.clone()))
223  }
224
225  pub fn add_external<T: Into<String>>(&mut self, name: T, reference: NodeReference, data: DATA) -> NodeIndex {
226    let name = name.into();
227    self.add_node(name, NodeKind::External(reference), data)
228  }
229
230  pub fn add_and_get_mut<T: Into<String>>(&mut self, name: T, reference: NodeReference, data: DATA) -> &mut Node<DATA> {
231    let index = self.add_node(name.into(), NodeKind::External(reference), data);
232    self.get_mut(index).unwrap()
233  }
234
235  pub fn add_inherent<T: Into<String>>(&mut self, name: T, reference: NodeReference, data: DATA) -> NodeIndex {
236    let name = name.into();
237    self.add_node(name, NodeKind::Inherent(reference), data)
238  }
239
240  fn add_node(&mut self, name: String, kind: NodeKind, data: DATA) -> NodeIndex {
241    let existing_index = self.node_map.get(&name);
242
243    match existing_index {
244      Some(index) => *index,
245      None => {
246        let index = self.nodes.len();
247        let node = Node::new(&name, index, kind, data);
248        self.nodes.push(node);
249        self.node_map.insert(name, index);
250        index
251      }
252    }
253  }
254
255  pub fn connect(&mut self, from: PortReference, to: PortReference, data: DATA) -> Result<(), Error> {
256    let connection_index = self.connections.len();
257    let downstream_node = &mut self.nodes[to.node_index];
258    downstream_node.connect_input(to.port_index, connection_index)?;
259    let upstream_node = &mut self.nodes[from.node_index];
260    upstream_node.connect_output(from.port_index, connection_index)?;
261    let connection = Connection::new(from, to, connection_index, data);
262    self.connections.push(connection);
263    Ok(())
264  }
265
266  pub fn walker(&self) -> SchematicWalker<DATA> {
267    SchematicWalker::new_from_input(self)
268  }
269
270  pub fn walk_from_output(&self) -> SchematicWalker<DATA> {
271    SchematicWalker::new_from_output(self)
272  }
273
274  pub fn walk_from_port<T: AsRef<PortReference>>(&self, port: T, direction: WalkDirection) -> SchematicWalker<DATA> {
275    let mut port = *port.as_ref();
276    // If we're walking from an output port, start at the complementary input port.
277    // This is to make sure we walk the upstream connections, not the entire tree.
278    // Use walk_from_output() to walk the entire tree upward.
279    if port.node_index == self.output {
280      port.direction = PortDirection::In;
281    }
282    SchematicWalker::from_port(self, port, direction)
283  }
284
285  #[must_use]
286  pub fn render_dot(&self) -> String {
287    dot::render(self)
288  }
289}