2 files changed, 27 insertions, 17 deletions
diff --git a/src/components.rs b/src/components.rs
index 8e25530..7ec88ba 100644
--- a/src/components.rs
+++ b/src/components.rs
@@ -1,19 +1,19 @@
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
 
 use bevy::{ecs::query::QueryItem, prelude::*, render::extract_component::ExtractComponent};
 
-use crate::render_assets::HeadlessRenderSource;
+use crate::render_assets;
 
 /// Headless render destination. Contains the image which the rendered frame is copied to.
 #[derive(Component, Default, Clone)]
 pub struct HeadlessRenderDestination(pub Arc<Mutex<Image>>);
 
 impl ExtractComponent for HeadlessRenderDestination {
-    type QueryData = (&'static Self, &'static Handle<HeadlessRenderSource>);
+    type QueryData = (&'static Self, &'static HeadlessRenderSource);
 
     type QueryFilter = ();
 
-    type Out = (Self, Handle<HeadlessRenderSource>);
+    type Out = (Self, HeadlessRenderSource);
 
     fn extract_component(
         (destination, source_handle): QueryItem<'_, Self::QueryData>,
@@ -22,11 +22,21 @@ impl ExtractComponent for HeadlessRenderDestination {
     }
 }
 
-/// Bundle containing both a source and destination for headless rendering.
-#[derive(Bundle)]
-pub struct HeadlessRenderBundle {
-    /// Source
-    pub source: Handle<HeadlessRenderSource>,
-    /// Destination
-    pub dest: HeadlessRenderDestination,
+impl HeadlessRenderDestination {
+    /// Get lock on this destination's image
+    pub fn image(&self) -> Result<MutexGuard<Image>, PoisonError<MutexGuard<Image>>> {
+        self.0.lock()
+    }
+}
+
+/// Headless render source
+#[derive(Component, Debug, Clone, DerefMut, Deref)]
+#[require(HeadlessRenderDestination, Camera3d)]
+pub struct HeadlessRenderSource(pub Handle<render_assets::HeadlessRenderSource>);
+
+impl HeadlessRenderSource {
+    /// Create a new headless render source from the provided image
+    pub fn new(asset_server: &AssetServer, image: Handle<Image>) -> Self {
+        Self(asset_server.add(render_assets::HeadlessRenderSource(image)))
+    }
 }
diff --git a/src/systems.rs b/src/systems.rs
index bb0efc7..64cf930 100644
--- a/src/systems.rs
+++ b/src/systems.rs
@@ -9,18 +9,18 @@ use bevy::{
 
 use pollster::FutureExt;
 
-use crate::{components::HeadlessRenderDestination, render_assets::{HeadlessRenderSource, GpuHeadlessRenderSource}};
+use crate::{
+    components::{HeadlessRenderDestination, HeadlessRenderSource},
+    render_assets::GpuHeadlessRenderSource,
+};
 
 pub fn copy_buffers(
-    mut headless_render_query: Query<(
-        &Handle<HeadlessRenderSource>,
-        &mut HeadlessRenderDestination,
-    )>,
+    mut headless_render_query: Query<(&HeadlessRenderSource, &mut HeadlessRenderDestination)>,
     sources: Res<RenderAssets<GpuHeadlessRenderSource>>,
     device: Res<RenderDevice>,
 ) {
     for (source_handle, destination_handle) in headless_render_query.iter_mut() {
-        let Some(gpu_source) = sources.get(source_handle) else {
+        let Some(gpu_source) = sources.get(source_handle.id()) else {
             continue;
         };