diff --git a/crates/bevy_render/src/render_graph/edge.rs b/crates/bevy_render/src/render_graph/edge.rs index bfb05eb04bf3c..0b3fe5432e079 100644 --- a/crates/bevy_render/src/render_graph/edge.rs +++ b/crates/bevy_render/src/render_graph/edge.rs @@ -48,3 +48,9 @@ impl Edge { } } } + +#[derive(PartialEq, Eq)] +pub enum EdgeExistence { + Exists, + DoesNotExist, +} diff --git a/crates/bevy_render/src/render_graph/graph.rs b/crates/bevy_render/src/render_graph/graph.rs index 54b6b106d847d..145e6cd917f88 100644 --- a/crates/bevy_render/src/render_graph/graph.rs +++ b/crates/bevy_render/src/render_graph/graph.rs @@ -9,6 +9,8 @@ use bevy_ecs::prelude::World; use bevy_utils::HashMap; use std::{borrow::Cow, fmt::Debug}; +use super::EdgeExistence; + /// The render graph configures the modular, parallel and re-usable render logic. /// It is a retained and stateless (nodes itself my have their internal state) structure, /// which can not be modified while it is executed by the graph runner. @@ -99,6 +101,69 @@ impl RenderGraph { id } + /// Removes the `node` with the `name` from the graph. + /// If the name is does not exist, nothing happens. + pub fn remove_node( + &mut self, + name: impl Into>, + ) -> Result<(), RenderGraphError> { + let name = name.into(); + if let Some(id) = self.node_names.remove(&name) { + if let Some(node_state) = self.nodes.remove(&id) { + // Remove all edges from other nodes to this one. Note that as we're removing this + // node, we don't need to remove its input edges + for input_edge in node_state.edges.input_edges().iter() { + match input_edge { + Edge::SlotEdge { + output_node, + output_index: _, + input_node: _, + input_index: _, + } => { + if let Ok(output_node) = self.get_node_state_mut(*output_node) { + output_node.edges.remove_output_edge(input_edge.clone())?; + } + } + Edge::NodeEdge { + input_node: _, + output_node, + } => { + if let Ok(output_node) = self.get_node_state_mut(*output_node) { + output_node.edges.remove_output_edge(input_edge.clone())?; + } + } + } + } + // Remove all edges from this node to other nodes. Note that as we're removing this + // node, we don't need to remove its output edges + for output_edge in node_state.edges.output_edges().iter() { + match output_edge { + Edge::SlotEdge { + output_node: _, + output_index: _, + input_node, + input_index: _, + } => { + if let Ok(input_node) = self.get_node_state_mut(*input_node) { + input_node.edges.remove_input_edge(output_edge.clone())?; + } + } + Edge::NodeEdge { + output_node: _, + input_node, + } => { + if let Ok(input_node) = self.get_node_state_mut(*input_node) { + input_node.edges.remove_input_edge(output_edge.clone())?; + } + } + } + } + } + } + + Ok(()) + } + /// Retrieves the [`NodeState`] referenced by the `label`. pub fn get_node_state( &self, @@ -187,7 +252,7 @@ impl RenderGraph { input_index, }; - self.validate_edge(&edge)?; + self.validate_edge(&edge, EdgeExistence::DoesNotExist)?; { let output_node = self.get_node_state_mut(output_node_id)?; @@ -199,6 +264,50 @@ impl RenderGraph { Ok(()) } + /// Removes the [`Edge::SlotEdge`] from the graph. If any nodes or slots do not exist then + /// nothing happens. + pub fn remove_slot_edge( + &mut self, + output_node: impl Into, + output_slot: impl Into, + input_node: impl Into, + input_slot: impl Into, + ) -> Result<(), RenderGraphError> { + let output_slot = output_slot.into(); + let input_slot = input_slot.into(); + let output_node_id = self.get_node_id(output_node)?; + let input_node_id = self.get_node_id(input_node)?; + + let output_index = self + .get_node_state(output_node_id)? + .output_slots + .get_slot_index(output_slot.clone()) + .ok_or(RenderGraphError::InvalidOutputNodeSlot(output_slot))?; + let input_index = self + .get_node_state(input_node_id)? + .input_slots + .get_slot_index(input_slot.clone()) + .ok_or(RenderGraphError::InvalidInputNodeSlot(input_slot))?; + + let edge = Edge::SlotEdge { + output_node: output_node_id, + output_index, + input_node: input_node_id, + input_index, + }; + + self.validate_edge(&edge, EdgeExistence::Exists)?; + + { + let output_node = self.get_node_state_mut(output_node_id)?; + output_node.edges.remove_output_edge(edge.clone())?; + } + let input_node = self.get_node_state_mut(input_node_id)?; + input_node.edges.remove_input_edge(edge)?; + + Ok(()) + } + /// Adds the [`Edge::NodeEdge`] to the graph. This guarantees that the `output_node` /// is run before the `input_node`. pub fn add_node_edge( @@ -214,7 +323,7 @@ impl RenderGraph { input_node: input_node_id, }; - self.validate_edge(&edge)?; + self.validate_edge(&edge, EdgeExistence::DoesNotExist)?; { let output_node = self.get_node_state_mut(output_node_id)?; @@ -226,10 +335,43 @@ impl RenderGraph { Ok(()) } - /// Verifies that the edge is not already existing and + /// Removes the [`Edge::NodeEdge`] from the graph. If either node does not exist then nothing + /// happens. + pub fn remove_node_edge( + &mut self, + output_node: impl Into, + input_node: impl Into, + ) -> Result<(), RenderGraphError> { + let output_node_id = self.get_node_id(output_node)?; + let input_node_id = self.get_node_id(input_node)?; + + let edge = Edge::NodeEdge { + output_node: output_node_id, + input_node: input_node_id, + }; + + self.validate_edge(&edge, EdgeExistence::Exists)?; + + { + let output_node = self.get_node_state_mut(output_node_id)?; + output_node.edges.remove_output_edge(edge.clone())?; + } + let input_node = self.get_node_state_mut(input_node_id)?; + input_node.edges.remove_input_edge(edge)?; + + Ok(()) + } + + /// Verifies that the edge existence is as expected and /// checks that slot edges are connected correctly. - pub fn validate_edge(&mut self, edge: &Edge) -> Result<(), RenderGraphError> { - if self.has_edge(edge) { + pub fn validate_edge( + &mut self, + edge: &Edge, + should_exist: EdgeExistence, + ) -> Result<(), RenderGraphError> { + if should_exist == EdgeExistence::Exists && !self.has_edge(edge) { + return Err(RenderGraphError::EdgeDoesNotExist(edge.clone())); + } else if should_exist == EdgeExistence::DoesNotExist && self.has_edge(edge) { return Err(RenderGraphError::EdgeAlreadyExists(edge.clone())); } @@ -256,7 +398,7 @@ impl RenderGraph { if let Some(Edge::SlotEdge { output_node: current_output_node, .. - }) = input_node_state.edges.input_edges.iter().find(|e| { + }) = input_node_state.edges.input_edges().iter().find(|e| { if let Edge::SlotEdge { input_index: current_input_index, .. @@ -267,11 +409,13 @@ impl RenderGraph { false } }) { - return Err(RenderGraphError::NodeInputSlotAlreadyOccupied { - node: input_node, - input_slot: input_index, - occupied_by_node: *current_output_node, - }); + if should_exist == EdgeExistence::DoesNotExist { + return Err(RenderGraphError::NodeInputSlotAlreadyOccupied { + node: input_node, + input_slot: input_index, + occupied_by_node: *current_output_node, + }); + } } if output_slot.slot_type != input_slot.slot_type { @@ -294,9 +438,9 @@ impl RenderGraph { let output_node_state = self.get_node_state(edge.get_output_node()); let input_node_state = self.get_node_state(edge.get_input_node()); if let Ok(output_node_state) = output_node_state { - if output_node_state.edges.output_edges.contains(edge) { + if output_node_state.edges.output_edges().contains(edge) { if let Ok(input_node_state) = input_node_state { - if input_node_state.edges.input_edges.contains(edge) { + if input_node_state.edges.input_edges().contains(edge) { return true; } } @@ -339,7 +483,7 @@ impl RenderGraph { let node = self.get_node_state(label)?; Ok(node .edges - .input_edges + .input_edges() .iter() .map(|edge| (edge, edge.get_output_node())) .map(move |(edge, output_node_id)| { @@ -356,7 +500,7 @@ impl RenderGraph { let node = self.get_node_state(label)?; Ok(node .edges - .output_edges + .output_edges() .iter() .map(|edge| (edge, edge.get_input_node())) .map(move |(edge, input_node_id)| (edge, self.get_node_state(input_node_id).unwrap()))) @@ -368,6 +512,12 @@ impl RenderGraph { self.sub_graphs.insert(name.into(), sub_graph); } + /// Removes the `sub_graph` with the `name` from the graph. + /// If the name does not exist then nothing happens. + pub fn remove_sub_graph(&mut self, name: impl Into>) { + self.sub_graphs.remove(&name.into()); + } + /// Retrieves the sub graph corresponding to the `name`. pub fn get_sub_graph(&self, name: impl AsRef) -> Option<&RenderGraph> { self.sub_graphs.get(name.as_ref()) diff --git a/crates/bevy_render/src/render_graph/mod.rs b/crates/bevy_render/src/render_graph/mod.rs index 0d204a5162dc1..9f5a5b7c712fb 100644 --- a/crates/bevy_render/src/render_graph/mod.rs +++ b/crates/bevy_render/src/render_graph/mod.rs @@ -31,6 +31,8 @@ pub enum RenderGraphError { }, #[error("attempted to add an edge that already exists")] EdgeAlreadyExists(Edge), + #[error("attempted to remove an edge that does not exist")] + EdgeDoesNotExist(Edge), #[error("node has an unconnected input slot")] UnconnectedNodeInputSlot { node: NodeId, input_slot: usize }, #[error("node has an unconnected output slot")] diff --git a/crates/bevy_render/src/render_graph/node.rs b/crates/bevy_render/src/render_graph/node.rs index eaaead9b0377c..5b806ad31b1e5 100644 --- a/crates/bevy_render/src/render_graph/node.rs +++ b/crates/bevy_render/src/render_graph/node.rs @@ -83,12 +83,30 @@ pub enum NodeRunError { /// A collection of input and output [`Edges`](Edge) for a [`Node`]. #[derive(Debug)] pub struct Edges { - pub id: NodeId, - pub input_edges: Vec, - pub output_edges: Vec, + id: NodeId, + input_edges: Vec, + output_edges: Vec, } impl Edges { + /// Returns all "input edges" (edges going "in") for this node . + #[inline] + pub fn input_edges(&self) -> &[Edge] { + &self.input_edges + } + + /// Returns all "output edges" (edges going "out") for this node . + #[inline] + pub fn output_edges(&self) -> &[Edge] { + &self.output_edges + } + + /// Returns this node's id. + #[inline] + pub fn id(&self) -> NodeId { + self.id + } + /// Adds an edge to the `input_edges` if it does not already exist. pub(crate) fn add_input_edge(&mut self, edge: Edge) -> Result<(), RenderGraphError> { if self.has_input_edge(&edge) { @@ -98,6 +116,21 @@ impl Edges { Ok(()) } + /// Removes an edge from the `input_edges` if it exists. + pub(crate) fn remove_input_edge(&mut self, edge: Edge) -> Result<(), RenderGraphError> { + if let Some((index, _)) = self + .input_edges + .iter() + .enumerate() + .find(|(_i, e)| **e == edge) + { + self.input_edges.swap_remove(index); + Ok(()) + } else { + Err(RenderGraphError::EdgeDoesNotExist(edge)) + } + } + /// Adds an edge to the `output_edges` if it does not already exist. pub(crate) fn add_output_edge(&mut self, edge: Edge) -> Result<(), RenderGraphError> { if self.has_output_edge(&edge) { @@ -107,6 +140,21 @@ impl Edges { Ok(()) } + /// Removes an edge from the `output_edges` if it exists. + pub(crate) fn remove_output_edge(&mut self, edge: Edge) -> Result<(), RenderGraphError> { + if let Some((index, _)) = self + .output_edges + .iter() + .enumerate() + .find(|(_i, e)| **e == edge) + { + self.output_edges.swap_remove(index); + Ok(()) + } else { + Err(RenderGraphError::EdgeDoesNotExist(edge)) + } + } + /// Checks whether the input edge already exists. pub fn has_input_edge(&self, edge: &Edge) -> bool { self.input_edges.contains(edge)