Skip to content

feat: optimize access map #395

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 29, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 49 additions & 88 deletions crates/bevy_mod_scripting_core/src/bindings/access_map.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
//! A map of access claims used to safely and dynamically access the world.
use std::thread::ThreadId;

use bevy::{
ecs::{component::ComponentId, world::unsafe_world_cell::UnsafeWorldCell},
Expand All @@ -16,7 +15,6 @@ use super::{ReflectAllocationId, ReflectBase};
#[derive(Debug, Clone, PartialEq, Eq)]
/// An owner of an access claim and the code location of the claim.
pub struct ClaimOwner {
id: ThreadId,
location: std::panic::Location<'static>,
}

Expand All @@ -35,6 +33,7 @@ impl Default for AccessCount {
}
}

#[profiling::all_functions]
impl AccessCount {
fn new() -> Self {
Self {
Expand Down Expand Up @@ -71,6 +70,7 @@ pub trait AccessMapKey {
fn from_index(value: u64) -> Self;
}

#[profiling::all_functions]
impl AccessMapKey for u64 {
fn as_index(&self) -> u64 {
*self
Expand Down Expand Up @@ -100,6 +100,7 @@ pub struct ReflectAccessId {
pub(crate) id: u64,
}

#[profiling::all_functions]
impl AccessMapKey for ReflectAccessId {
fn as_index(&self) -> u64 {
// project two linear non-negative ranges [0,inf] to a single linear non-negative range, offset by 1 to avoid 0
Expand Down Expand Up @@ -134,6 +135,7 @@ impl AccessMapKey for ReflectAccessId {
}
}

#[profiling::all_functions]
impl ReflectAccessId {
/// Creates a new access id for the global world
pub fn for_global() -> Self {
Expand Down Expand Up @@ -192,6 +194,7 @@ impl ReflectAccessId {
}
}

#[profiling::all_functions]
impl From<ComponentId> for ReflectAccessId {
fn from(value: ComponentId) -> Self {
Self {
Expand All @@ -201,6 +204,7 @@ impl From<ComponentId> for ReflectAccessId {
}
}

#[profiling::all_functions]
impl From<ReflectAllocationId> for ReflectAccessId {
fn from(value: ReflectAllocationId) -> Self {
Self {
Expand All @@ -210,12 +214,14 @@ impl From<ReflectAllocationId> for ReflectAccessId {
}
}

#[profiling::all_functions]
impl From<ReflectAccessId> for ComponentId {
fn from(val: ReflectAccessId) -> Self {
ComponentId::new(val.id as usize)
}
}

#[profiling::all_functions]
impl From<ReflectAccessId> for ReflectAllocationId {
fn from(val: ReflectAccessId) -> Self {
ReflectAllocationId::new(val.id)
Expand Down Expand Up @@ -315,6 +321,29 @@ struct AccessMapInner {
global_lock: AccessCount,
}

#[profiling::all_functions]
impl AccessMapInner {
#[inline]
fn entry(&self, key: u64) -> Option<&AccessCount> {
self.individual_accesses.get(&key)
}

#[inline]
fn entry_mut(&mut self, key: u64) -> Option<&mut AccessCount> {
self.individual_accesses.get_mut(&key)
}

#[inline]
fn entry_or_default(&mut self, key: u64) -> &mut AccessCount {
self.individual_accesses.entry(key).or_default()
}

#[inline]
fn remove(&mut self, key: u64) {
self.individual_accesses.remove(&key);
}
}

const GLOBAL_KEY: u64 = 0;

#[profiling::all_functions]
Expand Down Expand Up @@ -362,10 +391,10 @@ impl DynamicSystemMeta for AccessMap {
return false;
}

let entry = inner.individual_accesses.entry(key).or_default();
let entry = inner.entry_or_default(key);

if entry.can_read() {
entry.read_by.push(ClaimOwner {
id: std::thread::current().id(),
location: *std::panic::Location::caller(),
});
true
Expand All @@ -388,10 +417,10 @@ impl DynamicSystemMeta for AccessMap {
return false;
}

let entry = inner.individual_accesses.entry(key).or_default();
let entry = inner.entry_or_default(key);

if entry.can_write() {
entry.read_by.push(ClaimOwner {
id: std::thread::current().id(),
location: *std::panic::Location::caller(),
});
entry.written = true;
Expand All @@ -409,7 +438,6 @@ impl DynamicSystemMeta for AccessMap {
return false;
}
inner.global_lock.read_by.push(ClaimOwner {
id: std::thread::current().id(),
location: *std::panic::Location::caller(),
});
inner.global_lock.written = true;
Expand All @@ -420,39 +448,27 @@ impl DynamicSystemMeta for AccessMap {
let mut inner = self.0.lock();
let key = key.as_index();

if let Some(entry) = inner.individual_accesses.get_mut(&key) {
if let Some(entry) = inner.entry_mut(key) {
entry.written = false;
if let Some(claim) = entry.read_by.pop() {
assert!(
claim.id == std::thread::current().id(),
"Access released from wrong thread, claimed at {}",
claim.location.display_location()
);
}
entry.read_by.pop();
if entry.readers() == 0 {
inner.individual_accesses.remove(&key);
inner.remove(key);
}
}
}

fn release_global_access(&self) {
let mut inner = self.0.lock();
inner.global_lock.read_by.pop();
inner.global_lock.written = false;
if let Some(claim) = inner.global_lock.read_by.pop() {
assert!(
claim.id == std::thread::current().id(),
"Global access released from wrong thread, claimed at {}",
claim.location.display_location()
);
}
}

fn list_accesses<K: AccessMapKey>(&self) -> Vec<(K, AccessCount)> {
let inner = self.0.lock();
inner
.individual_accesses
.iter()
.map(|(&key, count)| (K::from_index(key), count.clone()))
.map(|(key, a)| (K::from_index(*key), a.clone()))
.collect()
}

Expand Down Expand Up @@ -486,8 +502,7 @@ impl DynamicSystemMeta for AccessMap {
})
} else {
inner
.individual_accesses
.get(&key.as_index())
.entry(key.as_index())
.and_then(|access| access.as_location())
}
}
Expand All @@ -496,8 +511,9 @@ impl DynamicSystemMeta for AccessMap {
let inner = self.0.lock();
inner
.individual_accesses
.values()
.find_map(|access| access.as_location())
.iter()
.next()
.and_then(|(_, access)| access.as_location())
}
}

Expand All @@ -507,6 +523,7 @@ pub struct SubsetAccessMap {
subset: Box<dyn Fn(u64) -> bool + Send + Sync + 'static>,
}

#[profiling::all_functions]
impl SubsetAccessMap {
/// Creates a new subset access map with the provided subset of ID's as well as a exception function.
pub fn new(
Expand All @@ -528,6 +545,7 @@ impl SubsetAccessMap {
}
}

#[profiling::all_functions]
impl DynamicSystemMeta for SubsetAccessMap {
fn with_scope<O, F: FnOnce() -> O>(&self, f: F) -> O {
self.inner.with_scope(f)
Expand Down Expand Up @@ -601,6 +619,7 @@ pub enum AnyAccessMap {
SubsetAccessMap(SubsetAccessMap),
}

#[profiling::all_functions]
impl DynamicSystemMeta for AnyAccessMap {
fn with_scope<O, F: FnOnce() -> O>(&self, f: F) -> O {
match self {
Expand Down Expand Up @@ -700,12 +719,14 @@ pub trait DisplayCodeLocation {
fn display_location(self) -> String;
}

#[profiling::all_functions]
impl DisplayCodeLocation for std::panic::Location<'_> {
fn display_location(self) -> String {
format!("\"{}:{}\"", self.file(), self.line())
}
}

#[profiling::all_functions]
impl DisplayCodeLocation for Option<std::panic::Location<'_>> {
fn display_location(self) -> String {
self.map(|l| l.display_location())
Expand Down Expand Up @@ -926,66 +947,6 @@ mod test {
assert!(!subset_access_map.claim_global_access());
}

#[test]
#[should_panic]
fn access_map_releasing_read_access_from_wrong_thread_panics() {
let access_map = AccessMap::default();

access_map.claim_read_access(1);
std::thread::spawn(move || {
access_map.release_access(1);
})
.join()
.unwrap();
}

#[test]
#[should_panic]
fn subset_map_releasing_read_access_from_wrong_thread_panics() {
let access_map = AccessMap::default();
let subset_access_map = SubsetAccessMap {
inner: access_map,
subset: Box::new(|id| id == 1),
};

subset_access_map.claim_read_access(1);
std::thread::spawn(move || {
subset_access_map.release_access(1);
})
.join()
.unwrap();
}

#[test]
#[should_panic]
fn access_map_releasing_write_access_from_wrong_thread_panics() {
let access_map = AccessMap::default();

access_map.claim_write_access(1);
std::thread::spawn(move || {
access_map.release_access(1);
})
.join()
.unwrap();
}

#[test]
#[should_panic]
fn subset_map_releasing_write_access_from_wrong_thread_panics() {
let access_map = AccessMap::default();
let subset_access_map = SubsetAccessMap {
inner: access_map,
subset: Box::new(|id| id == 1),
};

subset_access_map.claim_write_access(1);
std::thread::spawn(move || {
subset_access_map.release_access(1);
})
.join()
.unwrap();
}

#[test]
fn as_and_from_index_for_access_id_non_overlapping() {
let global = ReflectAccessId::for_global();
Expand Down
Loading