1use std::collections::{HashMap, HashSet};
2
3use gpui::{Bounds, Pixels, Point, Size, px};
4
5use crate::NodeBuilder;
6use crate::edge::{Edge, EdgeId};
7use crate::node::{Node, NodeId, Port, PortId};
8
9#[derive(Debug, Clone)]
10pub struct Graph {
11 pub(crate) nodes: HashMap<NodeId, Node>,
12 node_order: Vec<NodeId>,
13 pub ports: HashMap<PortId, Port>,
14 pub edges: HashMap<EdgeId, Edge>,
15
16 pub selected_edge: HashSet<EdgeId>,
17 pub selected_node: HashSet<NodeId>,
18}
19
20impl Graph {
21 pub fn new() -> Self {
22 Self {
23 nodes: HashMap::new(),
24 node_order: vec![],
25 ports: HashMap::new(),
26 edges: HashMap::new(),
27 selected_edge: HashSet::new(),
28 selected_node: HashSet::new(),
29 }
30 }
31
32 pub fn create_node(&self, node_type: &str) -> NodeBuilder {
33 NodeBuilder::new(node_type)
34 }
35
36 pub fn next_node_id(&self) -> NodeId {
37 let id = self.nodes.len() as u64 + 1;
38 NodeId(id)
39 }
40
41 pub fn next_port_id(&self) -> PortId {
42 let id = self.ports.len() as u64 + 1;
43 PortId(id)
44 }
45
46 pub fn next_edge_id(&self) -> EdgeId {
47 let id = self.edges.len() as u64 + 1;
48 EdgeId(id)
49 }
50
51 pub fn add_node(&mut self, node: Node) {
52 let node_id = node.id.clone();
53 self.nodes.insert(node.id, node);
54 self.node_order.push(node_id);
55 }
56
57 pub fn add_point(&mut self, port: Port) {
58 let ref mut map = self.ports;
59 map.insert(port.id, port);
60 }
61
62 pub fn nodes(&self) -> &HashMap<NodeId, Node> {
63 &self.nodes
64 }
65
66 pub fn node_order(&self) -> &Vec<NodeId> {
67 &self.node_order
68 }
69 pub fn node_order_mut(&mut self) -> &mut Vec<NodeId> {
70 &mut self.node_order
71 }
72
73 pub fn new_edge(&self) -> Edge {
74 let id = self.edges.len() + 1;
75
76 Edge::new(EdgeId(id as u64))
77 }
78
79 pub fn add_edge(&mut self, edge: Edge) {
80 self.edges.insert(edge.id, edge);
81 }
82
83 pub fn remove_edge(&mut self, edge_id: EdgeId) {
84 self.edges.remove(&edge_id);
85 }
86
87 pub fn get_node(&self, id: &NodeId) -> Option<&Node> {
88 self.nodes.get(id)
89 }
90
91 pub fn get_node_mut(&mut self, id: &NodeId) -> Option<&mut Node> {
92 self.nodes.get_mut(id)
93 }
94
95 pub fn remove_node(&mut self, id: &NodeId) {
96 let node = &self.nodes[id];
97 let mut port_ids = node.inputs.clone();
98 port_ids.extend(node.outputs.clone());
99
100 self.nodes.remove(id);
101 let index = self.node_order.iter().position(|v| *v == *id);
102 if let Some(index) = index {
103 self.node_order.remove(index);
104 }
105
106 for port_id in port_ids.iter() {
107 let edge1 = self
108 .edges
109 .iter()
110 .find(|(_, edge)| edge.source_port == *port_id);
111 if let Some((&edge_id, _)) = edge1 {
112 self.edges.remove(&edge_id);
113 }
114
115 let edge2 = self
116 .edges
117 .iter()
118 .find(|(_, edge)| edge.target_port == *port_id);
119 if let Some((&edge_id, _)) = edge2 {
120 self.edges.remove(&edge_id);
121 }
122
123 self.ports.remove(port_id);
124 }
125 }
126
127 pub fn add_selected_node(&mut self, id: NodeId, shift: bool) {
128 if shift {
129 if self.selected_node.contains(&id) {
130 self.selected_node.remove(&id);
131 } else {
132 self.selected_node.insert(id);
133 }
134 } else {
135 self.selected_node.clear();
136 self.selected_node.insert(id);
137 }
138 }
139 pub fn clear_selected_node(&mut self) {
140 self.selected_node.clear();
141 }
142
143 pub fn remove_selected_node(&mut self) -> bool {
144 if self.selected_node.len() == 0 {
145 return false;
146 }
147
148 let mut ids = vec![];
149 for id in self.selected_node.iter() {
150 ids.push(id.clone());
151 }
152 for id in ids.iter() {
153 self.remove_node(&id);
154 }
155 self.selected_node.clear();
156 return true;
157 }
158
159 pub fn add_selected_edge(&mut self, id: EdgeId, shift: bool) {
160 if shift {
161 if self.selected_edge.contains(&id) {
162 self.selected_edge.remove(&id);
163 } else {
164 self.selected_edge.insert(id);
165 }
166 } else {
167 self.selected_edge.clear();
168 self.selected_edge.insert(id);
169 }
170 }
171 pub fn clear_selected_edge(&mut self) {
172 self.selected_edge.clear();
173 }
174
175 pub fn remove_selected_edge(&mut self) -> bool {
176 if self.selected_edge.len() == 0 {
177 return false;
178 }
179
180 let mut ids = vec![];
181 for id in self.selected_edge.iter() {
182 ids.push(id.clone());
183 }
184 for id in ids.iter() {
185 self.edges.remove(id);
186 }
187 self.selected_edge.clear();
188 return true;
189 }
190
191 pub fn selection_bounds(&self) -> Option<Bounds<Pixels>> {
192 let mut min_x = f32::MAX;
193 let mut min_y = f32::MAX;
194 let mut max_x = f32::MIN;
195 let mut max_y = f32::MIN;
196
197 let mut found = false;
198
199 for id in &self.selected_node {
200 let node = &self.nodes[id];
201
202 min_x = min_x.min(node.x.into());
203 min_y = min_y.min(node.y.into());
204
205 max_x = max_x.max((node.x + node.size.width).into());
206 max_y = max_y.max((node.y + node.size.height).into());
207
208 found = true;
209 }
210
211 if !found {
212 return None;
213 }
214
215 Some(Bounds::new(
216 Point::new(px(min_x), px(min_y)),
217 Size::new(px(max_x - min_x), px(max_y - min_y)),
218 ))
219 }
220
221 pub fn selected_nodes_with_positions(&self) -> HashMap<NodeId, Point<Pixels>> {
222 self.selected_node
223 .iter()
224 .map(|id| {
225 let n = &self.nodes[id];
226 (*id, n.point())
227 })
228 .collect()
229 }
230
231 pub fn hit_node(&self, mouse: Point<Pixels>) -> Option<NodeId> {
232 self.nodes
233 .iter()
234 .find(|(_, n)| n.bounds().contains(&mouse))
235 .map(|(id, _)| *id)
236 }
237
238 pub fn bring_node_to_front(&mut self, node_id: NodeId) {
239 self.node_order_mut().retain(|id| *id != node_id);
240
241 self.node_order_mut().push(node_id);
242 }
243}