From c57f0588af783b458ffd1c66f1d0a8b071d9a674 Mon Sep 17 00:00:00 2001
From: Danny McGee <dannymcgee@gmail.com>
Date: Wed, 28 Feb 2024 01:30:58 -0500
Subject: [PATCH] Move `track_scene_elements` to `change_tracking`

Also track all `SceneElement`-bearing entities
---
 studio/src/change_tracking.rs | 15 +++++++++++++++
 studio/src/scene.rs           | 20 +++-----------------
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/studio/src/change_tracking.rs b/studio/src/change_tracking.rs
index a3f4e32..1b92f44 100644
--- a/studio/src/change_tracking.rs
+++ b/studio/src/change_tracking.rs
@@ -1,5 +1,7 @@
 use bevy::prelude::*;
 
+use crate::scene::{SceneElement, SceneSpawnerSystem, SceneState};
+
 pub struct ChangeTrackingPlugin;
 
 impl Plugin for ChangeTrackingPlugin {
@@ -8,6 +10,10 @@ impl Plugin for ChangeTrackingPlugin {
 
         // Add more of these systems for each component that we want to track
         app.add_systems(Update, send_change_events::<Transform>);
+        app.add_systems(
+            OnEnter(SceneState::Active),
+            track_scene_elements.after(SceneSpawnerSystem),
+        );
 
         #[cfg(feature = "wasm")]
         app.add_systems(
@@ -29,6 +35,15 @@ fn send_change_events<T: Component>(
     }
 }
 
+fn track_scene_elements(
+    mut cmd: Commands,
+    q_scene_elements: Query<Entity, (With<Transform>, With<SceneElement>, Without<ChangeTracker>)>,
+) {
+    for element in q_scene_elements.iter() {
+        cmd.entity(element).insert(ChangeTracker);
+    }
+}
+
 #[derive(Event, Deref)]
 pub struct EntityChangedEvent(pub Entity);
 
diff --git a/studio/src/scene.rs b/studio/src/scene.rs
index 5c51325..9cd07df 100644
--- a/studio/src/scene.rs
+++ b/studio/src/scene.rs
@@ -23,7 +23,6 @@ use wasm_bindgen::prelude::wasm_bindgen;
 use wgpu::{TextureDescriptor, TextureUsages};
 
 use crate::{
-    anim::SkeletalJoint,
     camera_controller::{CameraController, CameraControllerBundle},
     change_tracking::{ChangeTracker, ChangeTrackingPlugin},
     debug::DebugPoint,
@@ -63,13 +62,9 @@ impl Plugin for ScenePlugin {
             .add_systems(
                 OnEnter(SceneState::Active),
                 (
-                    (
-                        tag_scene_elements,
-                        frame_scene_in_view,
-                        start_capture.run_if(is_headless),
-                    ),
-                    apply_deferred,
-                    track_scene_elements,
+                    tag_scene_elements,
+                    frame_scene_in_view,
+                    start_capture.run_if(is_headless),
                 )
                     .chain()
                     .in_set(SceneSpawnerSystem),
@@ -524,15 +519,6 @@ fn start_capture(mut captures: Query<&mut ScreenCapture>) {
     }
 }
 
-fn track_scene_elements(
-    mut commands: Commands,
-    scene_elements: Query<Entity, (With<Transform>, With<SkeletalJoint>)>,
-) {
-    for element in scene_elements.iter() {
-        commands.entity(element).insert(ChangeTracker);
-    }
-}
-
 #[cfg(feature = "wasm")]
 mod wasm {
     use bevy::{prelude::*, window::PrimaryWindow};