From f06e94fc0b178ac6942a2857053d782b32aed9b3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Mockers?= <mockersf@gmail.com>
Date: Fri, 23 Oct 2020 03:29:02 +0200
Subject: [PATCH] allow easier override of default plugins

---
 src/add_default_plugins.rs | 115 ++++++++++++++++++++++++++++++-------
 1 file changed, 94 insertions(+), 21 deletions(-)

diff --git a/src/add_default_plugins.rs b/src/add_default_plugins.rs
index ab3ecb789cb8b..2ba326a00d6fc 100644
--- a/src/add_default_plugins.rs
+++ b/src/add_default_plugins.rs
@@ -1,49 +1,122 @@
-use crate::app::AppBuilder;
+use crate::app::{AppBuilder, Plugin};
+use std::any::TypeId;
 
 pub trait AddDefaultPlugins {
+    /// Add all default plugins
     fn add_default_plugins(&mut self) -> &mut Self;
+    /// Add default plugins using `default_plugin_builder` to control which are added
+    fn add_default_plugins_with_builder(
+        &mut self,
+        default_plugin_builder: DefaultPluginBuilder,
+    ) -> &mut Self;
 }
 
-impl AddDefaultPlugins for AppBuilder {
-    fn add_default_plugins(&mut self) -> &mut Self {
-        self.add_plugin(bevy_type_registry::TypeRegistryPlugin::default());
-        self.add_plugin(bevy_core::CorePlugin::default());
-        self.add_plugin(bevy_transform::TransformPlugin::default());
-        self.add_plugin(bevy_diagnostic::DiagnosticsPlugin::default());
-        self.add_plugin(bevy_input::InputPlugin::default());
-        self.add_plugin(bevy_window::WindowPlugin::default());
-        self.add_plugin(bevy_asset::AssetPlugin::default());
-        self.add_plugin(bevy_scene::ScenePlugin::default());
+enum PluginState {
+    Default,
+    Custom(fn(&mut AppBuilder)),
+    None,
+}
+
+impl Default for PluginState {
+    fn default() -> Self {
+        PluginState::Default
+    }
+}
+
+/// Helper to enable / disable default plugins, or override a default plugin with a custom that will be loaded at
+/// the same time as the default one.
+#[derive(Default)]
+pub struct DefaultPluginBuilder {
+    plugins: std::collections::HashMap<TypeId, PluginState>,
+}
+
+impl DefaultPluginBuilder {
+    /// Override the default plugin `T` with a custom builder
+    pub fn with_custom<T: Plugin>(&mut self, builder: fn(&mut AppBuilder)) -> &mut Self {
+        self.plugins
+            .insert(TypeId::of::<T>(), PluginState::Custom(builder));
+        self
+    }
+
+    /// Keep using the default plugin `T`. This is the default
+    pub fn with_default<T: Plugin>(&mut self) -> &mut Self {
+        self.plugins.insert(TypeId::of::<T>(), PluginState::Default);
+        self
+    }
+
+    /// Disable the default plugin `T`
+    pub fn disable<T: Plugin>(&mut self) -> &mut Self {
+        self.plugins.insert(TypeId::of::<T>(), PluginState::None);
+        self
+    }
+
+    fn build_plugin<T: Plugin + Default>(&self, app: &mut AppBuilder) {
+        match self.plugins.get(&TypeId::of::<T>()) {
+            None | Some(PluginState::Default) => {
+                app.add_plugin(T::default());
+            }
+            Some(PluginState::Custom(custom_builder)) => {
+                custom_builder(app);
+            }
+            Some(PluginState::None) => (),
+        }
+    }
+
+    fn build(self, app: &mut AppBuilder) -> &mut AppBuilder {
+        self.build_plugin::<bevy_type_registry::TypeRegistryPlugin>(app);
+        self.build_plugin::<bevy_core::CorePlugin>(app);
+        self.build_plugin::<bevy_transform::TransformPlugin>(app);
+        self.build_plugin::<bevy_diagnostic::DiagnosticsPlugin>(app);
+        self.build_plugin::<bevy_input::InputPlugin>(app);
+        self.build_plugin::<bevy_window::WindowPlugin>(app);
+        self.build_plugin::<bevy_asset::AssetPlugin>(app);
+        self.build_plugin::<bevy_scene::ScenePlugin>(app);
 
         #[cfg(feature = "bevy_render")]
-        self.add_plugin(bevy_render::RenderPlugin::default());
+        self.build_plugin::<bevy_render::RenderPlugin>(app);
 
         #[cfg(feature = "bevy_sprite")]
-        self.add_plugin(bevy_sprite::SpritePlugin::default());
+        self.build_plugin::<bevy_sprite::SpritePlugin>(app);
 
         #[cfg(feature = "bevy_pbr")]
-        self.add_plugin(bevy_pbr::PbrPlugin::default());
+        self.build_plugin::<bevy_pbr::PbrPlugin>(app);
 
         #[cfg(feature = "bevy_ui")]
-        self.add_plugin(bevy_ui::UiPlugin::default());
+        self.build_plugin::<bevy_ui::UiPlugin>(app);
 
         #[cfg(feature = "bevy_text")]
-        self.add_plugin(bevy_text::TextPlugin::default());
+        self.build_plugin::<bevy_text::TextPlugin>(app);
 
         #[cfg(feature = "bevy_audio")]
-        self.add_plugin(bevy_audio::AudioPlugin::default());
+        self.build_plugin::<bevy_audio::AudioPlugin>(app);
 
         #[cfg(feature = "bevy_gilrs")]
-        self.add_plugin(bevy_gilrs::GilrsPlugin::default());
+        self.build_plugin::<bevy_gilrs::GilrsPlugin>(app);
 
         #[cfg(feature = "bevy_gltf")]
-        self.add_plugin(bevy_gltf::GltfPlugin::default());
+        self.build_plugin::<bevy_gltf::GltfPlugin>(app);
 
         #[cfg(feature = "bevy_winit")]
-        self.add_plugin(bevy_winit::WinitPlugin::default());
+        self.build_plugin::<bevy_winit::WinitPlugin>(app);
 
         #[cfg(feature = "bevy_wgpu")]
-        self.add_plugin(bevy_wgpu::WgpuPlugin::default());
+        self.build_plugin::<bevy_wgpu::WgpuPlugin>(app);
+        app
+    }
+}
+
+impl AddDefaultPlugins for AppBuilder {
+    fn add_default_plugins(&mut self) -> &mut Self {
+        DefaultPluginBuilder::default().build(self);
+
+        self
+    }
+
+    fn add_default_plugins_with_builder(
+        &mut self,
+        default_plugin_builder: DefaultPluginBuilder,
+    ) -> &mut Self {
+        default_plugin_builder.build(self);
 
         self
     }