How to return mutable reference to component?

[haath]

Hey guys,

Recently I tried to implement some utility functions for quickly retrieving a single component from a child entity.

First I implemented the following function, which compiles fine.

fn get_child_component<'a, C, Q, F>(children: &Children, query: &'a Query<'_, '_, Q, F>) -> Option<&'a C>
    where
        C: Component,
        Q: WorldQuery,
        F: ReadOnlyWorldQuery,
{
    for &child in children
    {
        if let Ok(component) = query.get_component::<C>(child)
        {
            return Some(component);
        }
    }

    None
}

---

But then I tried to implement the same for getting a mutable reference to a component, and I ran into some issues.
I ended up with the following implementation:

fn get_child_component_mut<'a, C, Q, F>(
    children: &Children,
    query: &'a mut Query<'_, '_, Q, F>,
) -> Option<Mut<'a, C>>
where
    C: Component,
    Q: WorldQuery,
    F: ReadOnlyWorldQuery,
{
    for &child in children
    {
        if let Ok(component) = query.get_component_mut::<C>(child)
        {
            return Some(component);
        }
    }

    None
}

Specifically the error is the following:

error[E0499]: cannot borrow `*query` as mutable more than once at a time
    |
114 | fn get_child_component_mut<'a, C, Q, F>(
    |                            -- lifetime `'a` defined here
...
125 |         if let Ok(component) = query.get_component_mut::<C>(child)
    |                                ^^^^^ `*query` was mutably borrowed here in the previous iteration of the loop
126 |         {
127 |             return Some(component);
    |                    --------------- returning this value requires that `*query` is borrowed for `'a`

And here I'm a little stuck. If the function requires a mutable reference to the query with lifetime 'a, then shouldn't I be able to return a mutable reference to the component with lifetime 'a as well?

The error even says that query needs to be borrowed for 'a, but it's already borrowed for 'a as far as I understand...

Or how do you suggest I do this? Is it even possible?

[jostly]

The problem isn't the lifetime, error E0499 is about borrowing as mutable more than once. This happens because you do it in a loop, and it seems the compiler can't detect that there would only be at most one mutable borrow active at the time.

You'd have to do this in two stages, first identify which child has the component (if any), and then borrow mutably for just that one entity.

[haath]

@jostly right indeed! Splitting it up fixed it.

let mut child_with_component = None;

for &child in self
{
    if query.get_component_mut::<C>(child).is_ok()
    {
        child_with_component = Some(child);
        break;
    }
}

child_with_component.map(|child| query.component_mut::<C>(child))

Thanks!