Add support for undo in the viewer (alternative take)

crates/store/re_entity_db/src/entity_db.rs

@@ -18,7 +18,7 @@ use crate::{Error, TimesPerTimeline};
 // ----------------------------------------------------------------------------
 
 /// See [`GarbageCollectionOptions::time_budget`].
-const DEFAULT_GC_TIME_BUDGET: std::time::Duration = std::time::Duration::from_micros(3500); // empirical
+pub const DEFAULT_GC_TIME_BUDGET: std::time::Duration = std::time::Duration::from_micros(3500); // empirical
 
 // ----------------------------------------------------------------------------
 
@@ -387,19 +387,6 @@ impl EntityDb {
         self.set_store_info = Some(store_info);
     }
 
-    pub fn gc_everything_but_the_latest_row_on_non_default_timelines(
-        &mut self,
-    ) -> Vec<ChunkStoreEvent> {
-        re_tracing::profile_function!();
-
-        self.gc(&GarbageCollectionOptions {
-            target: GarbageCollectionTarget::Everything,
-            protect_latest: 1,
-            time_budget: DEFAULT_GC_TIME_BUDGET,
-            protected_time_ranges: Default::default(), // TODO(#3135): Use this for undo buffer
-        })
-    }
-
     /// Free up some RAM by forgetting the older parts of all timelines.
     pub fn purge_fraction_of_ram(&mut self, fraction_to_purge: f32) -> Vec<ChunkStoreEvent> {
         re_tracing::profile_function!();
@@ -430,7 +417,7 @@ impl EntityDb {
         store_events
     }
 
-    fn gc(&mut self, gc_options: &GarbageCollectionOptions) -> Vec<ChunkStoreEvent> {
+    pub fn gc(&mut self, gc_options: &GarbageCollectionOptions) -> Vec<ChunkStoreEvent> {
         re_tracing::profile_function!();
 
         let (store_events, stats_diff) = self.data_store.gc(gc_options);

crates/store/re_entity_db/src/lib.rs

@@ -13,7 +13,7 @@ mod times_per_timeline;
 mod versioned_instance_path;
 
 pub use self::{
-    entity_db::EntityDb,
+    entity_db::{EntityDb, DEFAULT_GC_TIME_BUDGET},
     entity_tree::EntityTree,
     instance_path::{InstancePath, InstancePathHash},
     store_bundle::{StoreBundle, StoreLoadError},

crates/viewer/re_ui/src/command.rs

@@ -20,6 +20,9 @@ pub enum UICommand {
     CloseCurrentRecording,
     CloseAllRecordings,
 
+    Undo,
+    Redo,
+
     #[cfg(not(target_arch = "wasm32"))]
     Quit,
 
@@ -119,7 +122,10 @@ impl UICommand {
             ),
 
             Self::CloseAllRecordings => ("Close all recordings",
-                "Close all open current recording (unsaved data will be lost)",),
+                "Close all open current recording (unsaved data will be lost)"),
+
+            Self::Undo => ("Undo", "Undo the last blueprint edit for the open recording"),
+            Self::Redo => ("Redo", "Redo the last undone thing"),
 
             #[cfg(not(target_arch = "wasm32"))]
             Self::Quit => ("Quit", "Close the Rerun Viewer"),
@@ -271,12 +277,16 @@ impl UICommand {
             KeyboardShortcut::new(Modifiers::COMMAND, key)
         }
 
+        fn cmd_shift(key: Key) -> KeyboardShortcut {
+            KeyboardShortcut::new(Modifiers::COMMAND | Modifiers::SHIFT, key)
+        }
+
         fn cmd_alt(key: Key) -> KeyboardShortcut {
-            KeyboardShortcut::new(Modifiers::COMMAND.plus(Modifiers::ALT), key)
+            KeyboardShortcut::new(Modifiers::COMMAND | Modifiers::ALT, key)
         }
 
         fn ctrl_shift(key: Key) -> KeyboardShortcut {
-            KeyboardShortcut::new(Modifiers::CTRL.plus(Modifiers::SHIFT), key)
+            KeyboardShortcut::new(Modifiers::CTRL | Modifiers::SHIFT, key)
         }
 
         match self {
@@ -287,6 +297,9 @@ impl UICommand {
             Self::CloseCurrentRecording => None,
             Self::CloseAllRecordings => None,
 
+            Self::Undo => Some(cmd(Key::Z)),
+            Self::Redo => Some(cmd_shift(Key::Z)),
+
             #[cfg(all(not(target_arch = "wasm32"), target_os = "windows"))]
             Self::Quit => Some(KeyboardShortcut::new(Modifiers::ALT, Key::F4)),
 

crates/viewer/re_viewer/src/app.rs

@@ -10,9 +10,9 @@ use re_ui::{toasts, DesignTokens, UICommand, UICommandSender};
 use re_viewer_context::{
     command_channel,
     store_hub::{BlueprintPersistence, StoreHub, StoreHubStats},
-    AppOptions, CommandReceiver, CommandSender, ComponentUiRegistry, PlayState, SpaceViewClass,
-    SpaceViewClassRegistry, SpaceViewClassRegistryError, StoreContext, SystemCommand,
-    SystemCommandSender,
+    AppOptions, BlueprintUndoState, CommandReceiver, CommandSender, ComponentUiRegistry, PlayState,
+    SpaceViewClass, SpaceViewClassRegistry, SpaceViewClassRegistryError, StoreContext,
+    SystemCommand, SystemCommandSender,
 };
 
 use crate::app_blueprint::PanelStateOverrides;
@@ -514,6 +514,15 @@ impl App {
                 // to apply updates here, but this needs more validation and testing to be safe.
                 if !self.state.app_options.inspect_blueprint_timeline {
                     let blueprint_db = store_hub.entity_db_mut(&blueprint_id);
+
+                    let undo_state = self
+                        .state
+                        .blueprint_undo_state
+                        .entry(blueprint_id)
+                        .or_default();
+
+                    undo_state.clear_redo(blueprint_db);
+
                     for chunk in updates {
                         match blueprint_db.add_chunk(&Arc::new(chunk)) {
                             Ok(_store_events) => {}
@@ -524,6 +533,23 @@ impl App {
                     }
                 }
             }
+            SystemCommand::UndoBlueprint { blueprint_id } => {
+                let blueprint_db = store_hub.entity_db_mut(&blueprint_id);
+                self.state
+                    .blueprint_undo_state
+                    .entry(blueprint_id)
+                    .or_default()
+                    .undo(blueprint_db);
+            }
+            SystemCommand::RedoBlueprint { blueprint_id } => {
+                let blueprint_db = store_hub.entity_db_mut(&blueprint_id);
+                self.state
+                    .blueprint_undo_state
+                    .entry(blueprint_id)
+                    .or_default()
+                    .redo(blueprint_db);
+            }
+
             SystemCommand::DropEntity(blueprint_id, entity_path) => {
                 let blueprint_db = store_hub.entity_db_mut(&blueprint_id);
                 blueprint_db.drop_entity_path_recursive(&entity_path);
@@ -617,6 +643,21 @@ impl App {
                     .send_system(SystemCommand::CloseAllRecordings);
             }
 
+            UICommand::Undo => {
+                if let Some(store_context) = store_context {
+                    let blueprint_id = store_context.blueprint.store_id().clone();
+                    self.command_sender
+                        .send_system(SystemCommand::UndoBlueprint { blueprint_id });
+                }
+            }
+            UICommand::Redo => {
+                if let Some(store_context) = store_context {
+                    let blueprint_id = store_context.blueprint.store_id().clone();
+                    self.command_sender
+                        .send_system(SystemCommand::RedoBlueprint { blueprint_id });
+                }
+            }
+
             #[cfg(not(target_arch = "wasm32"))]
             UICommand::Quit => {
                 egui_ctx.send_viewport_cmd(egui::ViewportCommand::Close);
@@ -1500,7 +1541,7 @@ impl eframe::App for App {
         // TODO(#2579): implement web-storage for blueprints as well
         if let Some(hub) = &mut self.store_hub {
             if self.state.app_options.blueprint_gc {
-                hub.gc_blueprints();
+                hub.gc_blueprints(&self.state.blueprint_undo_state);
             }
 
             if let Err(err) = hub.save_app_blueprints() {
@@ -1618,7 +1659,7 @@ impl eframe::App for App {
         self.receive_messages(&mut store_hub, egui_ctx);
 
         if self.app_options().blueprint_gc {
-            store_hub.gc_blueprints();
+            store_hub.gc_blueprints(&self.state.blueprint_undo_state);
         }
 
         store_hub.purge_empty();
@@ -1644,9 +1685,17 @@ impl eframe::App for App {
 
         let store_context = store_hub.read_context();
 
+        let blueprint_query =
+            store_context
+                .as_ref()
+                .map_or(BlueprintUndoState::default_query(), |store_context| {
+                    self.state
+                        .blueprint_query_for_viewer(store_context.blueprint)
+                });
+
         let app_blueprint = AppBlueprint::new(
             store_context.as_ref(),
-            &self.state.blueprint_query_for_viewer(),
+            &blueprint_query,
             egui_ctx,
             self.panel_state_overrides_active
                 .then_some(self.panel_state_overrides),

crates/viewer/re_viewer/src/app_state.rs

@@ -8,7 +8,7 @@ use re_smart_channel::ReceiveSet;
 use re_types::blueprint::components::PanelState;
 use re_ui::ContextExt as _;
 use re_viewer_context::{
-    blueprint_timeline, AppOptions, ApplicationSelectionState, CommandSender, ComponentUiRegistry,
+    AppOptions, ApplicationSelectionState, BlueprintUndoState, CommandSender, ComponentUiRegistry,
     PlayState, RecordingConfig, SpaceViewClassExt as _, SpaceViewClassRegistry, StoreContext,
     StoreHub, SystemCommandSender as _, ViewStates, ViewerContext,
 };
@@ -31,6 +31,9 @@ pub struct AppState {
     recording_configs: HashMap<StoreId, RecordingConfig>,
     blueprint_cfg: RecordingConfig,
 
+    /// Maps blueperint id to the current undo state for it.
+    pub blueprint_undo_state: HashMap<StoreId, BlueprintUndoState>,
+
     selection_panel: re_selection_panel::SelectionPanel,
     time_panel: re_time_panel::TimePanel,
     blueprint_panel: re_time_panel::TimePanel,
@@ -69,6 +72,7 @@ impl Default for AppState {
         Self {
             app_options: Default::default(),
             recording_configs: Default::default(),
+            blueprint_undo_state: Default::default(),
             blueprint_cfg: Default::default(),
             selection_panel: Default::default(),
             time_panel: Default::default(),
@@ -139,11 +143,12 @@ impl AppState {
     ) {
         re_tracing::profile_function!();
 
-        let blueprint_query = self.blueprint_query_for_viewer();
+        let blueprint_query = self.blueprint_query_for_viewer(store_context.blueprint);
 
         let Self {
             app_options,
             recording_configs,
+            blueprint_undo_state,
             blueprint_cfg,
             selection_panel,
             time_panel,
@@ -160,6 +165,11 @@ impl AppState {
         // check state early, before the UI has a chance to close these popups
         let is_any_popup_open = ui.memory(|m| m.any_popup_open());
 
+        blueprint_undo_state
+            .entry(store_context.blueprint.store_id().clone())
+            .or_default()
+            .update(ui.ctx(), store_context.blueprint);
+
         // Some of the mutations APIs of `ViewportBlueprints` are recorded as `Viewport::TreeAction`
         // and must be applied by `Viewport` at the end of the frame. We use a temporary channel for
         // this, which gives us interior mutability (only a shared reference of `ViewportBlueprint`
@@ -487,18 +497,25 @@ impl AppState {
 
         self.recording_configs
             .retain(|store_id, _| store_hub.store_bundle().contains(store_id));
+
+        self.blueprint_undo_state
+            .retain(|store_id, _| store_hub.store_bundle().contains(store_id));
     }
 
     /// Returns the blueprint query that should be used for generating the current
     /// layout of the viewer.
     ///
     /// If `inspect_blueprint_timeline` is enabled, we use the time selection from the
     /// blueprint `time_ctrl`. Otherwise, we use a latest query from the blueprint timeline.
-    pub fn blueprint_query_for_viewer(&self) -> LatestAtQuery {
+    pub fn blueprint_query_for_viewer(&mut self, blueprint: &EntityDb) -> LatestAtQuery {
         if self.app_options.inspect_blueprint_timeline {
             self.blueprint_cfg.time_ctrl.read().current_query().clone()
         } else {
-            LatestAtQuery::latest(blueprint_timeline())
+            let undo_state = self
+                .blueprint_undo_state
+                .entry(blueprint.store_id().clone())
+                .or_default();
+            undo_state.blueprint_query()
         }
     }
 }

crates/viewer/re_viewer/src/ui/rerun_menu.rs

@@ -32,6 +32,9 @@ impl App {
 
             ui.add_space(SPACING);
 
+            UICommand::Undo.menu_button_ui(ui, &self.command_sender); // TODO(emilk): only enabled if there is something to undo
+            UICommand::Redo.menu_button_ui(ui, &self.command_sender); // TODO(emilk): only enabled if there is something to redo
+
             UICommand::ToggleCommandPalette.menu_button_ui(ui, &self.command_sender);
 
             ui.add_space(SPACING);

crates/viewer/re_viewer_context/src/command_sender.rs

@@ -49,6 +49,13 @@ pub enum SystemCommand {
     /// is both modified and changed in the same frame.
     UpdateBlueprint(StoreId, Vec<Chunk>),
 
+    UndoBlueprint {
+        blueprint_id: StoreId,
+    },
+    RedoBlueprint {
+        blueprint_id: StoreId,
+    },
+
     /// Drop a specific entity from a store.
     ///
     /// Also drops all recursive children.

crates/viewer/re_viewer_context/src/lib.rs

@@ -28,6 +28,7 @@ pub mod test_context; //TODO(ab): this should be behind #[cfg(test)], but then `
 mod time_control;
 mod time_drag_value;
 mod typed_entity_collections;
+mod undo;
 mod utils;
 mod viewer_context;
 
@@ -83,6 +84,7 @@ pub use time_drag_value::TimeDragValue;
 pub use typed_entity_collections::{
     ApplicableEntities, IndicatedEntities, PerVisualizer, VisualizableEntities,
 };
+pub use undo::BlueprintUndoState;
 pub use utils::{auto_color_egui, auto_color_for_entity_path, level_to_rich_text};
 pub use viewer_context::{RecordingConfig, ViewerContext};
 

crates/viewer/re_viewer_context/src/store_hub.rs

@@ -3,12 +3,15 @@ use ahash::{HashMap, HashMapExt, HashSet};
 use anyhow::Context as _;
 use itertools::Itertools as _;
 
-use re_chunk_store::{ChunkStoreConfig, ChunkStoreGeneration, ChunkStoreStats};
+use re_chunk_store::{
+    ChunkStoreConfig, ChunkStoreGeneration, ChunkStoreStats, GarbageCollectionOptions,
+    GarbageCollectionTarget,
+};
 use re_entity_db::{EntityDb, StoreBundle};
-use re_log_types::{ApplicationId, StoreId, StoreKind};
+use re_log_types::{ApplicationId, ResolvedTimeRange, StoreId, StoreKind};
 use re_query::CachesStats;
 
-use crate::{Caches, StoreContext};
+use crate::{BlueprintUndoState, Caches, StoreContext};
 
 /// Interface for accessing all blueprints and recordings
 ///
@@ -673,7 +676,7 @@ impl StoreHub {
         });
     }
 
-    pub fn gc_blueprints(&mut self) {
+    pub fn gc_blueprints(&mut self, undo_state: &HashMap<StoreId, BlueprintUndoState>) {
         re_tracing::profile_function!();
 
         for blueprint_id in self
@@ -686,10 +689,23 @@ impl StoreHub {
                     continue; // no change since last gc
                 }
 
-                // TODO(jleibs): Decide a better tuning for this. Would like to save a
-                // reasonable amount of history, or incremental snapshots.
-                let store_events =
-                    blueprint.gc_everything_but_the_latest_row_on_non_default_timelines();
+                let mut protected_time_ranges = ahash::HashMap::default();
+                if let Some(undo) = undo_state.get(blueprint_id) {
+                    if let Some(time) = undo.oldest_undo_point() {
+                        // Save everything that we could want to undo to:
+                        protected_time_ranges.insert(
+                            crate::blueprint_timeline(),
+                            ResolvedTimeRange::new(time, re_chunk::TimeInt::MAX),
+                        );
+                    }
+                }
+
+                let store_events = blueprint.gc(&GarbageCollectionOptions {
+                    target: GarbageCollectionTarget::Everything,
+                    protect_latest: 0,
+                    time_budget: re_entity_db::DEFAULT_GC_TIME_BUDGET,
+                    protected_time_ranges,
+                });
                 if let Some(caches) = self.caches_per_recording.get_mut(blueprint_id) {
                     caches.on_store_events(&store_events);
                 }

crates/viewer/re_viewer_context/src/test_context.rs

@@ -166,6 +166,8 @@ impl TestContext {
                 | SystemCommand::ClearAndGenerateBlueprint
                 | SystemCommand::ActivateRecording(_)
                 | SystemCommand::CloseStore(_)
+                | SystemCommand::UndoBlueprint { .. }
+                | SystemCommand::RedoBlueprint { .. }
                 | SystemCommand::CloseAllRecordings => handled = false,
 
                 #[cfg(debug_assertions)]