From 5639312ddd9c57f29fff5974337c2546efb7a06f Mon Sep 17 00:00:00 2001
From: Garry O'Donnell <garry.o'donnell@diamond.ac.uk>
Date: Thu, 30 Mar 2023 16:23:10 +0100
Subject: [PATCH] Add component & direct resource modes

---
 bevy_matchbox/src/lib.rs | 84 ++++++++++++++++++++++++++++++++--------
 1 file changed, 67 insertions(+), 17 deletions(-)

diff --git a/bevy_matchbox/src/lib.rs b/bevy_matchbox/src/lib.rs
index a51fc6b2..857a81c2 100644
--- a/bevy_matchbox/src/lib.rs
+++ b/bevy_matchbox/src/lib.rs
@@ -4,40 +4,89 @@
 
 use bevy::{
     ecs::system::Command,
-    prelude::{Commands, Deref, DerefMut, Resource, World},
+    prelude::{Commands, Component, Deref, DerefMut, Resource, World},
     tasks::IoTaskPool,
 };
 pub use matchbox_socket;
 use matchbox_socket::{BuildablePlurality, WebRtcSocket, WebRtcSocketBuilder};
 use std::marker::PhantomData;
 
-/// A [`Resource`] wrapping [`WebRtcSocket`].
+/// A [`WebRtcSocket`] as a [`Component`] or [`Resource`].
 ///
+/// As a [`Component`], directly
 /// ```
-/// # use bevy_matchbox::prelude::*;
-/// # use bevy::prelude::*;
-/// # fn system(mut commands: Commands) {
-/// let room_url = "ws://matchbox.example.com";
-/// commands.open_socket(WebRtcSocketBuilder::new(room_url).add_channel(ChannelConfig::reliable()));
-/// commands.close_socket::<SingleChannel>();
-/// # }
+/// use bevy_matchbox::prelude::*;
+/// use bevy::prelude::*;
+///
+/// fn open(mut commands: Commands) {
+///     let room_url = "ws://matchbox.example.com";
+///     let builder = WebRtcSocketBuilder::new(room_url).add_channel(ChannelConfig::reliable());
+///     commands.spawn(MatchboxSocket::from(builder));
+/// }
+///
+/// fn close<C: BuildablePlurality + 'static>(
+///     mut commands: Commands,
+///     socket: Query<Entity, With<MatchboxSocket<C>>>
+/// ) {
+///     let socket = socket.single();
+///     commands.entity(socket).despawn()
+/// }
+///
+/// ```
+///
+/// As a [`Resource`], with [`Commands`]
+/// ```
+/// use bevy_matchbox::prelude::*;
+/// use bevy::prelude::*;
+///
+/// fn open(mut commands: Commands) {
+///     let room_url = "ws://matchbox.example.com";
+///     commands.open_socket(WebRtcSocketBuilder::new(room_url).add_channel(ChannelConfig::reliable()));
+/// }
+///
+/// fn close(mut commands: Commands) {
+///     commands.close_socket::<SingleChannel>();
+/// }
+/// ```
+///
+/// As a [`Resource`], directly
+///
+/// ```
+/// use bevy_matchbox::prelude::*;
+/// use bevy::prelude::*;
+///
+/// fn open(mut commands: Commands) {
+///     let room_url = "ws://matchbox.example.com";
+///     let builder = WebRtcSocketBuilder::new(room_url).add_channel(ChannelConfig::reliable());
+///     commands.insert_resource(MatchboxSocket::from(builder));
+/// }
+///
+/// fn close(mut commands: Commands) {
+///     commands.remove_resource::<MatchboxSocket<SingleChannel>>();
+/// }
 /// ```
 ///
 /// To create and destroy this resource use the [`OpenSocket`] and [`CloseSocket`] [`Command`]s respectively.
-#[derive(Resource, Debug, Deref, DerefMut)]
+#[derive(Resource, Component, Debug, Deref, DerefMut)]
 pub struct MatchboxSocket<C: BuildablePlurality>(WebRtcSocket<C>);
 
+impl<C: BuildablePlurality> From<WebRtcSocketBuilder<C>> for MatchboxSocket<C> {
+    fn from(builder: WebRtcSocketBuilder<C>) -> Self {
+        let (socket, message_loop) = builder.build();
+
+        let task_pool = IoTaskPool::get();
+        task_pool.spawn(message_loop).detach();
+
+        MatchboxSocket(socket)
+    }
+}
+
 /// A [`Command`] used to open a [`MatchboxSocket`] and allocate it as a resource.
 struct OpenSocket<C: BuildablePlurality>(WebRtcSocketBuilder<C>);
 
 impl<C: BuildablePlurality + 'static> Command for OpenSocket<C> {
     fn write(self, world: &mut World) {
-        let (socket, message_loop) = self.0.build();
-
-        let task_pool = IoTaskPool::get();
-        task_pool.spawn(message_loop).detach();
-
-        world.insert_resource(MatchboxSocket(socket));
+        world.insert_resource(MatchboxSocket::from(self.0));
     }
 }
 
@@ -78,6 +127,7 @@ impl<'w, 's> CloseSocketExt for Commands<'w, 's> {
 pub mod prelude {
     pub use crate::{CloseSocketExt, MatchboxSocket, OpenSocketExt};
     pub use matchbox_socket::{
-        ChannelConfig, MultipleChannels, PeerId, PeerState, SingleChannel, WebRtcSocketBuilder,
+        BuildablePlurality, ChannelConfig, MultipleChannels, PeerId, PeerState, SingleChannel,
+        WebRtcSocketBuilder,
     };
 }