diff --git a/crates/border-wars/src/main.rs b/crates/border-wars/src/main.rs
index 85fb565..ef67c5e 100644
--- a/crates/border-wars/src/main.rs
+++ b/crates/border-wars/src/main.rs
@@ -2,12 +2,10 @@
 
 use bevy::prelude::*;
 use border_wars::scenes::ScenesPlugin;
-use border_wars::CurrentScene;
 
 fn main() {
     App::new()
         .add_plugins(DefaultPlugins)
-        .add_state::<CurrentScene>()
         .add_plugins(ScenesPlugin)
         .run();
 }
diff --git a/crates/border-wars/src/scenes/menu.rs b/crates/border-wars/src/scenes/menu.rs
index 8ede900..3efa8a6 100644
--- a/crates/border-wars/src/scenes/menu.rs
+++ b/crates/border-wars/src/scenes/menu.rs
@@ -10,10 +10,7 @@ pub struct MenuPlugin;
 
 impl Plugin for MenuPlugin {
     fn build(&self, app: &mut App) {
-        app.add_systems(
-            Update,
-            menu_ui.run_if(state_exists_and_equals(CurrentScene::Menu)),
-        );
+        app.add_systems(Update, menu_ui.run_if(in_state(CurrentScene::Menu)));
     }
 }
 /// Display the UI of the menu to host a game or join one.
diff --git a/crates/border-wars/src/scenes/mod.rs b/crates/border-wars/src/scenes/mod.rs
index 01865e3..06e9603 100644
--- a/crates/border-wars/src/scenes/mod.rs
+++ b/crates/border-wars/src/scenes/mod.rs
@@ -3,6 +3,8 @@
 use bevy::prelude::*;
 use bevy_egui::EguiPlugin;
 
+use crate::CurrentScene;
+
 pub mod menu;
 
 /// The plugin for all scenes.
@@ -10,6 +12,8 @@ pub struct ScenesPlugin;
 
 impl Plugin for ScenesPlugin {
     fn build(&self, app: &mut App) {
-        app.add_plugins(EguiPlugin).add_plugins(menu::MenuPlugin);
+        app.add_plugins(EguiPlugin)
+            .add_state::<CurrentScene>()
+            .add_plugins(menu::MenuPlugin);
     }
 }