Fast Component Lookup-by-Value with Indexes

[alice-i-cecile]

Motivation

When writing games, it's common to want to find an entity that has a component with a particular value. In one particularly important case, suppose we want to find all entities in a particular Position on our grid.

In Bevy, this is straightforward enough: query for that component, iterate through it, then see which values match. But doing so incurs a pretty heavy cost: every time we need to do this, we're scanning through every single entity with that component, and this operation takes linear time.

In a game where this is a common operation, and the number of entities to search is large, it would be nice to be able to have a quick way to look this up. Create some sort of hashmap from your component's value to their entity, insert your components when they're changed, and then look it up in nice constant time.

Sounds easy right? But getting all the details right is surprisingly devilish, performance is at a premium and bugs deep in your internal code can be a nightmare to debug.

Hence: a proposal for an Official Bevy Index.

Constraints

1. Must be faster than simply iterating through entities for reasonably large collections (at least n=10000, ideally n=1000).
2. The overhead of creating and maintaining the index shouldn't outweigh the benefits gained, especially for slowly-changing components.
3. Using these indexes should be ergonomic: low-boilerplate, no obvious pitfalls, no need for manual updating.
4. The values provided must always reflect the underlying game state.
5. When multiple entities have a component with the same value, choose whether to extract all of those entities, or the first one.
6. As shown below, any system that uses an Index<T> needs the same scheduling constraints as one that reads T directly.

Approaches

In all cases, our goal is to provide our system a set of all entities for which a specified component has the given value.

Iteration

As a baseline, here's how you might do this with iteration.

fn my_system(query: Query<&MyComponent>) {
	let mut set = HashSet::new();
	for c in query.iter().cloned(){
		if *c == target {
			set.insert(*c);
		}
	}
}

Clean and easy, but the linear time will be a real deal-breaker for some applications.

Naive Hashmap

So, let's use a hashmap!
We want to look up our entities by their component value, so we'll go with a Hashmap<MyComponent, Entity>. Stick in your target value, and get out an entity.

But wait: what if we have multiple entities with the same value? We'll only get one of them back out, and updating our data structure will overwrite other good values.

Simple Multimap

Alright, so we'll swap to a multimap: so we can store multiple values for each given key. Problem solved, right?

Now we just need to add some nice updating logic and we can get back to actually making our game. So what do we need to do?

We don't want to completely rewrite our index every time we update it (otherwise why aren't you just using iteration?), so we need to modify only the values that we've found. We can grab a list of components that will need to be updated with Bevy's handy Changed functionality. We'll need to be particularly careful here to remove the values for entities that have been removed: otherwise we'll end up calling .get on invalid entities and throwing panics.

But then how do we know where the old values that we need to remove are?

Once-per-frame Multimap-Powered Index

We could iterate through the dictionary in order, but now we're back to linear (in the number of unique component values) time. Instead, we need a reverse mapping too. We can't really on Bevy's built-in queries with .get here, since we need the old value of our component. So instead, we're left with creating and maintaining a reverse mapping on our own, which can be a simple hashmap since we know every entity will have exactly one component.

Wrap both the forward and reverse mappings in a custom Index struct and we can move on to scheduling our updates! Update it after startup, and then run it at the end of each frame, spin up a nice little utility system to initialize all this and ta-da, effortless indexing right?

Not so much: the values of these components can readily change within the frame, resulting in us working off of stale values to bizarre and catastrophic effect.

Manually-timed Multimap-Powered Index

Alright, so we need to make sure that our index is fresh every time that we're using it. There are two broad approaches here:

1. Update the index after each system that changes the underlying component.
2. Update the index before each system that uses the index.

If a single value of our component is changed multiple times before we use it though, we're going to be wasting effort, meaning that the second option is preferred.

So then toss an update_index<T>.system() before each time we want to call our index, ensuring that it's fresh.
The boilerplate isn't great but surely this should always work (until we forget to add the boilerplate), right?

Unfortunately, the magic of parallelism can work in unexpected ways. Suppose our system that uses the resource doesn't actually care about our IndexComponent directly: it might just be trying to find all units with a given position, rather than modifying position directly. But that means that it's not querying for IndexComponent, merely accessing Res<Index<IndexComponent>> and so isn't getting a hard or soft lock on the values from our scheduler in the same way as if it was reading or writing the component through a query.

This means that we can have another system running in parallel, modifying the values of our IndexComponent out from underneath us as we're relying on them in a way that could never happen with the iterative approach!
Like with most race-condition-flavored problems, tracking down a bug caused by this behavior sounds particularly nasty.

As an extra wart, we still need to keep the end-of-tick updates too, otherwise we might drop changes, since Changed is cleared at the end of each tick. Unfortunately, this will duplicate all of the work that we've done previously in the current tick. Keeping around a list of entities that we've updated this tick won't solve this duplication: we can't be sure that the value hasn't changed again on us in the meantime.

Index as a Bevy Primitive

I hope that this design discussion has convinced you, dear reader, that this index functionality is natural, appealing and deceptively tricky. If we're going to be mucking around with engine internals anyways, let's make a nice API for our end users.

1. You can use an index for the component type <T> in your systems by adding an Index<T> to your function arguments. We don't want to use a Res<Index<T>> here, because manually mucking around with your indexes is a recipe for disaster and it results in the race-condition bug described above.
2. An index for the type T should be automatically initialized and maintained (like Local) whenever we have at least one system is added that has an argument of Index<T>. Doing this manually is boilerplate, but also easy to screw up if you add it after values of the component start changing.
3. Indexes have one public method .get(k: T) -> HashSet<Entity> (which may be empty). The order that entities are stored in our index is completely arbitrary, so we don't want to expose it lest people try an rely on it.
4. All of the updating should happen magically under the hood, and should Just Work no matter what weird systems the user might dream up.

This lets us write our initial example code as:

fn my_system(index: Index<MyComponent>) {
	let set = index.get(target);
}

Of course, if n is small or entities change often, performance may be better with simple iteration.

Steps 1, 2 and 3 of this design are easy enough, if we're privileged with access to Bevy's internals.
The real challenge is designing an efficient, foolproof approach to updating.
Here are some ideas:

1. Dedicated update systems: Automatically run an updating system before every system with an Index parameter (as long as there's been an intervening system that could have changed the component). This is much less painful without the boilerplate, and won't lead to race conditions now that we can tweak our scheduler to be smarter. We'll still need to run the end-of-frame update system though.
2. Metaprogramming system modification: Modify the code of every system with an index argument to actually expand out to a Query on the appropriate component and combine it into the same step. This lets us avoid modifying the scheduler logic, but it's very hacky and we still need the end-of-tick update system.
3. DerefMut magic: Whenever DerefMut is called on a component, also update our index if it exists. This wonderful bit of magic is already used by our Changed component flag. This works flawlessly, but could have some painful cache issues as we're only inserting one change at a time, and could result in us wasting effort if a component is changed multiple times before our index is used.
4. ToIndex component flag: Whenever DerefMut is called on a component, set its ToIndex component flag (analogous to Changed) to true. Automatically schedule a system like in 1, then process all of the updates at once in a batch. Unlike Changed, don't reset these flags at the end of the tick. This may result in better branch-prediction, as the work done in our tight component-changing loops is always the same.
5. Time-stamped Changed flags: One of the possible solutions to System-order-independent ECS change tracking #68 involves recording what system a change was made in, and ensuring that changes always persist for exactly one frame. If we had this behavior, we could use this in combination with Dedicated System Updates in order to avoid double-processing, eliminating its largest drawback without needing to muck around with DerefMut, making Index a zero-cost feature when not used, and allowing much cleaner extension for similar behavior.

I prefer 5, but that's contingent on a modification to the underlying Changed behavior. Without that, I suspect that option 3 or 4 is going to be correct, as they're elegant, foolproof and avoid likely-redundant end-of-tick cleanup work. Which one is better will likely come down to benchmarking under realistic-ish workloads. I hate 2, and would really like to avoid it.

As an extra note: option 3 is the only choice that is fully robust to arbitrarily inserting systems once the app is running. Support of that feature is currently an open question, but as you can see here comes with significant constraints.

Steps to Implement

1. Rewrite the IntoSystem macro to accommodate arguments of type Index. Add the SystemParam trait to Index to get it to play nicely with the IntoSystem macro.
2. Ensure that the scheduler treats systems with an Index<T> argument as reading T.
3. When a system with an Index argument is run, hand it the appropriate Index struct.
4. Write a nice little struct for holding our forward and reverse maps. See here for an initial attempt.
5. Automatically update the index at the appropriate time, as discussed directly above this section.
6. Write some tests that cover the edge cases listed in this design doc to ensure that it actually works.

[alice-i-cecile]

As a bit of a digression, this is one particularly useful case of caching the results of some computation on our components. Pretty much any attempt to solve similar problems via memoization will run into the same nasty timing / freshness issues, that will be nearly impossible to solve elegantly from outside of the engine itself.

Do we want to create a reusable public interface for these sorts of tasks, and then implement Index as one example of this?

[mockersf]

There is also a Labels component that can be added to an entity, and that can be used to search entities by label using the resource EntityLabels

It's maintained by the system entity_labels_system that uses a query Query<(Entity, &Labels), Changed<Labels>> to update the hashmap

[alice-i-cecile]

TIL, thanks. It's a nice little implementation of this idea in general, but it looks like it can suffer from staleness issues in the same way, because it's only running when you manually add it.

I think that ideally we solve this computation-caching problem in general, and update Labels along with it.

[Ratysz]

I've spent some time experimenting with the pattern we talked about on Discord (foregoing housekeeping systems and instead keeping data always valid through API or reference coupling of the index and indexed component), and in the process of trying to make it more concurrency-friendly landed on this unsavory snippet:

click me

use bevy::prelude::*;
use dashmap::{DashMap, DashSet};
use std::{hash::Hash, sync::Arc};

pub struct Index<T: Hash + Eq> {
    map: Arc<DashMap<T, DashSet<Entity>>>,
}

impl<T: Hash + Eq> Default for Index<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: Hash + Eq> Index<T> {
    pub fn new() -> Self {
        Self {
            map: Arc::new(DashMap::new()),
        }
    }

    pub fn entities(&self, key: &T) -> Option<Vec<Entity>> {
        self.map
            .as_ref()
            .get(key)
            .map(|entities| entities.iter().map(|entity| *entity).collect())
    }

    pub fn entities_len(&self, key: &T) -> usize {
        self.map
            .as_ref()
            .get(key)
            .map(|entities| entities.len())
            .unwrap_or(0)
    }
}

pub struct Indexed<T: Hash + Eq + Clone> {
    map: Arc<DashMap<T, DashSet<Entity>>>,
    entity: Entity,
    component: T,
}

impl<T: Hash + Eq + Clone> Indexed<T> {
    pub fn new(index: Index<T>, entity: Entity, component: T) -> Self {
        index
            .map
            .as_ref()
            .entry(component.clone())
            .or_insert_with(DashSet::new)
            .insert(entity);
        Self {
            map: index.map,
            entity,
            component,
        }
    }

    pub fn get(&self) -> &T {
        &self.component
    }

    pub fn set(&mut self, value: T) {
        self.map
            .as_ref()
            .get(&self.component)
            .unwrap()
            .remove(&self.entity);
        self.component = value;
        self.map
            .as_ref()
            .entry(self.component.clone())
            .or_insert_with(DashSet::new)
            .insert(self.entity);
    }

    pub fn unwrap(self) -> T {
        self.component.clone()
    }
}

impl<T: Hash + Eq + Clone> Drop for Indexed<T> {
    fn drop(&mut self) {
        self.map
            .as_ref()
            .get(&self.component)
            .unwrap()
            .remove(&self.entity);
    }
}

The Indexed is a wrapper for the component that is to be used as the index, i.e. the actual component associated with an entity is Indexed<MyPosition> instead of the naked MyPosition. Also, reverse mapping is not addressed because that's just random access into the world or a query with Indexed<MyPosition>, no need to duplicate that data.

I don't particularly like this snippet: creating a new indexed component is awkward; there are no safeguards from folks doing remove_one() and immediately attaching it to a different entity instead of unwrapping and re-wrapping it; inspecting entities at an index allocates a vector, although I'm pretty sure data copying to some degree is inevitable if we want to keep this degree of concurrency.

A path that might lend some improvements: give query::Fetch access to resources, like SystemParam. This will enable all kinds of wrappers/accessors for individual components, allowing things like Query<Indexed<&MyPosition>> instead of Query<&Indexed<MyPosition>> - the component is naked MyPosition again and Indexed only exists for as long as the query does, so no more awkward initialization. (Obvious caveat: users can access the component directly, desynchronizing it from the index.) Side note, this could also serve as a building block for a #68 solution.

Why am I still convinced against deferred re-indexing (aka housekeeping system): this allows data to exist in an invalid, desynchronized state. It's not unreasonable to imagine a system that makes changes to index-like components in its first half, then tries to do something with them in its latter half, and ends up creating bugs because the changes made in the first half are not observed by the latter one.
Perhaps a middle of the road approach is the way to go: during random access, modify index and component at the same time. During query access, defer re-indexing until the end of query... somehow. Query sets will likely need some more thinking.

Note on "Steps to Implement": as I've said, you don't need to touch the macro; unlike I claimed, you don't need to touch SystemParam either - Index<T> will be just fine being a plain resource.

I should have probably split this up into several posts...

[Ratysz]

Interesting, that's surprising to me. My feeling was that with a plain resource, you risk other users altering the index in unintended ways.

How would a special-cased system parameter be different? If it looks like a resource, works like a resource, and is used like a resource, it's probably just a resource, no need to reinvent the wheel; and encapsulation should be handled by Rust's own mechanisms, pub(crate) being the lazy but effective example.

Could we expose a .rebuild method on Index to allow for manual updates mid-system? Not as nice as it working completely magically, but would be pretty simple to use in practice.

That should work. I would still prefer it if doing the wrong thing was either impossible or heavily telegraphed in the API, but this is starting to look like the lesser evil.

I think the minimalist take would be a good first implementation: Index<T> is a resource, it is considered stale as soon as a T is changed, it can be refreshed manually mid-system through a mutable reference, and it comes with a refresh system that can be put wherever by the user; refreshing works via Changed<T>, no fancy automation beyond that.

[alice-i-cecile]

I agree with that as a solid goal for a first implementation; I have a prototype that I could get to that stage quite quickly.

Is there a way to build in a bit more safety? It won't be immediately clear that without manually refreshing before use you're likely to get stale values.

Definitely worth turning into a draft PR to examine, but I'm not sure I'd be comfortable shipping it to end users without some stronger guarantees on safe behavior.

[Ratysz]

Track index generation, perhaps, like the archetype generation. Somehow. But good docs will have to be sufficient for a prototype.

As I've said, this could be an external crate just fine, you don't have to PR that right away.

[cart]

Yeah a working external crate would be a good way to battle test an implementation first.

[alice-i-cecile]

Deal: I'll get bevy_index live and see how far I can push it without needing access to any internals :)

[alice-i-cecile]

With the addition of specialized component types in the upcoming ECS changes, we may be able to selectively impl DerefMut to force index updates, ensuring that the index is never stale.

Getting the details right on this would be tricky but it would be a very elegant and perfectly foolproof solution, which cannot be said of any of the other proposals.

[alice-i-cecile]

There was some additional discussion about indexes on Discord. This solution is basically the "pluggable storage" model. As discussed there, it couples storage and search in an uncomfortable manner, and probably isn't correct.

[dmurph]

The solution we used for 'old' index entries in IndexedDB in chromium (built on leveldb) was use just allow the 'old' values to be tombstones in the index, and clean them periodically & during cursor iterator. Each index entry would store the 'version' of the value, and every time the value was written that version would increment, so it was easy to tell if an index entry was a 'tombstone' when iterating the index.

I cannot, unfortunately, vouch for how 'better' that solution is than others. But it is a solution. The tricky part is figuring out how to do your tombstone sweeping (essentially garbage collection) at good times. I believe, assuming components are stored in flat arrays, this would be cache friendly, as you would load the index vec (containing the index key, entity, version) and the 'version' field for the entities, and that's all. So two blocks of memory? I'm not an expert in this particular area 😅

Or - perhaps you can guarantee that the index will be iterated frequently enough that tombstones are cleaned up anyways. In which case, this solution might be really good.

Edit: I didn't think about detecting when to change the 'version' of the entity, but it looks like this maybe be a universal problem / maybe solved by detecting 'mut' grabs.

[alice-i-cecile]

Yep, the change detection solution we already have will handle that.

[dmurph]

There is some chat on discord about this.

I guess there are some core questions about use cases. I'm guessing folks will want sorted iteration as well as efficient lookup, which in my experience means something like a b+tree / sorted array. However, insertions can be an issue there sometimes, and if you know the points querying occurs, then you can amortize out some sorting too if you do a custom b+tree (e.g. you know when lookups / reading will happen, so you can sort the blocks then - but maybe rusts's btreemap already does that?).

So - if you're doing that, then tombstones allow you to be append-only for modifications, and amortize out when you need to remove a bunch of entries (which then would definitely call for a custom b+tree implementation).

If you don't need iteration, then solutions can be simpler / hash-map.

[alice-i-cecile]

Follow-up from the future: Bevy 0.14 has observers, which work great for maintaining indexes if you forbid direct mutation.