WIP: disconnect spaces when their categories differ

crates/re_viewport/Cargo.toml

@@ -42,5 +42,6 @@ glam.workspace = true
 image = { workspace = true, default-features = false, features = ["png"] }
 itertools.workspace = true
 nohash-hasher.workspace = true
+once_cell.workspace = true
 serde.workspace = true
 tinyvec.workspace = true

crates/re_viewport/src/space_info.rs

@@ -1,12 +1,14 @@
+use itertools::Itertools;
 use nohash_hasher::IntSet;
 use re_space_view::UnreachableTransformReason;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 
 use re_arrow_store::{LatestAtQuery, TimeInt, Timeline};
 use re_data_store::{store_db::EntityDb, EntityPath, EntityTree};
 use re_types::{
-    archetypes::Pinhole,
+    archetypes::{Boxes2D, Boxes3D, Image, Pinhole, Points2D, Points3D, TextDocument, TextLog},
     components::{DisconnectedSpace, Transform3D},
+    Archetype, ComponentName,
 };
 
 use crate::query_pinhole;
@@ -23,6 +25,14 @@ pub enum SpaceInfoConnection {
     Disconnected,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Ord)]
+enum SpaceInfoCategory {
+    Spatial2D,
+    Spatial3D,
+    Other,
+    Empty,
+}
+
 /// Information about one "space".
 ///
 /// This is gathered by analyzing the transform hierarchy of the entities.
@@ -79,16 +89,19 @@ impl SpaceInfo {
                     continue;
                 };
 
-                // don't allow nested pinhole
-                let has_pinhole = matches!(
-                    connection,
+                let mut has_pinhole = false;
+
+                match connection {
                     SpaceInfoConnection::Connected {
-                        pinhole: Some(_),
-                        ..
+                        pinhole: Some(_), ..
+                    } => {
+                        has_pinhole = true;
+                        if encountered_pinhole {
+                            continue;
+                        }
                     }
-                );
-                if encountered_pinhole && has_pinhole {
-                    continue;
+                    SpaceInfoConnection::Connected { .. } => {}
+                    SpaceInfoConnection::Disconnected => continue,
                 }
 
                 visit_descendants_with_reachable_transform_recursively(
@@ -116,18 +129,102 @@ pub struct SpaceInfoCollection {
     spaces: BTreeMap<EntityPath, SpaceInfo>,
 }
 
+impl SpaceInfoCollection {
+    pub fn len(&self) -> usize {
+        self.spaces.len()
+    }
+}
+
+// TODO(jleibs): These need to get pulled from a registry of some form
+
+static INDICATORS_2D: once_cell::sync::Lazy<BTreeSet<ComponentName>> =
+    once_cell::sync::Lazy::new(|| {
+        [
+            Points2D::indicator().name(),
+            Boxes2D::indicator().name(),
+            Image::indicator().name(),
+        ]
+        .iter()
+        .cloned()
+        .collect()
+    });
+static INDICATORS_3D: once_cell::sync::Lazy<BTreeSet<ComponentName>> =
+    once_cell::sync::Lazy::new(|| {
+        [Points3D::indicator().name(), Boxes3D::indicator().name()]
+            .iter()
+            .cloned()
+            .collect()
+    });
+static INDICATORS_OTHER: once_cell::sync::Lazy<BTreeSet<ComponentName>> =
+    once_cell::sync::Lazy::new(|| {
+        [
+            TextLog::indicator().name(),
+            TextDocument::indicator().name(),
+        ]
+        .iter()
+        .cloned()
+        .collect()
+    });
+
+fn populate_categories(
+    categories: &mut BTreeMap<EntityPath, BTreeMap<SpaceInfoCategory, usize>>,
+    tree: &EntityTree,
+) -> BTreeMap<SpaceInfoCategory, usize> {
+    let key_set: BTreeSet<ComponentName> = tree.components.keys().cloned().collect();
+
+    let mut categories_for_subtree = BTreeMap::default();
+
+    for child in tree.children.values() {
+        for (cat, count) in populate_categories(categories, child) {
+            *categories_for_subtree.entry(cat).or_default() += count;
+        }
+    }
+
+    if !key_set.is_disjoint(&INDICATORS_2D) {
+        *categories_for_subtree
+            .entry(SpaceInfoCategory::Spatial2D)
+            .or_default() += 1;
+    }
+    if !key_set.is_disjoint(&INDICATORS_3D) {
+        *categories_for_subtree
+            .entry(SpaceInfoCategory::Spatial3D)
+            .or_default() += 1;
+    }
+    if !key_set.is_disjoint(&INDICATORS_OTHER) {
+        *categories_for_subtree
+            .entry(SpaceInfoCategory::Other)
+            .or_default() += 1;
+    }
+
+    categories.insert(tree.path.clone(), categories_for_subtree.clone());
+
+    categories_for_subtree
+}
+
+fn best_category(category_counts: &BTreeMap<SpaceInfoCategory, usize>) -> SpaceInfoCategory {
+    category_counts
+        .iter()
+        .sorted_by(|(_, v1), (_, v2)| v2.cmp(v1))
+        .next()
+        .map_or(SpaceInfoCategory::Empty, |(k, _)| *k)
+}
+
 impl SpaceInfoCollection {
     /// Do a graph analysis of the transform hierarchy, and create cuts
     /// wherever we find a non-identity transform.
     pub fn new(entity_db: &EntityDb) -> Self {
         re_tracing::profile_function!();
 
+        let mut categories = BTreeMap::default();
+        populate_categories(&mut categories, &entity_db.tree);
+
         fn add_children(
             entity_db: &EntityDb,
             spaces_info: &mut SpaceInfoCollection,
             parent_space: &mut SpaceInfo,
             tree: &EntityTree,
             query: &LatestAtQuery,
+            categories: &BTreeMap<EntityPath, BTreeMap<SpaceInfoCategory, usize>>,
         ) {
             // Determine how the paths are connected.
             let store = &entity_db.data_store;
@@ -146,6 +243,20 @@ impl SpaceInfoCollection {
                 .map_or(false, |dp| dp.0)
             {
                 Some(SpaceInfoConnection::Disconnected)
+            } else if let (Some(parent_categories), Some(space_categories)) = (
+                categories.get(&parent_space.path),
+                categories.get(&tree.path),
+            ) {
+                let best_parent_category = best_category(parent_categories);
+                let best_space_category = best_category(space_categories);
+                if best_parent_category != best_space_category
+                    || parent_categories.get(&best_parent_category)
+                        == space_categories.get(&best_space_category)
+                {
+                    Some(SpaceInfoConnection::Disconnected)
+                } else {
+                    None
+                }
             } else {
                 None
             };
@@ -169,6 +280,7 @@ impl SpaceInfoCollection {
                         &mut child_space_info,
                         child_tree,
                         query,
+                        categories,
                     );
                 }
                 spaces_info
@@ -181,7 +293,14 @@ impl SpaceInfoCollection {
                     .insert(tree.path.clone()); // spaces includes self
 
                 for child_tree in tree.children.values() {
-                    add_children(entity_db, spaces_info, parent_space, child_tree, query);
+                    add_children(
+                        entity_db,
+                        spaces_info,
+                        parent_space,
+                        child_tree,
+                        query,
+                        categories,
+                    );
                 }
             }
         }
@@ -208,6 +327,7 @@ impl SpaceInfoCollection {
             &mut root_space_info,
             &entity_db.tree,
             &query,
+            &categories,
         );
         spaces_info
             .spaces

crates/re_viewport/src/space_view_heuristics.rs

@@ -48,7 +48,9 @@ fn candidate_space_view_paths<'a>(
     spaces_info
         .iter()
         .map(|info| &info.path)
-        .chain(root_children.values().map(|sub_tree| &sub_tree.path))
+        // TODO(jleibs): These root children shouldn't need manual adding. If they are interesting
+        // they should have been generated by the SpaceInfos
+        //.chain(root_children.values().map(|sub_tree| &sub_tree.path))
         .unique()
 }
 
@@ -150,11 +152,15 @@ fn is_interesting_space_view_at_root(
     //
     // If there are any images directly under the root, don't create root space either.
     // -> For images we want more fine grained control and resort to child-of-root spaces only.
+
+    // TODO(jleibs): This gets in the way
+    /*
     for entity_path in &candidate.contents.root_group().entities {
         if contains_any_image(entity_path, data_store) {
             return false;
         }
     }
+    */
 
     true
 }
@@ -165,6 +171,9 @@ fn is_interesting_space_view_not_at_root(
     classes_with_interesting_roots: &[SpaceViewClassName],
     query: &LatestAtQuery,
 ) -> bool {
+    // TODO(jleibs): Why not make everything interesting?
+    return true;
+
     // TODO(andreas): Can we express this with [`AutoSpawnHeuristic`] instead?
 
     // Consider children of the root interesting, *unless* a root with the same category was already considered interesting!