Improve bevy_input_focus (#16749)

# Objective

I was curious to use the newly created `bevy_input_focus`, but I found
some issues with it
  - It was only implementing traits for `World`.
  - Lack of tests
  - `is_focus_within` logic was incorrect.


## Solution
 This PR includes some improvements to the `bevy_input_focus` crate: 
- Add new `IsFocusedHelper` that doesn't require access to `&World`. It
implements `IsFocused`
- Remove `IsFocused` impl for `DeferredWorld`. Since it already
implements `Deref<Target=World>` it was just duplication of code.
- impl `SetInputFocus` for `Commands`. There was no way to use
`SetFocusCommand` directly. This allows it.
- The `is_focus_within` logic has been fixed to check descendants.
Previously it was checking if any of the ancestors had focus which is
not correct according to the documentation.
  - Added a bunch of unit tests to verify the logic of the crate.

## Testing

- Did you test these changes? If so, how? Yes, running newly added unit
tests.

---

crates/bevy_input_focus/Cargo.toml

@@ -6,7 +6,7 @@ description = "Keyboard focus management"
 homepage = "https://bevyengine.org"
 repository = "https://github.com/bevyengine/bevy"
 license = "MIT OR Apache-2.0"
-keywords = ["bevy", "color"]
+keywords = ["bevy"]
 rust-version = "1.76.0"
 
 [dependencies]
@@ -14,9 +14,11 @@ bevy_app = { path = "../bevy_app", version = "0.15.0-dev", default-features = fa
 bevy_ecs = { path = "../bevy_ecs", version = "0.15.0-dev", default-features = false }
 bevy_input = { path = "../bevy_input", version = "0.15.0-dev", default-features = false }
 bevy_hierarchy = { path = "../bevy_hierarchy", version = "0.15.0-dev", default-features = false }
-bevy_utils = { path = "../bevy_utils", version = "0.15.0-dev", default-features = false }
 bevy_window = { path = "../bevy_window", version = "0.15.0-dev", default-features = false }
 
+[dev-dependencies]
+smol_str = "0.2"
+
 [lints]
 workspace = true
 

crates/bevy_input_focus/src/lib.rs

@@ -16,16 +16,16 @@
 //! This crate does *not* provide any integration with UI widgets, or provide functions for
 //! tab navigation or gamepad-based focus navigation, as those are typically application-specific.
 
-use bevy_app::{App, Plugin, Update};
+use bevy_app::{App, Plugin, PreUpdate};
 use bevy_ecs::{
     component::Component,
     entity::Entity,
     event::{Event, EventReader},
     query::With,
-    system::{Commands, Query, Res, Resource},
+    system::{Commands, Query, Res, Resource, SystemParam},
     world::{Command, DeferredWorld, World},
 };
-use bevy_hierarchy::Parent;
+use bevy_hierarchy::{HierarchyQueryExt, Parent};
 use bevy_input::keyboard::KeyboardInput;
 use bevy_window::PrimaryWindow;
 
@@ -41,7 +41,7 @@ pub struct InputFocus(pub Option<Entity>);
 #[derive(Clone, Debug, Resource)]
 pub struct InputFocusVisible(pub bool);
 
-/// Helper functions for [`World`] and [`DeferredWorld`] to set and clear input focus.
+/// Helper functions for [`World`],  [`DeferredWorld`] and [`Commands`] to set and clear input focus.
 pub trait SetInputFocus {
     /// Set input focus to the given entity.
     fn set_input_focus(&mut self, entity: Entity);
@@ -88,6 +88,16 @@ impl Command for SetFocusCommand {
     }
 }
 
+impl SetInputFocus for Commands<'_, '_> {
+    fn set_input_focus(&mut self, entity: Entity) {
+        self.queue(SetFocusCommand(Some(entity)));
+    }
+
+    fn clear_input_focus(&mut self) {
+        self.queue(SetFocusCommand(None));
+    }
+}
+
 /// A bubble-able event for keyboard input. This event is normally dispatched to the current
 /// input focus entity, if any. If no entity has input focus, then the event is dispatched to
 /// the main window.
@@ -108,7 +118,7 @@ impl Plugin for InputDispatchPlugin {
     fn build(&self, app: &mut App) {
         app.insert_resource(InputFocus(None))
             .insert_resource(InputFocusVisible(false))
-            .add_systems(Update, dispatch_keyboard_input);
+            .add_systems(PreUpdate, dispatch_keyboard_input);
     }
 }
 
@@ -137,7 +147,10 @@ fn dispatch_keyboard_input(
 }
 
 /// Trait which defines methods to check if an entity currently has focus. This is implemented
-/// for both [`World`] and [`DeferredWorld`].
+/// for [`World`] and [`IsFocusedHelper`].
+/// [`DeferredWorld`] indirectly implements it through [`Deref`].
+///
+/// [`Deref`]: std::ops::Deref
 pub trait IsFocused {
     /// Returns true if the given entity has input focus.
     fn is_focused(&self, entity: Entity) -> bool;
@@ -148,97 +161,270 @@ pub trait IsFocused {
     /// Returns true if the given entity has input focus and the focus indicator is visible.
     fn is_focus_visible(&self, entity: Entity) -> bool;
 
-    /// Returns true if the given entity, or any descenant, has input focus and the focus
+    /// Returns true if the given entity, or any descendant, has input focus and the focus
     /// indicator is visible.
     fn is_focus_within_visible(&self, entity: Entity) -> bool;
 }
 
-impl IsFocused for DeferredWorld<'_> {
+/// System param that helps get information about the current focused entity.
+#[derive(SystemParam)]
+pub struct IsFocusedHelper<'w, 's> {
+    parent_query: Query<'w, 's, &'static Parent>,
+    input_focus: Option<Res<'w, InputFocus>>,
+    input_focus_visible: Option<Res<'w, InputFocusVisible>>,
+}
+
+impl IsFocused for IsFocusedHelper<'_, '_> {
     fn is_focused(&self, entity: Entity) -> bool {
-        self.get_resource::<InputFocus>()
-            .map(|f| f.0)
-            .unwrap_or_default()
-            .map(|f| f == entity)
-            .unwrap_or_default()
+        self.input_focus
+            .as_deref()
+            .and_then(|f| f.0)
+            .is_some_and(|e| e == entity)
     }
 
     fn is_focus_within(&self, entity: Entity) -> bool {
-        let Some(focus_resource) = self.get_resource::<InputFocus>() else {
+        let Some(focus) = self.input_focus.as_deref().and_then(|f| f.0) else {
             return false;
         };
-        let Some(focus) = focus_resource.0 else {
-            return false;
-        };
-        let mut e = entity;
-        loop {
-            if e == focus {
-                return true;
-            }
-            if let Some(parent) = self.entity(e).get::<Parent>() {
-                e = parent.get();
-            } else {
-                break;
-            }
+        if focus == entity {
+            return true;
         }
-        false
+        self.parent_query.iter_ancestors(focus).any(|e| e == entity)
     }
 
     fn is_focus_visible(&self, entity: Entity) -> bool {
-        self.get_resource::<InputFocusVisible>()
-            .map(|vis| vis.0)
-            .unwrap_or_default()
-            && self.is_focused(entity)
+        self.input_focus_visible.as_deref().is_some_and(|vis| vis.0) && self.is_focused(entity)
     }
 
     fn is_focus_within_visible(&self, entity: Entity) -> bool {
-        self.get_resource::<InputFocusVisible>()
-            .map(|vis| vis.0)
-            .unwrap_or_default()
-            && self.is_focus_within(entity)
+        self.input_focus_visible.as_deref().is_some_and(|vis| vis.0) && self.is_focus_within(entity)
     }
 }
 
 impl IsFocused for World {
     fn is_focused(&self, entity: Entity) -> bool {
         self.get_resource::<InputFocus>()
-            .map(|f| f.0)
-            .unwrap_or_default()
-            .map(|f| f == entity)
-            .unwrap_or_default()
+            .and_then(|f| f.0)
+            .is_some_and(|f| f == entity)
     }
 
     fn is_focus_within(&self, entity: Entity) -> bool {
-        let Some(focus_resource) = self.get_resource::<InputFocus>() else {
+        let Some(focus) = self.get_resource::<InputFocus>().and_then(|f| f.0) else {
             return false;
         };
-        let Some(focus) = focus_resource.0 else {
-            return false;
-        };
-        let mut e = entity;
+        let mut e = focus;
         loop {
-            if e == focus {
+            if e == entity {
                 return true;
             }
-            if let Some(parent) = self.entity(e).get::<Parent>() {
-                e = parent.get();
+            if let Some(parent) = self.entity(e).get::<Parent>().map(Parent::get) {
+                e = parent;
             } else {
-                break;
+                return false;
             }
         }
-        false
     }
 
     fn is_focus_visible(&self, entity: Entity) -> bool {
         self.get_resource::<InputFocusVisible>()
-            .map(|vis| vis.0)
-            .unwrap_or_default()
+            .is_some_and(|vis| vis.0)
             && self.is_focused(entity)
     }
 
     fn is_focus_within_visible(&self, entity: Entity) -> bool {
         self.get_resource::<InputFocusVisible>()
-            .map(|vis| vis.0)
-            .unwrap_or_default()
+            .is_some_and(|vis| vis.0)
             && self.is_focus_within(entity)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use bevy_ecs::{component::ComponentId, observer::Trigger, system::RunSystemOnce};
+    use bevy_hierarchy::BuildChildren;
+    use bevy_input::{
+        keyboard::{Key, KeyCode},
+        ButtonState, InputPlugin,
+    };
+    use smol_str::SmolStr;
+
+    #[derive(Component)]
+    #[component(on_add = set_focus_on_add)]
+    struct SetFocusOnAdd;
+
+    fn set_focus_on_add(mut world: DeferredWorld, entity: Entity, _: ComponentId) {
+        world.set_input_focus(entity);
+    }
+
+    #[derive(Component, Default)]
+    struct GatherKeyboardEvents(String);
+
+    fn gather_keyboard_events(
+        trigger: Trigger<FocusKeyboardInput>,
+        mut query: Query<&mut GatherKeyboardEvents>,
+    ) {
+        if let Ok(mut gather) = query.get_mut(trigger.target()) {
+            if let Key::Character(c) = &trigger.0.logical_key {
+                gather.0.push_str(c.as_str());
+            }
+        }
+    }
+
+    const KEY_A_EVENT: KeyboardInput = KeyboardInput {
+        key_code: KeyCode::KeyA,
+        logical_key: Key::Character(SmolStr::new_static("A")),
+        state: ButtonState::Pressed,
+        repeat: false,
+        window: Entity::PLACEHOLDER,
+    };
+
+    #[test]
+    fn test_without_plugin() {
+        let mut app = App::new();
+
+        let entity = app.world_mut().spawn_empty().id();
+
+        app.world_mut().set_input_focus(entity);
+        assert!(!app.world().is_focused(entity));
+
+        app.world_mut()
+            .run_system_once(move |helper: IsFocusedHelper| {
+                assert!(!helper.is_focused(entity));
+                assert!(!helper.is_focus_within(entity));
+                assert!(!helper.is_focus_visible(entity));
+                assert!(!helper.is_focus_within_visible(entity));
+            })
+            .unwrap();
+
+        app.world_mut()
+            .run_system_once(move |world: DeferredWorld| {
+                assert!(!world.is_focused(entity));
+                assert!(!world.is_focus_within(entity));
+                assert!(!world.is_focus_visible(entity));
+                assert!(!world.is_focus_within_visible(entity));
+            })
+            .unwrap();
+    }
+
+    #[test]
+    fn test_keyboard_events() {
+        fn get_gathered(app: &App, entity: Entity) -> &str {
+            app.world()
+                .entity(entity)
+                .get::<GatherKeyboardEvents>()
+                .unwrap()
+                .0
+                .as_str()
+        }
+
+        let mut app = App::new();
+
+        app.add_plugins((InputPlugin, InputDispatchPlugin))
+            .add_observer(gather_keyboard_events);
+
+        let entity_a = app
+            .world_mut()
+            .spawn((GatherKeyboardEvents::default(), SetFocusOnAdd))
+            .id();
+
+        let child_of_b = app
+            .world_mut()
+            .spawn((GatherKeyboardEvents::default(),))
+            .id();
+
+        let entity_b = app
+            .world_mut()
+            .spawn((GatherKeyboardEvents::default(),))
+            .add_child(child_of_b)
+            .id();
+
+        assert!(app.world().is_focused(entity_a));
+        assert!(!app.world().is_focused(entity_b));
+        assert!(!app.world().is_focused(child_of_b));
+        assert!(!app.world().is_focus_visible(entity_a));
+        assert!(!app.world().is_focus_visible(entity_b));
+        assert!(!app.world().is_focus_visible(child_of_b));
+
+        // entity_a should receive this event
+        app.world_mut().send_event(KEY_A_EVENT);
+        app.update();
+
+        assert_eq!(get_gathered(&app, entity_a), "A");
+        assert_eq!(get_gathered(&app, entity_b), "");
+        assert_eq!(get_gathered(&app, child_of_b), "");
+
+        app.world_mut().clear_input_focus();
+
+        assert!(!app.world().is_focused(entity_a));
+        assert!(!app.world().is_focus_visible(entity_a));
+
+        // This event should be lost
+        app.world_mut().send_event(KEY_A_EVENT);
+        app.update();
+
+        assert_eq!(get_gathered(&app, entity_a), "A");
+        assert_eq!(get_gathered(&app, entity_b), "");
+        assert_eq!(get_gathered(&app, child_of_b), "");
+
+        app.world_mut().set_input_focus(entity_b);
+        assert!(app.world().is_focused(entity_b));
+        assert!(!app.world().is_focused(child_of_b));
+
+        app.world_mut()
+            .run_system_once(move |mut commands: Commands| {
+                commands.set_input_focus(child_of_b);
+            })
+            .unwrap();
+        assert!(app.world().is_focus_within(entity_b));
+
+        // These events should be received by entity_b and child_of_b
+        app.world_mut().send_event_batch([KEY_A_EVENT; 4]);
+        app.update();
+
+        assert_eq!(get_gathered(&app, entity_a), "A");
+        assert_eq!(get_gathered(&app, entity_b), "AAAA");
+        assert_eq!(get_gathered(&app, child_of_b), "AAAA");
+
+        app.world_mut().resource_mut::<InputFocusVisible>().0 = true;
+
+        app.world_mut()
+            .run_system_once(move |helper: IsFocusedHelper| {
+                assert!(!helper.is_focused(entity_a));
+                assert!(!helper.is_focus_within(entity_a));
+                assert!(!helper.is_focus_visible(entity_a));
+                assert!(!helper.is_focus_within_visible(entity_a));
+
+                assert!(!helper.is_focused(entity_b));
+                assert!(helper.is_focus_within(entity_b));
+                assert!(!helper.is_focus_visible(entity_b));
+                assert!(helper.is_focus_within_visible(entity_b));
+
+                assert!(helper.is_focused(child_of_b));
+                assert!(helper.is_focus_within(child_of_b));
+                assert!(helper.is_focus_visible(child_of_b));
+                assert!(helper.is_focus_within_visible(child_of_b));
+            })
+            .unwrap();
+
+        app.world_mut()
+            .run_system_once(move |world: DeferredWorld| {
+                assert!(!world.is_focused(entity_a));
+                assert!(!world.is_focus_within(entity_a));
+                assert!(!world.is_focus_visible(entity_a));
+                assert!(!world.is_focus_within_visible(entity_a));
+
+                assert!(!world.is_focused(entity_b));
+                assert!(world.is_focus_within(entity_b));
+                assert!(!world.is_focus_visible(entity_b));
+                assert!(world.is_focus_within_visible(entity_b));
+
+                assert!(world.is_focused(child_of_b));
+                assert!(world.is_focus_within(child_of_b));
+                assert!(world.is_focus_visible(child_of_b));
+                assert!(world.is_focus_within_visible(child_of_b));
+            })
+            .unwrap();
+    }
+}