Implement RenderGraphExt for World.

Author: andriyDev

crates/bevy_render/src/render_graph/app.rs

@@ -1,5 +1,5 @@
 use bevy_app::{App, SubApp};
-use bevy_ecs::world::FromWorld;
+use bevy_ecs::world::{FromWorld, World};
 use tracing::warn;
 
 use super::{IntoRenderNodeArray, Node, RenderGraph, RenderLabel, RenderSubGraph};
@@ -32,15 +32,15 @@ pub trait RenderGraphExt {
     ) -> &mut Self;
 }
 
-impl RenderGraphExt for SubApp {
+impl RenderGraphExt for World {
     fn add_render_graph_node<T: Node + FromWorld>(
         &mut self,
         sub_graph: impl RenderSubGraph,
         node_label: impl RenderLabel,
     ) -> &mut Self {
         let sub_graph = sub_graph.intern();
-        let node = T::from_world(self.world_mut());
-        let mut render_graph = self.world_mut().get_resource_mut::<RenderGraph>().expect(
+        let node = T::from_world(self);
+        let mut render_graph = self.get_resource_mut::<RenderGraph>().expect(
             "RenderGraph not found. Make sure you are using add_render_graph_node on the RenderApp",
         );
         if let Some(graph) = render_graph.get_sub_graph_mut(sub_graph) {
@@ -59,7 +59,7 @@ impl RenderGraphExt for SubApp {
         edges: impl IntoRenderNodeArray<N>,
     ) -> &mut Self {
         let sub_graph = sub_graph.intern();
-        let mut render_graph = self.world_mut().get_resource_mut::<RenderGraph>().expect(
+        let mut render_graph = self.get_resource_mut::<RenderGraph>().expect(
             "RenderGraph not found. Make sure you are using add_render_graph_edges on the RenderApp",
         );
         if let Some(graph) = render_graph.get_sub_graph_mut(sub_graph) {
@@ -79,7 +79,7 @@ impl RenderGraphExt for SubApp {
         input_node: impl RenderLabel,
     ) -> &mut Self {
         let sub_graph = sub_graph.intern();
-        let mut render_graph = self.world_mut().get_resource_mut::<RenderGraph>().expect(
+        let mut render_graph = self.get_resource_mut::<RenderGraph>().expect(
             "RenderGraph not found. Make sure you are using add_render_graph_edge on the RenderApp",
         );
         if let Some(graph) = render_graph.get_sub_graph_mut(sub_graph) {
@@ -93,21 +93,56 @@ impl RenderGraphExt for SubApp {
     }
 
     fn add_render_sub_graph(&mut self, sub_graph: impl RenderSubGraph) -> &mut Self {
-        let mut render_graph = self.world_mut().get_resource_mut::<RenderGraph>().expect(
+        let mut render_graph = self.get_resource_mut::<RenderGraph>().expect(
             "RenderGraph not found. Make sure you are using add_render_sub_graph on the RenderApp",
         );
         render_graph.add_sub_graph(sub_graph, RenderGraph::default());
         self
     }
 }
 
+impl RenderGraphExt for SubApp {
+    fn add_render_graph_node<T: Node + FromWorld>(
+        &mut self,
+        sub_graph: impl RenderSubGraph,
+        node_label: impl RenderLabel,
+    ) -> &mut Self {
+        World::add_render_graph_node::<T>(self.world_mut(), sub_graph, node_label);
+        self
+    }
+
+    fn add_render_graph_edge(
+        &mut self,
+        sub_graph: impl RenderSubGraph,
+        output_node: impl RenderLabel,
+        input_node: impl RenderLabel,
+    ) -> &mut Self {
+        World::add_render_graph_edge(self.world_mut(), sub_graph, output_node, input_node);
+        self
+    }
+
+    fn add_render_graph_edges<const N: usize>(
+        &mut self,
+        sub_graph: impl RenderSubGraph,
+        edges: impl IntoRenderNodeArray<N>,
+    ) -> &mut Self {
+        World::add_render_graph_edges(self.world_mut(), sub_graph, edges);
+        self
+    }
+
+    fn add_render_sub_graph(&mut self, sub_graph: impl RenderSubGraph) -> &mut Self {
+        World::add_render_sub_graph(self.world_mut(), sub_graph);
+        self
+    }
+}
+
 impl RenderGraphExt for App {
     fn add_render_graph_node<T: Node + FromWorld>(
         &mut self,
         sub_graph: impl RenderSubGraph,
         node_label: impl RenderLabel,
     ) -> &mut Self {
-        SubApp::add_render_graph_node::<T>(self.main_mut(), sub_graph, node_label);
+        World::add_render_graph_node::<T>(self.world_mut(), sub_graph, node_label);
         self
     }
 
@@ -117,7 +152,7 @@ impl RenderGraphExt for App {
         output_node: impl RenderLabel,
         input_node: impl RenderLabel,
     ) -> &mut Self {
-        SubApp::add_render_graph_edge(self.main_mut(), sub_graph, output_node, input_node);
+        World::add_render_graph_edge(self.world_mut(), sub_graph, output_node, input_node);
         self
     }
 
@@ -126,12 +161,12 @@ impl RenderGraphExt for App {
         sub_graph: impl RenderSubGraph,
         edges: impl IntoRenderNodeArray<N>,
     ) -> &mut Self {
-        SubApp::add_render_graph_edges(self.main_mut(), sub_graph, edges);
+        World::add_render_graph_edges(self.world_mut(), sub_graph, edges);
         self
     }
 
     fn add_render_sub_graph(&mut self, sub_graph: impl RenderSubGraph) -> &mut Self {
-        SubApp::add_render_sub_graph(self.main_mut(), sub_graph);
+        World::add_render_sub_graph(self.world_mut(), sub_graph);
         self
     }
 }