render graph utils

crates/bevy_core_pipeline/src/bloom/mod.rs

@@ -16,7 +16,7 @@ use bevy_render::{
         ComponentUniforms, DynamicUniformIndex, ExtractComponentPlugin, UniformComponentPlugin,
     },
     prelude::Color,
-    render_graph::{Node, NodeRunError, RenderGraph, RenderGraphContext},
+    render_graph::{add_node, Node, NodeRunError, RenderGraphContext},
     render_resource::*,
     renderer::{RenderContext, RenderDevice},
     texture::{CachedTexture, TextureCache},
@@ -76,42 +76,28 @@ impl Plugin for BloomPlugin {
             );
 
         // Add bloom to the 3d render graph
-        {
-            let bloom_node = BloomNode::new(&mut render_app.world);
-            let mut graph = render_app.world.resource_mut::<RenderGraph>();
-            let draw_3d_graph = graph
-                .get_sub_graph_mut(crate::core_3d::graph::NAME)
-                .unwrap();
-            draw_3d_graph.add_node(core_3d::graph::node::BLOOM, bloom_node);
-            // MAIN_PASS -> BLOOM -> TONEMAPPING
-            draw_3d_graph.add_node_edge(
-                crate::core_3d::graph::node::MAIN_PASS,
-                core_3d::graph::node::BLOOM,
-            );
-            draw_3d_graph.add_node_edge(
+        add_node::<BloomNode>(
+            render_app,
+            core_3d::graph::NAME,
+            core_3d::graph::node::BLOOM,
+            &[
+                core_3d::graph::node::MAIN_PASS,
                 core_3d::graph::node::BLOOM,
-                crate::core_3d::graph::node::TONEMAPPING,
-            );
-        }
+                core_3d::graph::node::TONEMAPPING,
+            ],
+        );
 
         // Add bloom to the 2d render graph
-        {
-            let bloom_node = BloomNode::new(&mut render_app.world);
-            let mut graph = render_app.world.resource_mut::<RenderGraph>();
-            let draw_2d_graph = graph
-                .get_sub_graph_mut(crate::core_2d::graph::NAME)
-                .unwrap();
-            draw_2d_graph.add_node(core_2d::graph::node::BLOOM, bloom_node);
-            // MAIN_PASS -> BLOOM -> TONEMAPPING
-            draw_2d_graph.add_node_edge(
-                crate::core_2d::graph::node::MAIN_PASS,
-                core_2d::graph::node::BLOOM,
-            );
-            draw_2d_graph.add_node_edge(
+        add_node::<BloomNode>(
+            render_app,
+            core_2d::graph::NAME,
+            core_2d::graph::node::BLOOM,
+            &[
+                core_2d::graph::node::MAIN_PASS,
                 core_2d::graph::node::BLOOM,
-                crate::core_2d::graph::node::TONEMAPPING,
-            );
-        }
+                core_2d::graph::node::TONEMAPPING,
+            ],
+        );
     }
 }
 
@@ -128,8 +114,8 @@ pub struct BloomNode {
     )>,
 }
 
-impl BloomNode {
-    pub fn new(world: &mut World) -> Self {
+impl FromWorld for BloomNode {
+    fn from_world(world: &mut World) -> Self {
         Self {
             view_query: QueryState::new(world),
         }

crates/bevy_core_pipeline/src/core_2d/mod.rs

@@ -73,15 +73,14 @@ impl Plugin for Core2dPlugin {
         draw_2d_graph.add_node(graph::node::TONEMAPPING, tonemapping);
         draw_2d_graph.add_node(graph::node::END_MAIN_PASS_POST_PROCESSING, EmptyNode);
         draw_2d_graph.add_node(graph::node::UPSCALING, upscaling);
-        draw_2d_graph.add_node_edge(graph::node::MAIN_PASS, graph::node::TONEMAPPING);
-        draw_2d_graph.add_node_edge(
+
+        draw_2d_graph.add_node_edges(&[
+            graph::node::MAIN_PASS,
             graph::node::TONEMAPPING,
             graph::node::END_MAIN_PASS_POST_PROCESSING,
-        );
-        draw_2d_graph.add_node_edge(
-            graph::node::END_MAIN_PASS_POST_PROCESSING,
             graph::node::UPSCALING,
-        );
+        ]);
+
         graph.add_sub_graph(graph::NAME, draw_2d_graph);
     }
 }

crates/bevy_core_pipeline/src/core_3d/mod.rs

@@ -94,16 +94,13 @@ impl Plugin for Core3dPlugin {
         draw_3d_graph.add_node(graph::node::END_MAIN_PASS_POST_PROCESSING, EmptyNode);
         draw_3d_graph.add_node(graph::node::UPSCALING, upscaling);
 
-        draw_3d_graph.add_node_edge(graph::node::PREPASS, graph::node::MAIN_PASS);
-        draw_3d_graph.add_node_edge(graph::node::MAIN_PASS, graph::node::TONEMAPPING);
-        draw_3d_graph.add_node_edge(
+        draw_3d_graph.add_node_edges(&[
+            graph::node::MAIN_PASS,
             graph::node::TONEMAPPING,
             graph::node::END_MAIN_PASS_POST_PROCESSING,
-        );
-        draw_3d_graph.add_node_edge(
-            graph::node::END_MAIN_PASS_POST_PROCESSING,
             graph::node::UPSCALING,
-        );
+        ]);
+
         graph.add_sub_graph(graph::NAME, draw_3d_graph);
     }
 }

crates/bevy_core_pipeline/src/fxaa/mod.rs

@@ -9,7 +9,7 @@ use bevy_reflect::{
 use bevy_render::{
     extract_component::{ExtractComponent, ExtractComponentPlugin},
     prelude::Camera,
-    render_graph::RenderGraph,
+    render_graph::add_node,
     render_resource::*,
     renderer::RenderDevice,
     texture::BevyDefault,
@@ -92,38 +92,27 @@ impl Plugin for FxaaPlugin {
             .init_resource::<SpecializedRenderPipelines<FxaaPipeline>>()
             .add_systems(Render, prepare_fxaa_pipelines.in_set(RenderSet::Prepare));
 
-        {
-            let fxaa_node = FxaaNode::new(&mut render_app.world);
-            let mut binding = render_app.world.resource_mut::<RenderGraph>();
-            let graph = binding.get_sub_graph_mut(core_3d::graph::NAME).unwrap();
-
-            graph.add_node(core_3d::graph::node::FXAA, fxaa_node);
-
-            graph.add_node_edge(
+        add_node::<FxaaNode>(
+            render_app,
+            core_3d::graph::NAME,
+            core_3d::graph::node::FXAA,
+            &[
                 core_3d::graph::node::TONEMAPPING,
                 core_3d::graph::node::FXAA,
-            );
-            graph.add_node_edge(
-                core_3d::graph::node::FXAA,
                 core_3d::graph::node::END_MAIN_PASS_POST_PROCESSING,
-            );
-        }
-        {
-            let fxaa_node = FxaaNode::new(&mut render_app.world);
-            let mut binding = render_app.world.resource_mut::<RenderGraph>();
-            let graph = binding.get_sub_graph_mut(core_2d::graph::NAME).unwrap();
-
-            graph.add_node(core_2d::graph::node::FXAA, fxaa_node);
+            ],
+        );
 
-            graph.add_node_edge(
+        add_node::<FxaaNode>(
+            render_app,
+            core_2d::graph::NAME,
+            core_2d::graph::node::FXAA,
+            &[
                 core_2d::graph::node::TONEMAPPING,
                 core_2d::graph::node::FXAA,
-            );
-            graph.add_node_edge(
-                core_2d::graph::node::FXAA,
                 core_2d::graph::node::END_MAIN_PASS_POST_PROCESSING,
-            );
-        }
+            ],
+        );
     }
 }
 

crates/bevy_core_pipeline/src/fxaa/node.rs

@@ -27,8 +27,8 @@ pub struct FxaaNode {
     cached_texture_bind_group: Mutex<Option<(TextureViewId, BindGroup)>>,
 }
 
-impl FxaaNode {
-    pub fn new(world: &mut World) -> Self {
+impl FromWorld for FxaaNode {
+    fn from_world(world: &mut World) -> Self {
         Self {
             query: QueryState::new(world),
             cached_texture_bind_group: Mutex::new(None),

crates/bevy_render/src/render_graph/graph.rs

@@ -119,6 +119,14 @@ impl RenderGraph {
         id
     }
 
+    /// Add `node_edge` based on the order of the given `edges` array.
+    pub fn add_node_edges(&mut self, edges: &[&'static str]) {
+        for window in edges.windows(2) {
+            let [a, b] = window else { break; };
+            self.add_node_edge(*a, *b);
+        }
+    }
+
     /// Removes the `node` with the `name` from the graph.
     /// If the name is does not exist, nothing happens.
     pub fn remove_node(
@@ -583,6 +591,36 @@ impl RenderGraph {
     pub fn get_sub_graph_mut(&mut self, name: impl AsRef<str>) -> Option<&mut RenderGraph> {
         self.sub_graphs.get_mut(name.as_ref())
     }
+
+    /// Retrieves the sub graph corresponding to the `name`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any invalid node name is given.
+    ///
+    /// # See also
+    ///
+    /// - [`get_sub_graph`](Self::get_sub_graph) for a fallible version.
+    pub fn sub_graph(&self, name: impl AsRef<str>) -> &RenderGraph {
+        self.sub_graphs
+            .get(name.as_ref())
+            .unwrap_or_else(|| panic!("Node {} not found in sub_graph", name.as_ref()))
+    }
+
+    /// Retrieves the sub graph corresponding to the `name` mutably.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any invalid node name is given.
+    ///
+    /// # See also
+    ///
+    /// - [`get_sub_graph_mut`](Self::get_sub_graph_mut) for a fallible version.
+    pub fn sub_graph_mut(&mut self, name: impl AsRef<str>) -> &mut RenderGraph {
+        self.sub_graphs
+            .get_mut(name.as_ref())
+            .unwrap_or_else(|| panic!("Node {} not found in sub_graph", name.as_ref()))
+    }
 }
 
 impl Debug for RenderGraph {
@@ -635,7 +673,7 @@ mod tests {
         },
         renderer::RenderContext,
     };
-    use bevy_ecs::world::World;
+    use bevy_ecs::world::{FromWorld, World};
     use bevy_utils::HashSet;
 
     #[derive(Debug)]
@@ -676,6 +714,22 @@ mod tests {
         }
     }
 
+    fn input_nodes(name: &'static str, graph: &RenderGraph) -> HashSet<NodeId> {
+        graph
+            .iter_node_inputs(name)
+            .unwrap()
+            .map(|(_edge, node)| node.id)
+            .collect::<HashSet<NodeId>>()
+    }
+
+    fn output_nodes(name: &'static str, graph: &RenderGraph) -> HashSet<NodeId> {
+        graph
+            .iter_node_outputs(name)
+            .unwrap()
+            .map(|(_edge, node)| node.id)
+            .collect::<HashSet<NodeId>>()
+    }
+
     #[test]
     fn test_graph_edges() {
         let mut graph = RenderGraph::default();
@@ -688,22 +742,6 @@ mod tests {
         graph.add_node_edge("B", "C");
         graph.add_slot_edge("C", 0, "D", 0);
 
-        fn input_nodes(name: &'static str, graph: &RenderGraph) -> HashSet<NodeId> {
-            graph
-                .iter_node_inputs(name)
-                .unwrap()
-                .map(|(_edge, node)| node.id)
-                .collect::<HashSet<NodeId>>()
-        }
-
-        fn output_nodes(name: &'static str, graph: &RenderGraph) -> HashSet<NodeId> {
-            graph
-                .iter_node_outputs(name)
-                .unwrap()
-                .map(|(_edge, node)| node.id)
-                .collect::<HashSet<NodeId>>()
-        }
-
         assert!(input_nodes("A", &graph).is_empty(), "A has no inputs");
         assert!(
             output_nodes("A", &graph) == HashSet::from_iter(vec![c_id]),
@@ -803,4 +841,48 @@ mod tests {
             "Adding to a duplicate edge should return an error"
         );
     }
+
+    #[test]
+    fn test_add_node_edges() {
+        struct SimpleNode;
+        impl Node for SimpleNode {
+            fn run(
+                &self,
+                _graph: &mut RenderGraphContext,
+                _render_context: &mut RenderContext,
+                _world: &World,
+            ) -> Result<(), NodeRunError> {
+                Ok(())
+            }
+        }
+        impl FromWorld for SimpleNode {
+            fn from_world(_world: &mut World) -> Self {
+                Self
+            }
+        }
+
+        let mut graph = RenderGraph::default();
+        let a_id = graph.add_node("A", SimpleNode);
+        let b_id = graph.add_node("B", SimpleNode);
+        let c_id = graph.add_node("C", SimpleNode);
+
+        graph.add_node_edges(&["A", "B", "C"]);
+
+        assert!(
+            output_nodes("A", &graph) == HashSet::from_iter(vec![b_id]),
+            "A -> B"
+        );
+        assert!(
+            input_nodes("B", &graph) == HashSet::from_iter(vec![a_id]),
+            "B -> C"
+        );
+        assert!(
+            output_nodes("B", &graph) == HashSet::from_iter(vec![c_id]),
+            "B -> C"
+        );
+        assert!(
+            input_nodes("C", &graph) == HashSet::from_iter(vec![b_id]),
+            "B -> C"
+        );
+    }
 }

crates/bevy_render/src/render_graph/helpers.rs

@@ -0,0 +1,22 @@
+use bevy_app::App;
+use bevy_ecs::world::FromWorld;
+
+use super::{Node, RenderGraph};
+
+/// Utility function to add a [`Node`] to the [`RenderGraph`]
+/// * Create the [`Node`] using the [`FromWorld`] implementation
+/// * Add it to the graph
+/// * Automatically add the required node edges based on the given ordering
+pub fn add_node<T: Node + FromWorld>(
+    render_app: &mut App,
+    sub_graph_name: &'static str,
+    node_name: &'static str,
+    edges: &[&'static str],
+) {
+    let node = T::from_world(&mut render_app.world);
+    let mut render_graph = render_app.world.resource_mut::<RenderGraph>();
+
+    let graph = render_graph.get_sub_graph_mut(sub_graph_name).unwrap();
+    graph.add_node(node_name, node);
+    graph.add_node_edges(edges);
+}

crates/bevy_render/src/render_graph/mod.rs

@@ -1,12 +1,14 @@
 mod context;
 mod edge;
 mod graph;
+mod helpers;
 mod node;
 mod node_slot;
 
 pub use context::*;
 pub use edge::*;
 pub use graph::*;
+pub use helpers::*;
 pub use node::*;
 pub use node_slot::*;