More context menu 6: Add "Expand/Collapse all" actions (#5433)

### What

What the title says ☝🏻 


https://github.com/rerun-io/rerun/assets/49431240/a5dcd8c4-a05c-4159-a854-0ebec52a34a2


### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested the web demo (if applicable):
* Using newly built examples:
[app.rerun.io](https://app.rerun.io/pr/5433/index.html)
* Using examples from latest `main` build:
[app.rerun.io](https://app.rerun.io/pr/5433/index.html?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[app.rerun.io](https://app.rerun.io/pr/5433/index.html?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG
* [x] If applicable, add a new check to the [release
checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)!

- [PR Build Summary](https://build.rerun.io/pr/5433)
- [Docs
preview](https://rerun.io/preview/c276bcecb906512bfabf6eb41be63ba433731ff8/docs)
<!--DOCS-PREVIEW-->
- [Examples
preview](https://rerun.io/preview/c276bcecb906512bfabf6eb41be63ba433731ff8/examples)
<!--EXAMPLES-PREVIEW-->
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)

crates/re_viewport/src/context_menu/actions.rs

@@ -5,7 +5,7 @@ use nohash_hasher::IntSet;
 use re_entity_db::InstancePath;
 use re_log_types::{EntityPath, EntityPathFilter, EntityPathRule, RuleEffect};
 use re_space_view::{determine_visualizable_entities, SpaceViewBlueprint};
-use re_viewer_context::{ContainerId, Item, SpaceViewClassIdentifier, SpaceViewId};
+use re_viewer_context::{CollapseScope, ContainerId, Item, SpaceViewClassIdentifier, SpaceViewId};
 
 use super::{ContextMenuAction, ContextMenuContext};
 use crate::Contents;
@@ -544,3 +544,122 @@ fn create_space_view_for_selected_entities(
     ctx.viewport_blueprint
         .mark_user_interaction(ctx.viewer_context);
 }
+
+// ---
+
+/// Collapse or expand all items in the selection.
+// TODO(ab): the current implementation makes strong assumptions of which CollapseScope to use based
+// on the item type. This is brittle and will not scale if/when we add more trees to the UI. When
+// that happens, we will have to pass the scope to `context_menu_ui_for_item` and use it here.
+pub(super) enum CollapseExpandAllAction {
+    CollapseAll,
+    ExpandAll,
+}
+
+impl ContextMenuAction for CollapseExpandAllAction {
+    fn supports_selection(&self, ctx: &ContextMenuContext<'_>) -> bool {
+        // let's allow this action if at least one item supports it
+        ctx.selection
+            .iter()
+            .any(|(item, _)| self.supports_item(ctx, item))
+    }
+
+    fn supports_item(&self, ctx: &ContextMenuContext<'_>, item: &Item) -> bool {
+        // TODO(ab): in an ideal world, we'd check the fully expended/collapsed state of the item to
+        // avoid showing a command that wouldn't have an effect but that's lots of added complexity.
+        match item {
+            Item::StoreId(_) | Item::ComponentPath(_) => false,
+            Item::SpaceView(_) | Item::Container(_) | Item::InstancePath(_) => true,
+            //TODO(ab): for DataResult, walk the data result tree instead!
+            Item::DataResult(_, instance_path) => ctx
+                .viewer_context
+                .entity_db
+                .tree()
+                .subtree(&instance_path.entity_path)
+                .is_some_and(|subtree| !subtree.is_leaf()),
+        }
+    }
+
+    fn label(&self, _ctx: &ContextMenuContext<'_>) -> String {
+        match self {
+            CollapseExpandAllAction::CollapseAll => "Collapse all".to_owned(),
+            CollapseExpandAllAction::ExpandAll => "Expand all".to_owned(),
+        }
+    }
+
+    fn process_container(&self, ctx: &ContextMenuContext<'_>, container_id: &ContainerId) {
+        ctx.viewport_blueprint
+            .visit_contents_in_container(container_id, &mut |contents| match contents {
+                Contents::Container(container_id) => CollapseScope::BlueprintTree
+                    .container(*container_id)
+                    .set_open(&ctx.egui_context, self.open()),
+                Contents::SpaceView(space_view_id) => self.process_space_view(ctx, space_view_id),
+            });
+    }
+
+    fn process_space_view(&self, ctx: &ContextMenuContext<'_>, space_view_id: &SpaceViewId) {
+        CollapseScope::BlueprintTree
+            .space_view(*space_view_id)
+            .set_open(&ctx.egui_context, self.open());
+
+        let query_result = ctx.viewer_context.lookup_query_result(*space_view_id);
+        let result_tree = &query_result.tree;
+        if let Some(root_node) = result_tree.root_node() {
+            self.process_data_result(
+                ctx,
+                space_view_id,
+                &InstancePath::entity_splat(root_node.data_result.entity_path.clone()),
+            );
+        }
+    }
+
+    fn process_data_result(
+        &self,
+        ctx: &ContextMenuContext<'_>,
+        space_view_id: &SpaceViewId,
+        instance_path: &InstancePath,
+    ) {
+        //TODO(ab): here we should in principle walk the DataResult tree instead of the entity tree
+        // but the current API isn't super ergonomic.
+        let Some(subtree) = ctx
+            .viewer_context
+            .entity_db
+            .tree()
+            .subtree(&instance_path.entity_path)
+        else {
+            return;
+        };
+
+        subtree.visit_children_recursively(&mut |entity_path, _| {
+            CollapseScope::BlueprintTree
+                .data_result(*space_view_id, entity_path.clone())
+                .set_open(&ctx.egui_context, self.open());
+        });
+    }
+
+    fn process_instance_path(&self, ctx: &ContextMenuContext<'_>, instance_path: &InstancePath) {
+        let Some(subtree) = ctx
+            .viewer_context
+            .entity_db
+            .tree()
+            .subtree(&instance_path.entity_path)
+        else {
+            return;
+        };
+
+        subtree.visit_children_recursively(&mut |entity_path, _| {
+            CollapseScope::StreamsTree
+                .entity(entity_path.clone())
+                .set_open(&ctx.egui_context, self.open());
+        });
+    }
+}
+
+impl CollapseExpandAllAction {
+    fn open(&self) -> bool {
+        match self {
+            CollapseExpandAllAction::CollapseAll => false,
+            CollapseExpandAllAction::ExpandAll => true,
+        }
+    }
+}

crates/re_viewport/src/context_menu/mod.rs

@@ -11,7 +11,8 @@ mod sub_menu;
 
 use actions::{
     AddContainerAction, AddEntitiesToNewSpaceViewAction, AddSpaceViewAction, CloneSpaceViewAction,
-    HideAction, MoveContentsToNewContainerAction, RemoveAction, ShowAction,
+    CollapseExpandAllAction, HideAction, MoveContentsToNewContainerAction, RemoveAction,
+    ShowAction,
 };
 use sub_menu::SubMenu;
 
@@ -46,6 +47,7 @@ pub fn context_menu_ui_for_item(
             let context_menu_ctx = ContextMenuContext {
                 viewer_context: ctx,
                 viewport_blueprint,
+                egui_context: ui.ctx().clone(),
                 selection,
                 clicked_item: item,
             };
@@ -102,6 +104,10 @@ fn action_list(
                 Box::new(HideAction),
                 Box::new(RemoveAction),
             ],
+            vec![
+                Box::new(CollapseExpandAllAction::ExpandAll),
+                Box::new(CollapseExpandAllAction::CollapseAll),
+            ],
             vec![Box::new(CloneSpaceViewAction)],
             vec![
                 Box::new(SubMenu {
@@ -186,6 +192,7 @@ fn show_context_menu_for_selection(ctx: &ContextMenuContext<'_>, ui: &mut egui::
 struct ContextMenuContext<'a> {
     viewer_context: &'a ViewerContext<'a>,
     viewport_blueprint: &'a ViewportBlueprint,
+    egui_context: egui::Context,
     selection: &'a Selection,
     clicked_item: &'a Item,
 }

tests/python/release_checklist/check_context_menu_collapse_expand_all.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import os
+from argparse import Namespace
+from uuid import uuid4
+
+import rerun as rr
+
+README = """
+# Context Menu - Add entity to new space view
+
+## Blueprint tree
+
+* Reset the blueprint.
+* Right-click on Viewport and select "Collapse all". Check everything is collapsed by manually expending everything.
+* Right-click on Viewport and select "Collapse all" and then "Expend all". Check everything is expanded.
+
+## Streams tree
+
+* Same as above, with the `world/` entity.
+
+
+## Multi-selection
+
+* Same as above, with both the viewport (blueprint tree) and `world/` (streams tree) selected.
+"""
+
+
+def log_readme() -> None:
+    rr.log("readme", rr.TextDocument(README, media_type=rr.MediaType.MARKDOWN), timeless=True)
+
+
+def log_some_space_views() -> None:
+    # TODO(ab): add a deep-ish container hierarchy blueprint for more collapse/expand fun
+
+    rr.set_time_sequence("frame_nr", 0)
+
+    rr.log("/", rr.Boxes3D(centers=[0, 0, 0], half_sizes=[1, 1, 1]))
+    rr.log("/world/robot/arm/actuator/thing", rr.Boxes3D(centers=[0.5, 0, 0], half_sizes=[0.1, 0.1, 0.1]))
+
+
+def run(args: Namespace) -> None:
+    # TODO(cmc): I have no idea why this works without specifying a `recording_id`, but
+    # I'm not gonna rely on it anyway.
+    rr.script_setup(args, f"{os.path.basename(__file__)}", recording_id=uuid4())
+
+    log_readme()
+    log_some_space_views()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Interactive release checklist")
+    rr.script_add_args(parser)
+    args = parser.parse_args()
+    run(args)