From 3220ad6b0bb11ba293c9e3186d90ee1ca090689f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois?= <mockersf@gmail.com>
Date: Tue, 18 Apr 2023 18:31:55 +0200
Subject: [PATCH] do not crash when rendering only one gizmo (#8434)

# Objective

- Fixes #8432

## Solution

- Do not even create mesh when one is not needed
---
 crates/bevy_gizmos/src/gizmos.rs |   1 -
 crates/bevy_gizmos/src/lib.rs    | 102 ++++++++++++++++++-------------
 2 files changed, 61 insertions(+), 42 deletions(-)

diff --git a/crates/bevy_gizmos/src/gizmos.rs b/crates/bevy_gizmos/src/gizmos.rs
index 04a65c77f5f14..10e398d095d26 100644
--- a/crates/bevy_gizmos/src/gizmos.rs
+++ b/crates/bevy_gizmos/src/gizmos.rs
@@ -20,7 +20,6 @@ pub(crate) struct GizmoStorage {
     pub list_colors: Vec<ColorItem>,
     pub strip_positions: Vec<PositionItem>,
     pub strip_colors: Vec<ColorItem>,
-    pub in_use: bool,
 }
 
 /// A [`SystemParam`](bevy_ecs::system::SystemParam) for drawing gizmos.
diff --git a/crates/bevy_gizmos/src/lib.rs b/crates/bevy_gizmos/src/lib.rs
index 65c7fa51f532d..01a2067a7a08d 100644
--- a/crates/bevy_gizmos/src/lib.rs
+++ b/crates/bevy_gizmos/src/lib.rs
@@ -126,17 +126,15 @@ impl Default for GizmoConfig {
 
 #[derive(Resource)]
 struct MeshHandles {
-    list: Handle<Mesh>,
-    strip: Handle<Mesh>,
+    list: Option<Handle<Mesh>>,
+    strip: Option<Handle<Mesh>>,
 }
 
 impl FromWorld for MeshHandles {
-    fn from_world(world: &mut World) -> Self {
-        let mut meshes = world.resource_mut::<Assets<Mesh>>();
-
+    fn from_world(_world: &mut World) -> Self {
         MeshHandles {
-            list: meshes.add(Mesh::new(PrimitiveTopology::LineList)),
-            strip: meshes.add(Mesh::new(PrimitiveTopology::LineStrip)),
+            list: None,
+            strip: None,
         }
     }
 }
@@ -146,33 +144,51 @@ struct GizmoMesh;
 
 fn update_gizmo_meshes(
     mut meshes: ResMut<Assets<Mesh>>,
-    handles: Res<MeshHandles>,
+    mut handles: ResMut<MeshHandles>,
     mut storage: ResMut<GizmoStorage>,
 ) {
-    let list_mesh = meshes.get_mut(&handles.list).unwrap();
+    if storage.list_positions.is_empty() {
+        handles.list = None;
+    } else if let Some(handle) = handles.list.as_ref() {
+        let list_mesh = meshes.get_mut(handle).unwrap();
 
-    storage.in_use = false;
+        let positions = mem::take(&mut storage.list_positions);
+        list_mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
 
-    if !storage.list_positions.is_empty() {
-        storage.in_use = true;
+        let colors = mem::take(&mut storage.list_colors);
+        list_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, colors);
+    } else {
+        let mut list_mesh = Mesh::new(PrimitiveTopology::LineList);
 
         let positions = mem::take(&mut storage.list_positions);
         list_mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
 
         let colors = mem::take(&mut storage.list_colors);
         list_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, colors);
+
+        handles.list = Some(meshes.add(list_mesh));
     }
 
-    if !storage.strip_positions.is_empty() {
-        storage.in_use = true;
+    if storage.strip_positions.is_empty() {
+        handles.strip = None;
+    } else if let Some(handle) = handles.strip.as_ref() {
+        let strip_mesh = meshes.get_mut(handle).unwrap();
 
-        let strip_mesh = meshes.get_mut(&handles.strip).unwrap();
+        let positions = mem::take(&mut storage.strip_positions);
+        strip_mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
+
+        let colors = mem::take(&mut storage.strip_colors);
+        strip_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, colors);
+    } else {
+        let mut strip_mesh = Mesh::new(PrimitiveTopology::LineStrip);
 
         let positions = mem::take(&mut storage.strip_positions);
         strip_mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
 
         let colors = mem::take(&mut storage.strip_colors);
         strip_mesh.insert_attribute(Mesh::ATTRIBUTE_COLOR, colors);
+
+        handles.strip = Some(meshes.add(strip_mesh));
     }
 }
 
@@ -180,40 +196,44 @@ fn extract_gizmo_data(
     mut commands: Commands,
     handles: Extract<Res<MeshHandles>>,
     config: Extract<Res<GizmoConfig>>,
-    storage: Extract<Res<GizmoStorage>>,
 ) {
     if config.is_changed() {
         commands.insert_resource(**config);
     }
 
-    if !config.enabled || !storage.in_use {
+    if !config.enabled {
         return;
     }
 
     let transform = Mat4::IDENTITY;
     let inverse_transpose_model = transform.inverse().transpose();
-    commands.spawn_batch([&handles.list, &handles.strip].map(|handle| {
-        (
-            GizmoMesh,
-            #[cfg(feature = "bevy_pbr")]
-            (
-                handle.clone_weak(),
-                MeshUniform {
-                    flags: 0,
-                    transform,
-                    previous_transform: transform,
-                    inverse_transpose_model,
-                },
-            ),
-            #[cfg(feature = "bevy_sprite")]
-            (
-                Mesh2dHandle(handle.clone_weak()),
-                Mesh2dUniform {
-                    flags: 0,
-                    transform,
-                    inverse_transpose_model,
-                },
-            ),
-        )
-    }));
+    commands.spawn_batch(
+        [handles.list.clone(), handles.strip.clone()]
+            .into_iter()
+            .flatten()
+            .map(move |handle| {
+                (
+                    GizmoMesh,
+                    #[cfg(feature = "bevy_pbr")]
+                    (
+                        handle.clone_weak(),
+                        MeshUniform {
+                            flags: 0,
+                            transform,
+                            previous_transform: transform,
+                            inverse_transpose_model,
+                        },
+                    ),
+                    #[cfg(feature = "bevy_sprite")]
+                    (
+                        Mesh2dHandle(handle),
+                        Mesh2dUniform {
+                            flags: 0,
+                            transform,
+                            inverse_transpose_model,
+                        },
+                    ),
+                )
+            }),
+    );
 }