use std::sync::OnceLock;
use ahash::HashMap;
use nohash_hasher::{IntMap, IntSet};
use re_chunk_store::{
ChunkStore, ChunkStoreDiff, ChunkStoreEvent, ChunkStoreSubscriber, ChunkStoreSubscriberHandle,
};
use re_log::debug_assert;
use re_log_types::{EntityPath, EntityPathHash, StoreId};
use re_sdk_types::Component as _;
use re_sdk_types::components::{PinholeProjection, ViewCoordinates};
bitflags::bitflags! {
#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub struct SubSpaceConnectionFlags: u8 {
const Pinhole = 0b0000001;
}
}
bitflags::bitflags! {
#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub struct HeuristicHints: u8 {
const ViewCoordinates3d = 0b0000001;
}
}
#[derive(Debug)]
pub struct SubSpace {
pub origin: EntityPath,
pub entities: IntSet<EntityPath>,
pub child_spaces: IntSet<EntityPath>,
pub parent_space: EntityPathHash,
pub connection_to_parent: SubSpaceConnectionFlags,
pub heuristic_hints: IntMap<EntityPath, HeuristicHints>,
}
impl SubSpace {
#[inline]
pub fn supports_3d_content(&self) -> bool {
!self
.connection_to_parent
.contains(SubSpaceConnectionFlags::Pinhole)
}
#[inline]
#[expect(clippy::unused_self)]
pub fn supports_2d_content(&self) -> bool {
true
}
}
#[derive(Default)]
pub struct SpatialTopologyStoreSubscriber {
topologies: HashMap<StoreId, SpatialTopology>,
}
impl SpatialTopologyStoreSubscriber {
pub fn subscription_handle() -> ChunkStoreSubscriberHandle {
static SUBSCRIPTION: OnceLock<ChunkStoreSubscriberHandle> = OnceLock::new();
*SUBSCRIPTION.get_or_init(|| ChunkStore::register_subscriber(Box::<Self>::default()))
}
}
impl re_byte_size::MemUsageTreeCapture for SpatialTopologyStoreSubscriber {
fn capture_mem_usage_tree(&self) -> re_byte_size::MemUsageTree {
use re_byte_size::SizeBytes as _;
let mut node = re_byte_size::MemUsageNode::new();
for (store_id, topology) in &self.topologies {
let name = format!(
"{}/{}",
store_id.application_id().as_str(),
store_id.recording_id().as_str()
);
node.add(name, topology.total_size_bytes());
}
node.into_tree()
}
}
impl ChunkStoreSubscriber for SpatialTopologyStoreSubscriber {
#[inline]
fn name(&self) -> String {
"SpatialTopologyStoreSubscriber".to_owned()
}
#[inline]
fn as_any(&self) -> &dyn std::any::Any {
self
}
#[inline]
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
fn on_events(&mut self, events: &[ChunkStoreEvent]) {
re_tracing::profile_function!();
for event in events {
let ChunkStoreDiff::SchemaAddition(add) = &event.diff else {
continue;
};
for meta in &add.new_columns {
self.topologies
.entry(event.store_id.clone())
.or_default()
.on_store_diff(
&meta.entity_path,
meta.components
.iter()
.map(|component| &component.descriptor),
);
}
}
}
}
#[derive(Debug)]
pub struct SpatialTopology {
subspaces: IntMap<EntityPathHash, SubSpace>,
subspace_origin_per_logged_entity: IntMap<EntityPathHash, EntityPathHash>,
has_explicit_coordinate_frame: bool,
}
impl re_byte_size::SizeBytes for SubSpaceConnectionFlags {
#[inline]
fn heap_size_bytes(&self) -> u64 {
0
}
#[inline]
fn is_pod() -> bool {
true
}
}
impl re_byte_size::SizeBytes for HeuristicHints {
#[inline]
fn heap_size_bytes(&self) -> u64 {
0
}
#[inline]
fn is_pod() -> bool {
true
}
}
impl re_byte_size::SizeBytes for SubSpace {
fn heap_size_bytes(&self) -> u64 {
let Self {
origin,
entities,
child_spaces,
parent_space: _,
connection_to_parent: _,
heuristic_hints,
} = self;
origin.heap_size_bytes()
+ entities.heap_size_bytes()
+ child_spaces.heap_size_bytes()
+ heuristic_hints.heap_size_bytes()
}
}
impl re_byte_size::SizeBytes for SpatialTopology {
fn heap_size_bytes(&self) -> u64 {
let Self {
subspaces,
subspace_origin_per_logged_entity,
has_explicit_coordinate_frame: _,
} = self;
subspaces.heap_size_bytes() + subspace_origin_per_logged_entity.heap_size_bytes()
}
}
impl Default for SpatialTopology {
fn default() -> Self {
Self {
subspaces: std::iter::once((
EntityPath::root().hash(),
SubSpace {
origin: EntityPath::root(),
entities: IntSet::default(), child_spaces: IntSet::default(),
parent_space: EntityPathHash::NONE,
connection_to_parent: SubSpaceConnectionFlags::empty(),
heuristic_hints: IntMap::default(),
},
))
.collect(),
subspace_origin_per_logged_entity: Default::default(),
has_explicit_coordinate_frame: false,
}
}
}
impl SpatialTopology {
pub fn access<T>(store_id: &StoreId, f: impl FnOnce(&Self) -> T) -> Option<T> {
ChunkStore::with_subscriber_once(
SpatialTopologyStoreSubscriber::subscription_handle(),
move |topology_subscriber: &SpatialTopologyStoreSubscriber| {
topology_subscriber.topologies.get(store_id).map(f)
},
)
.flatten()
}
#[inline]
pub fn subspace_for_entity(&self, entity: &EntityPath) -> &SubSpace {
self.subspaces
.get(&self.subspace_origin_hash_for_entity(entity))
.expect("unknown subspace origin, `SpatialTopology` is in an invalid state")
}
#[inline]
pub fn iter_subspaces(&self) -> impl Iterator<Item = &SubSpace> {
self.subspaces.values()
}
fn subspace_origin_hash_for_entity(&self, entity: &EntityPath) -> EntityPathHash {
let mut entity_reference = entity;
let mut entity_storage: EntityPath;
loop {
if let Some(origin_hash) = self
.subspace_origin_per_logged_entity
.get(&entity_reference.hash())
{
return *origin_hash;
}
if let Some(parent) = entity_reference.parent() {
entity_storage = parent;
entity_reference = &entity_storage;
} else {
return EntityPath::root().hash();
}
}
}
#[inline]
pub fn subspace_for_subspace_origin(&self, origin: EntityPathHash) -> Option<&SubSpace> {
self.subspaces.get(&origin)
}
fn on_store_diff<'a>(
&mut self,
entity_path: &EntityPath,
added_components: impl Iterator<Item = &'a re_sdk_types::ComponentDescriptor>,
) {
re_tracing::profile_function!();
let mut new_subspace_connections = SubSpaceConnectionFlags::empty();
let mut new_heuristic_hints = HeuristicHints::empty();
for added_component in added_components.filter_map(|descr| descr.component_type) {
if added_component == PinholeProjection::name() {
new_subspace_connections.insert(SubSpaceConnectionFlags::Pinhole);
} else if added_component == ViewCoordinates::name() {
new_heuristic_hints.insert(HeuristicHints::ViewCoordinates3d);
} else if added_component == re_tf::TransformFrameId::name() {
self.has_explicit_coordinate_frame = true;
}
}
if let Some(subspace_origin_hash) = self
.subspace_origin_per_logged_entity
.get(&entity_path.hash())
{
if !new_subspace_connections.is_empty() {
self.update_space_with_new_connections(
entity_path,
*subspace_origin_hash,
new_subspace_connections,
);
}
} else {
self.add_new_entity(entity_path, new_subspace_connections);
}
if !new_heuristic_hints.is_empty() {
let subspace = self
.subspaces
.get_mut(&self.subspace_origin_hash_for_entity(entity_path))
.expect("unknown subspace origin, `SpatialTopology` is in an invalid state");
subspace
.heuristic_hints
.entry(entity_path.clone())
.or_insert(HeuristicHints::empty())
.insert(new_heuristic_hints);
}
}
fn update_space_with_new_connections(
&mut self,
entity_path: &EntityPath,
subspace_origin_hash: EntityPathHash,
new_connections: SubSpaceConnectionFlags,
) {
if subspace_origin_hash == entity_path.hash() {
let subspace = self
.subspaces
.get_mut(&subspace_origin_hash)
.expect("Subspace origin not part of origin->subspace map.");
subspace.connection_to_parent.insert(new_connections);
} else {
self.split_subspace(subspace_origin_hash, entity_path, new_connections);
}
}
fn add_new_entity(
&mut self,
entity_path: &EntityPath,
subspace_connections: SubSpaceConnectionFlags,
) {
let subspace_origin_hash = self.subspace_origin_hash_for_entity(entity_path);
let target_space_origin_hash =
if subspace_connections.is_empty() || entity_path.hash() == subspace_origin_hash {
let subspace = self
.subspaces
.get_mut(&subspace_origin_hash)
.expect("Subspace origin not part of origin->subspace map.");
subspace.entities.insert(entity_path.clone());
subspace.connection_to_parent.insert(subspace_connections);
subspace.origin.hash()
} else {
self.split_subspace(subspace_origin_hash, entity_path, subspace_connections);
entity_path.hash()
};
self.subspace_origin_per_logged_entity
.insert(entity_path.hash(), target_space_origin_hash);
}
fn split_subspace(
&mut self,
split_subspace_origin_hash: EntityPathHash,
new_space_origin: &EntityPath,
connection_to_parent: SubSpaceConnectionFlags,
) {
let split_subspace = self
.subspaces
.get_mut(&split_subspace_origin_hash)
.expect("Subspace origin not part of origin->subspace map.");
debug_assert!(new_space_origin.is_descendant_of(&split_subspace.origin));
let mut new_space = SubSpace {
origin: new_space_origin.clone(),
entities: std::iter::once(new_space_origin.clone()).collect(),
child_spaces: Default::default(),
parent_space: split_subspace_origin_hash,
connection_to_parent,
heuristic_hints: Default::default(),
};
split_subspace.entities.retain(|e| {
if e.starts_with(new_space_origin) {
self.subspace_origin_per_logged_entity
.insert(e.hash(), new_space.origin.hash());
new_space.entities.insert(e.clone());
false
} else {
true
}
});
split_subspace.child_spaces.retain(|child_origin| {
debug_assert!(child_origin != new_space_origin);
if child_origin.is_descendant_of(new_space_origin) {
new_space.child_spaces.insert(child_origin.clone());
false
} else {
true
}
});
split_subspace.child_spaces.insert(new_space_origin.clone());
for child_origin in &new_space.child_spaces {
let child_space = self
.subspaces
.get_mut(&child_origin.hash())
.expect("Child origin not part of origin->subspace map.");
child_space.parent_space = new_space.origin.hash();
}
self.subspaces.insert(new_space.origin.hash(), new_space);
}
pub fn has_explicit_coordinate_frame(&self) -> bool {
self.has_explicit_coordinate_frame
}
}
#[cfg(test)]
mod tests {
use re_log_types::EntityPath;
use re_sdk_types::components::{PinholeProjection, ViewCoordinates};
use re_sdk_types::{Component as _, ComponentDescriptor};
use super::SpatialTopology;
use crate::spatial_topology::{HeuristicHints, SubSpaceConnectionFlags};
#[test]
fn no_splits() {
let mut topo = SpatialTopology::default();
assert_eq!(topo.subspaces.len(), 1);
assert_eq!(topo.subspace_origin_per_logged_entity.len(), 0);
add_diff(&mut topo, "robo", &[]);
add_diff(&mut topo, "robo/arm", &[]);
add_diff(&mut topo, "robo/eyes/cam", &[]);
check_paths_in_space(&topo, &["robo", "robo/arm", "robo/eyes/cam"], "/");
let subspace = topo.subspace_for_entity(&"robo".into());
assert!(subspace.child_spaces.is_empty());
assert!(subspace.parent_space.is_none());
assert_eq!(
topo.subspace_for_entity(&EntityPath::root()).origin,
EntityPath::root()
);
assert_eq!(
topo.subspace_for_entity(&"robo/eyes".into()).origin,
EntityPath::root()
);
assert_eq!(
topo.subspace_for_entity(&"robo/leg".into()).origin,
EntityPath::root()
);
#[expect(clippy::single_element_loop)]
for (name, flags) in [
(PinholeProjection::name(), SubSpaceConnectionFlags::Pinhole),
] {
add_diff(
&mut topo,
"",
&[ComponentDescriptor::partial("whatever").with_component_type(name)],
);
let subspace = topo.subspace_for_entity(&"robo".into());
assert_eq!(subspace.connection_to_parent, flags);
assert!(subspace.child_spaces.is_empty());
assert!(subspace.parent_space.is_none());
}
}
#[test]
fn valid_splits() {
let mut topo = SpatialTopology::default();
add_diff(&mut topo, "robo", &[]);
add_diff(&mut topo, "robo/eyes/left/cam/annotation", &[]);
add_diff(&mut topo, "robo/arm", &[]);
add_diff(
&mut topo,
"robo/eyes/left/cam",
&[ComponentDescriptor::partial("whatever")
.with_component_type(PinholeProjection::name())],
);
add_diff(&mut topo, "robo/eyes/right/cam/annotation", &[]);
add_diff(&mut topo, "robo/eyes/right/cam", &[]);
{
check_paths_in_space(
&topo,
&[
"robo",
"robo/arm",
"robo/eyes/right/cam",
"robo/eyes/right/cam/annotation",
],
"/",
);
check_paths_in_space(
&topo,
&["robo/eyes/left/cam", "robo/eyes/left/cam/annotation"],
"robo/eyes/left/cam",
);
let root = topo.subspace_for_entity(&"robo".into());
let left_camera = topo.subspace_for_entity(&"robo/eyes/left/cam".into());
assert_eq!(left_camera.origin, "robo/eyes/left/cam".into());
assert_eq!(left_camera.parent_space, root.origin.hash());
assert_eq!(
left_camera.connection_to_parent,
SubSpaceConnectionFlags::Pinhole
);
assert_eq!(root.connection_to_parent, SubSpaceConnectionFlags::empty());
assert!(root.parent_space.is_none());
}
add_diff(
&mut topo,
"robo/eyes/right/cam",
&[ComponentDescriptor::partial("whatever")
.with_component_type(PinholeProjection::name())],
);
{
check_paths_in_space(&topo, &["robo", "robo/arm"], "/");
check_paths_in_space(
&topo,
&["robo/eyes/right/cam", "robo/eyes/right/cam/annotation"],
"robo/eyes/right/cam",
);
let root = topo.subspace_for_entity(&"robo".into());
let left_camera = topo.subspace_for_entity(&"robo/eyes/left/cam".into());
let right_camera = topo.subspace_for_entity(&"robo/eyes/right/cam".into());
assert_eq!(right_camera.origin, "robo/eyes/right/cam".into());
assert_eq!(right_camera.parent_space, root.origin.hash());
assert_eq!(
right_camera.connection_to_parent,
SubSpaceConnectionFlags::Pinhole
);
assert_eq!(left_camera.origin, "robo/eyes/left/cam".into());
assert_eq!(left_camera.parent_space, root.origin.hash());
assert_eq!(
left_camera.connection_to_parent,
SubSpaceConnectionFlags::Pinhole
);
assert_eq!(root.connection_to_parent, SubSpaceConnectionFlags::empty());
assert!(root.parent_space.is_none());
assert_eq!(
topo.subspace_for_entity(&"robo/eyes/right/cam/unheard".into())
.origin,
"robo/eyes/right/cam".into()
);
assert_eq!(
topo.subspace_for_entity(&"bonkers".into()).origin,
EntityPath::root()
);
}
add_diff(
&mut topo,
"robo",
&[ComponentDescriptor::partial("whatever")
.with_component_type(ViewCoordinates::name())],
);
{
let root = topo.subspace_for_entity(&EntityPath::root());
assert!(root.parent_space.is_none());
assert_eq!(root.connection_to_parent, SubSpaceConnectionFlags::empty());
assert_eq!(
root.heuristic_hints,
std::iter::once((EntityPath::from("robo"), HeuristicHints::ViewCoordinates3d))
.collect()
);
}
}
#[test]
fn handle_invalid_splits_gracefully() {
for nested_first in [false, true] {
let mut topo = SpatialTopology::default();
if nested_first {
add_diff(
&mut topo,
"cam0/cam1",
&[ComponentDescriptor::partial("whatever1")
.with_component_type(PinholeProjection::name())],
);
add_diff(
&mut topo,
"cam0",
&[ComponentDescriptor::partial("whatever2")
.with_component_type(PinholeProjection::name())],
);
} else {
add_diff(
&mut topo,
"cam0",
&[ComponentDescriptor::partial("whatever3")
.with_component_type(PinholeProjection::name())],
);
add_diff(
&mut topo,
"cam0/cam1",
&[ComponentDescriptor::partial("whatever4")
.with_component_type(PinholeProjection::name())],
);
}
check_paths_in_space(&topo, &["cam0"], "cam0");
check_paths_in_space(&topo, &["cam0/cam1"], "cam0/cam1");
let root = topo.subspace_for_entity(&EntityPath::root());
let cam0 = topo.subspace_for_entity(&"cam0".into());
let cam1 = topo.subspace_for_entity(&"cam0/cam1".into());
assert_eq!(root.connection_to_parent, SubSpaceConnectionFlags::empty());
assert_eq!(cam0.connection_to_parent, SubSpaceConnectionFlags::Pinhole);
assert_eq!(cam1.connection_to_parent, SubSpaceConnectionFlags::Pinhole);
assert_eq!(cam0.parent_space, EntityPath::root().hash());
assert_eq!(cam1.parent_space, cam0.origin.hash());
assert!(cam1.child_spaces.is_empty());
}
}
fn add_diff(topo: &mut SpatialTopology, path: &str, components: &[ComponentDescriptor]) {
topo.on_store_diff(&path.into(), components.iter());
}
fn check_paths_in_space(topo: &SpatialTopology, paths: &[&str], expected_origin: &str) {
for path in paths {
let path = *path;
assert_eq!(
topo.subspace_for_entity(&path.into()).origin,
expected_origin.into()
);
}
let space = topo.subspace_for_entity(&paths[0].into());
for path in paths {
let path = *path;
assert!(space.entities.contains(&path.into()));
}
assert_eq!(space.entities.len(), paths.len());
}
}