1use std::collections::{HashMap, HashSet};
2
3use gpui::{Bounds, Pixels, Point, Size, px};
4use serde::{Deserialize, Serialize};
5
6use crate::edge::{Edge, EdgeId};
7use crate::node::{Node, NodeId, Port, PortId};
8use crate::{EdgeBuilder, NodeBuilder, PortKind, PortPosition};
9
10mod store;
11
12pub use store::{ChangeSource, GraphChange, GraphChangeKind, GraphOp};
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct Graph {
16 pub(crate) nodes: HashMap<NodeId, Node>,
17 node_order: Vec<NodeId>,
18 pub ports: HashMap<PortId, Port>,
19 pub edges: HashMap<EdgeId, Edge>,
20
21 pub selected_edge: HashSet<EdgeId>,
22 pub selected_node: HashSet<NodeId>,
23}
24
25impl Graph {
26 pub fn new() -> Self {
27 Self {
28 nodes: HashMap::new(),
29 node_order: vec![],
30 ports: HashMap::new(),
31 edges: HashMap::new(),
32 selected_edge: HashSet::new(),
33 selected_node: HashSet::new(),
34 }
35 }
36
37 pub fn from_json(json: &str) -> serde_json::Result<Self> {
38 serde_json::from_str(json)
39 }
40
41 pub fn to_json(&self) -> serde_json::Result<String> {
42 serde_json::to_string(self)
43 }
44
45 pub fn is_empty(&self) -> bool {
46 self.nodes.is_empty()
47 && self.ports.is_empty()
48 && self.edges.is_empty()
49 && self.node_order.is_empty()
50 && self.selected_edge.is_empty()
51 && self.selected_node.is_empty()
52 }
53
54 pub fn apply(&mut self, op: GraphChangeKind) {
55 match op {
56 GraphChangeKind::NodeAdded(node) => self.add_node(node),
57 GraphChangeKind::NodeRemoved { id } => self.remove_node(&id),
58 GraphChangeKind::NodeMoved { id, x, y } => {
59 if let Some(node) = self.nodes.get_mut(&id) {
60 node.x = px(x);
61 node.y = px(y);
62 }
63 }
64 GraphChangeKind::NodeSetWidthed { id, width } => {
65 if let Some(node) = self.nodes.get_mut(&id) {
66 node.size.width = px(width);
67 }
68 }
69 GraphChangeKind::NodeSetHeighted { id, height } => {
70 if let Some(node) = self.nodes.get_mut(&id) {
71 node.size.height = px(height);
72 }
73 }
74 GraphChangeKind::NodeDataUpdated { id, data } => {
75 if let Some(node) = self.nodes.get_mut(&id) {
76 node.data = data;
77 }
78 }
79 GraphChangeKind::NodeOrderUpdate(vec) => {
80 self.node_order = vec;
81 }
82 GraphChangeKind::PortAdded(port) => self.add_port(port),
83 GraphChangeKind::PortRemoved { id } => {
84 self.remove_port(&id);
85 }
86 GraphChangeKind::EdgeAdded(edge) => self.add_edge(edge),
87 GraphChangeKind::EdgeRemoved { id } => self.remove_edge(id),
88 GraphChangeKind::RedrawRequested => {}
89 GraphChangeKind::Batch(graph_change_kinds) => {
90 for change in graph_change_kinds {
91 self.apply(change);
92 }
93 }
94 }
95 }
96
97 pub fn create_node(&self, node_type: &str) -> NodeBuilder {
98 NodeBuilder::new(node_type)
99 }
100
101 pub fn create_dege(&self) -> EdgeBuilder {
102 EdgeBuilder::new()
103 }
104
105 pub fn next_node_id(&self) -> NodeId {
106 NodeId::new()
107 }
108
109 pub fn next_port_id(&self) -> PortId {
110 PortId::new()
111 }
112
113 pub fn next_edge_id(&self) -> EdgeId {
114 EdgeId::new()
115 }
116
117 pub fn add_node(&mut self, node: Node) {
118 let node_id = node.id;
119 self.nodes.insert(node.id, node);
120 self.node_order.push(node_id);
121 }
122
123 pub fn add_port(&mut self, port: Port) {
124 let ref mut map = self.ports;
125 map.insert(port.id, port);
126 }
127
128 pub fn remove_port(&mut self, id: &PortId) {
129 self.ports.remove(id);
130 }
131
132 pub fn nodes(&self) -> &HashMap<NodeId, Node> {
133 &self.nodes
134 }
135
136 pub fn node_order(&self) -> &Vec<NodeId> {
137 &self.node_order
138 }
139 pub fn node_order_mut(&mut self) -> &mut Vec<NodeId> {
140 &mut self.node_order
141 }
142
143 pub fn new_edge(&self) -> Edge {
144 Edge::new()
145 }
146
147 pub fn add_edge(&mut self, edge: Edge) {
148 self.edges.insert(edge.id, edge);
149 }
150
151 pub fn remove_edge(&mut self, edge_id: EdgeId) {
152 self.edges.remove(&edge_id);
153 }
154
155 pub fn get_node(&self, id: &NodeId) -> Option<&Node> {
156 self.nodes.get(id)
157 }
158
159 pub fn get_node_mut(&mut self, id: &NodeId) -> Option<&mut Node> {
160 self.nodes.get_mut(id)
161 }
162
163 pub fn remove_node(&mut self, id: &NodeId) {
164 let Some(node) = &self.nodes.get(id) else {
165 return;
166 };
167 let mut port_ids = node.inputs.clone();
168 port_ids.extend(node.outputs.clone());
169
170 self.nodes.remove(id);
171 let index = self.node_order.iter().position(|v| *v == *id);
172 if let Some(index) = index {
173 self.node_order.remove(index);
174 }
175
176 for port_id in port_ids.iter() {
177 let edge1 = self
178 .edges
179 .iter()
180 .find(|(_, edge)| edge.source_port == *port_id);
181 if let Some((&edge_id, _)) = edge1 {
182 self.edges.remove(&edge_id);
183 }
184
185 let edge2 = self
186 .edges
187 .iter()
188 .find(|(_, edge)| edge.target_port == *port_id);
189 if let Some((&edge_id, _)) = edge2 {
190 self.edges.remove(&edge_id);
191 }
192
193 self.ports.remove(port_id);
194 }
195 }
196
197 pub fn add_selected_node(&mut self, id: NodeId, shift: bool) {
198 if shift {
199 if self.selected_node.contains(&id) {
200 self.selected_node.remove(&id);
201 } else {
202 self.selected_node.insert(id);
203 }
204 } else {
205 self.selected_node.clear();
206 self.selected_node.insert(id);
207 }
208 }
209 pub fn clear_selected_node(&mut self) {
210 self.selected_node.clear();
211 }
212
213 pub fn remove_selected_node(&mut self) -> bool {
214 if self.selected_node.len() == 0 {
215 return false;
216 }
217
218 let mut ids = vec![];
219 for id in self.selected_node.iter() {
220 ids.push(id.clone());
221 }
222 for id in ids.iter() {
223 self.remove_node(&id);
224 }
225 self.selected_node.clear();
226 return true;
227 }
228
229 pub fn add_selected_edge(&mut self, id: EdgeId, shift: bool) {
230 if shift {
231 if self.selected_edge.contains(&id) {
232 self.selected_edge.remove(&id);
233 } else {
234 self.selected_edge.insert(id);
235 }
236 } else {
237 self.selected_edge.clear();
238 self.selected_edge.insert(id);
239 }
240 }
241 pub fn clear_selected_edge(&mut self) {
242 self.selected_edge.clear();
243 }
244
245 pub fn remove_selected_edge(&mut self) -> bool {
246 if self.selected_edge.len() == 0 {
247 return false;
248 }
249
250 let mut ids = vec![];
251 for id in self.selected_edge.iter() {
252 ids.push(id.clone());
253 }
254 for id in ids.iter() {
255 self.edges.remove(id);
256 }
257 self.selected_edge.clear();
258 return true;
259 }
260
261 pub fn selection_bounds(&self) -> Option<Bounds<Pixels>> {
262 let mut min_x = f32::MAX;
263 let mut min_y = f32::MAX;
264 let mut max_x = f32::MIN;
265 let mut max_y = f32::MIN;
266
267 let mut found = false;
268
269 for id in &self.selected_node {
270 let node = &self.nodes.get(id)?;
271
272 min_x = min_x.min(node.x.into());
273 min_y = min_y.min(node.y.into());
274
275 max_x = max_x.max((node.x + node.size.width).into());
276 max_y = max_y.max((node.y + node.size.height).into());
277
278 found = true;
279 }
280
281 if !found {
282 return None;
283 }
284
285 Some(Bounds::new(
286 Point::new(px(min_x), px(min_y)),
287 Size::new(px(max_x - min_x), px(max_y - min_y)),
288 ))
289 }
290
291 pub fn selected_nodes_with_positions(&self) -> HashMap<NodeId, Point<Pixels>> {
292 self.selected_node
293 .iter()
294 .filter_map(|id| {
295 let n = &self.nodes.get(id)?;
296 Some((*id, n.point()))
297 })
298 .collect()
299 }
300
301 pub fn hit_node(&self, mouse: Point<Pixels>) -> Option<NodeId> {
302 self.nodes
303 .iter()
304 .find(|(_, n)| n.bounds().contains(&mouse))
305 .map(|(id, _)| *id)
306 }
307
308 pub fn bring_node_to_front(&mut self, node_id: NodeId) {
309 self.node_order_mut().retain(|id| *id != node_id);
310
311 self.node_order_mut().push(node_id);
312 }
313
314 pub fn ports_on_node_side(
315 &self,
316 node_id: NodeId,
317 kind: PortKind,
318 position: PortPosition,
319 ) -> Vec<&Port> {
320 self.ports
321 .values()
322 .filter(|p| p.node_id == node_id && p.kind == kind && p.position == position)
323 .collect()
324 }
325}