Move StoreHub out of the Viewer during Update

crates/re_viewer/src/app_state.rs

@@ -1,15 +1,15 @@
 use ahash::HashMap;
 
 use re_data_store::StoreDb;
-use re_log_types::{ApplicationId, LogMsg, StoreId, TimeRangeF};
+use re_log_types::{LogMsg, StoreId, TimeRangeF};
 use re_smart_channel::Receiver;
 use re_viewer_context::{
     AppOptions, Caches, ComponentUiRegistry, PlayState, RecordingConfig, SpaceViewClassRegistry,
-    ViewerContext,
+    StoreContext, ViewerContext,
 };
 use re_viewport::ViewportState;
 
-use crate::ui::Blueprint;
+use crate::{store_hub::StoreHub, ui::Blueprint};
 
 const WATERMARK: bool = false; // Nice for recording media material
 
@@ -23,11 +23,6 @@ pub struct AppState {
     #[serde(skip)]
     pub(crate) cache: Caches,
 
-    #[serde(skip)]
-    selected_rec_id: Option<StoreId>,
-    #[serde(skip)]
-    pub(crate) selected_blueprint_by_app: HashMap<ApplicationId, StoreId>,
-
     /// Configuration for the current recording (found in [`StoreDb`]).
     recording_configs: HashMap<StoreId, RecordingConfig>,
 
@@ -41,16 +36,6 @@ pub struct AppState {
 }
 
 impl AppState {
-    /// The selected/visible recording id, if any.
-    pub fn recording_id(&self) -> Option<StoreId> {
-        self.selected_rec_id.clone()
-    }
-
-    /// The selected/visible recording id, if any.
-    pub fn set_recording_id(&mut self, recording_id: StoreId) {
-        self.selected_rec_id = Some(recording_id);
-    }
-
     pub fn app_options(&self) -> &AppOptions {
         &self.app_options
     }
@@ -61,18 +46,25 @@ impl AppState {
 
     /// Currently selected section of time, if any.
     #[cfg_attr(target_arch = "wasm32", allow(dead_code))]
-    pub fn loop_selection(&self) -> Option<(re_data_store::Timeline, TimeRangeF)> {
-        self.selected_rec_id.as_ref().and_then(|rec_id| {
-            self.recording_configs
-                .get(rec_id)
-                // is there an active loop selection?
-                .and_then(|rec_cfg| {
-                    rec_cfg
-                        .time_ctrl
-                        .loop_selection()
-                        .map(|q| (*rec_cfg.time_ctrl.timeline(), q))
-                })
-        })
+    pub fn loop_selection(
+        &self,
+        store_context: Option<&StoreContext<'_>>,
+    ) -> Option<(re_data_store::Timeline, TimeRangeF)> {
+        store_context
+            .as_ref()
+            .and_then(|ctx| ctx.recording)
+            .map(|rec| rec.store_id())
+            .and_then(|rec_id| {
+                self.recording_configs
+                    .get(rec_id)
+                    // is there an active loop selection?
+                    .and_then(|rec_cfg| {
+                        rec_cfg
+                            .time_ctrl
+                            .loop_selection()
+                            .map(|q| (*rec_cfg.time_ctrl.timeline(), q))
+                    })
+            })
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -82,6 +74,7 @@ impl AppState {
         ui: &mut egui::Ui,
         render_ctx: &mut re_renderer::RenderContext,
         store_db: &StoreDb,
+        store_context: &StoreContext<'_>,
         re_ui: &re_ui::ReUi,
         component_ui_registry: &ComponentUiRegistry,
         space_view_class_registry: &SpaceViewClassRegistry,
@@ -92,8 +85,6 @@ impl AppState {
         let Self {
             app_options,
             cache,
-            selected_rec_id: _,
-            selected_blueprint_by_app: _,
             recording_configs,
             selection_panel,
             time_panel,
@@ -113,6 +104,7 @@ impl AppState {
             space_view_class_registry,
             component_ui_registry,
             store_db,
+            store_context,
             rec_cfg,
             re_ui,
             render_ctx,
@@ -166,21 +158,9 @@ impl AppState {
         recording_config_entry(&mut self.recording_configs, id, data_source, store_db)
     }
 
-    pub fn cleanup(&mut self, store_hub: &crate::StoreHub) {
+    pub fn cleanup(&mut self, store_hub: &StoreHub) {
         re_tracing::profile_function!();
 
-        if !self
-            .selected_rec_id
-            .as_ref()
-            .map_or(false, |rec_id| store_hub.contains_recording(rec_id))
-        {
-            // Pick any:
-            self.selected_rec_id = store_hub
-                .recordings()
-                .next()
-                .map(|log| log.store_id().clone());
-        }
-
         self.recording_configs
             .retain(|store_id, _| store_hub.contains_recording(store_id));
     }

crates/re_viewer/src/lib.rs

@@ -24,7 +24,7 @@ pub(crate) use {
 
 pub use app::{App, StartupOptions};
 pub use remote_viewer_app::RemoteViewerApp;
-pub use store_hub::StoreHub;
+pub use store_hub::StoreBundle;
 
 pub mod external {
     pub use {eframe, egui};

crates/re_viewer/src/loading.rs

@@ -1,13 +1,13 @@
-use crate::StoreHub;
+use crate::StoreBundle;
 
 #[cfg(not(target_arch = "wasm32"))]
 #[must_use]
-pub fn load_file_path(path: &std::path::Path) -> Option<StoreHub> {
-    fn load_file_path_impl(path: &std::path::Path) -> anyhow::Result<StoreHub> {
+pub fn load_file_path(path: &std::path::Path) -> Option<StoreBundle> {
+    fn load_file_path_impl(path: &std::path::Path) -> anyhow::Result<StoreBundle> {
         re_tracing::profile_function!();
         use anyhow::Context as _;
         let file = std::fs::File::open(path).context("Failed to open file")?;
-        StoreHub::from_rrd(file)
+        StoreBundle::from_rrd(file)
     }
 
     re_log::info!("Loading {path:?}…");
@@ -35,8 +35,8 @@ pub fn load_file_path(path: &std::path::Path) -> Option<StoreHub> {
 }
 
 #[must_use]
-pub fn load_file_contents(name: &str, read: impl std::io::Read) -> Option<StoreHub> {
-    match StoreHub::from_rrd(read) {
+pub fn load_file_contents(name: &str, read: impl std::io::Read) -> Option<StoreBundle> {
+    match StoreBundle::from_rrd(read) {
         Ok(mut rrd) => {
             re_log::info!("Loaded {name:?}");
             for store_db in rrd.store_dbs_mut() {

crates/re_viewer/src/store_hub.rs

@@ -1,14 +1,195 @@
+use ahash::HashMap;
+use itertools::Itertools;
+use re_arrow_store::{DataStoreConfig, DataStoreStats};
 use re_data_store::StoreDb;
-use re_log_types::{StoreId, StoreKind};
+use re_log_types::{ApplicationId, StoreId, StoreKind};
+use re_viewer_context::StoreContext;
 
-/// Stores many [`StoreDb`]s of recordings and blueprints.
+/// Interface for accessing all blueprints and recordings
+///
+/// The [`StoreHub`] provides access to the [`StoreDb`] instances that are used
+/// to store both blueprints and recordings.
+///
+/// Internally, the [`StoreHub`] tracks which [`ApplicationId`] and `recording
+/// id` ([`StoreId`]) are currently selected in the viewer. These can be configured
+/// using [`StoreHub::set_recording_id`] and [`StoreHub::set_app_id`] respectively.
 #[derive(Default)]
 pub struct StoreHub {
+    selected_rec_id: Option<StoreId>,
+    application_id: Option<ApplicationId>,
+    blueprint_by_app_id: HashMap<ApplicationId, StoreId>,
+    store_dbs: StoreBundle,
+}
+
+/// Convenient information used for `MemoryPanel`
+#[derive(Default)]
+pub struct StoreHubStats {
+    pub blueprint_stats: DataStoreStats,
+    pub blueprint_config: DataStoreConfig,
+    pub recording_stats: DataStoreStats,
+    pub recording_config: DataStoreConfig,
+}
+
+impl StoreHub {
+    /// Add a [`StoreBundle`] to the [`StoreHub`]
+    pub fn add_bundle(&mut self, bundle: StoreBundle) {
+        self.store_dbs.append(bundle);
+    }
+
+    /// Get a read-only [`StoreContext`] from the [`StoreHub`] if one is available.
+    ///
+    /// All of the returned references to blueprints and recordings will have a
+    /// matching [`ApplicationId`].
+    pub fn read_context(&mut self) -> Option<StoreContext<'_>> {
+        // If we have an app-id, then use it to look up the blueprint.
+        let blueprint_id = self.application_id.as_ref().map(|app_id| {
+            self.blueprint_by_app_id
+                .entry(app_id.clone())
+                .or_insert_with(|| StoreId::from_string(StoreKind::Blueprint, app_id.clone().0))
+        });
+
+        // As long as we have a blueprint-id, create the blueprint.
+        blueprint_id
+            .as_ref()
+            .map(|id| self.store_dbs.blueprint_entry(id));
+
+        // If we have a blueprint, we can return the `StoreContext`. In most
+        // cases it should have already existed or been created above.
+        blueprint_id
+            .and_then(|id| self.store_dbs.blueprint(id))
+            .map(|blueprint| {
+                let recording = self
+                    .selected_rec_id
+                    .as_ref()
+                    .and_then(|id| self.store_dbs.recording(id));
+
+                // TODO(antoine): The below filter will limit our recording view to the current
+                // `ApplicationId`. Leaving this commented out for now since that is a bigger
+                // behavioral change we might want to plan/communicate around as it breaks things
+                // like --split-recordings in the api_demo.
+                StoreContext {
+                    blueprint,
+                    recording,
+                    alternate_recordings: self
+                        .store_dbs
+                        .recordings()
+                        //.filter(|rec| rec.app_id() == self.application_id.as_ref())
+                        .collect_vec(),
+                }
+            })
+    }
+
+    /// Change the selected/visible recording id.
+    /// This will also change the application-id to match the newly selected recording.
+    pub fn set_recording_id(&mut self, recording_id: StoreId) {
+        // If this recording corresponds to an app that we know about, then apdate the app-id.
+        if let Some(app_id) = self
+            .store_dbs
+            .recording(&recording_id)
+            .as_ref()
+            .and_then(|recording| recording.app_id())
+        {
+            self.set_app_id(app_id.clone());
+        }
+
+        self.selected_rec_id = Some(recording_id);
+    }
+
+    /// Change the selected [`ApplicationId`]
+    pub fn set_app_id(&mut self, app_id: ApplicationId) {
+        self.application_id = Some(app_id);
+    }
+
+    /// Change which blueprint is active for a given [`ApplicationId`]
+    pub fn set_blueprint_for_app_id(&mut self, blueprint_id: StoreId, app_id: ApplicationId) {
+        self.blueprint_by_app_id.insert(app_id, blueprint_id);
+    }
+
+    /// Clear the current blueprint
+    pub fn clear_blueprint(&mut self) {
+        if let Some(app_id) = &self.application_id {
+            if let Some(blueprint_id) = self.blueprint_by_app_id.remove(app_id) {
+                self.store_dbs.remove(&blueprint_id);
+            }
+        }
+    }
+
+    /// Mutable access to a [`StoreDb`] by id
+    pub fn store_db_mut(&mut self, store_id: &StoreId) -> &mut StoreDb {
+        self.store_dbs.store_db_entry(store_id)
+    }
+
+    /// Remove any empty [`StoreDb`]s from the hub
+    pub fn purge_empty(&mut self) {
+        self.store_dbs.purge_empty();
+    }
+
+    /// Call [`StoreDb::purge_fraction_of_ram`] on every recording
+    pub fn purge_fraction_of_ram(&mut self, fraction_to_purge: f32) {
+        self.store_dbs.purge_fraction_of_ram(fraction_to_purge);
+    }
+
+    /// Directly access the [`StoreDb`] for the selected recording
+    pub fn current_recording(&self) -> Option<&StoreDb> {
+        self.selected_rec_id
+            .as_ref()
+            .and_then(|id| self.store_dbs.recording(id))
+    }
+
+    /// Check whether the [`StoreHub`] contains the referenced recording
+    pub fn contains_recording(&self, id: &StoreId) -> bool {
+        self.store_dbs.contains_recording(id)
+    }
+
+    /// Populate a [`StoreHubStats`] based on the selected app.
+    // TODO(jleibs): We probably want stats for all recordings, not just
+    // the currently selected recording.
+    pub fn stats(&self) -> StoreHubStats {
+        // If we have an app-id, then use it to look up the blueprint.
+        let blueprint = self
+            .application_id
+            .as_ref()
+            .and_then(|app_id| self.blueprint_by_app_id.get(app_id))
+            .and_then(|blueprint_id| self.store_dbs.blueprint(blueprint_id));
+
+        let blueprint_stats = blueprint
+            .map(|store_db| DataStoreStats::from_store(&store_db.entity_db.data_store))
+            .unwrap_or_default();
+
+        let blueprint_config = blueprint
+            .map(|store_db| store_db.entity_db.data_store.config().clone())
+            .unwrap_or_default();
+
+        let recording = self
+            .selected_rec_id
+            .as_ref()
+            .and_then(|rec_id| self.store_dbs.recording(rec_id));
+
+        let recording_stats = recording
+            .map(|store_db| DataStoreStats::from_store(&store_db.entity_db.data_store))
+            .unwrap_or_default();
+
+        let recording_config = recording
+            .map(|store_db| store_db.entity_db.data_store.config().clone())
+            .unwrap_or_default();
+
+        StoreHubStats {
+            blueprint_stats,
+            blueprint_config,
+            recording_stats,
+            recording_config,
+        }
+    }
+}
+
+/// Stores many [`StoreDb`]s of recordings and blueprints.
+#[derive(Default)]
+pub struct StoreBundle {
     // TODO(emilk): two separate maps per [`StoreKind`].
     store_dbs: ahash::HashMap<StoreId, StoreDb>,
 }
 
-impl StoreHub {
+impl StoreBundle {
     /// Decode an rrd stream.
     /// It can theoretically contain multiple recordings, and blueprints.
     pub fn from_rrd(read: impl std::io::Read) -> anyhow::Result<Self> {
@@ -49,6 +230,10 @@ impl StoreHub {
         }
     }
 
+    pub fn remove(&mut self, id: &StoreId) {
+        self.store_dbs.remove(id);
+    }
+
     // --
 
     pub fn contains_recording(&self, id: &StoreId) -> bool {