Fix stuck menu when submenu vanishes (#7589)

* Closes https://github.com/rerun-io/rerun/issues/11301

This fixes a bug where a menu could get stuck, not closing at all, when
the currently open submenu stops being shown.
I also added a way to reproduce this to the demo, as well as a test
ensuring that there is no race condition in the fix.

crates/egui/src/containers/menu.rs

@@ -152,7 +152,8 @@ impl MenuState {
     pub fn from_id<R>(ctx: &Context, id: Id, f: impl FnOnce(&mut Self) -> R) -> R {
         let pass_nr = ctx.cumulative_pass_nr();
         ctx.data_mut(|data| {
-            let state = data.get_temp_mut_or_insert_with(id.with(Self::ID), || Self {
+            let state_id = id.with(Self::ID);
+            let mut state = data.get_temp(state_id).unwrap_or(Self {
                 open_item: None,
                 last_visible_pass: pass_nr,
             });
@@ -160,14 +161,38 @@ impl MenuState {
             if state.last_visible_pass + 1 < pass_nr {
                 state.open_item = None;
             }
-            state.last_visible_pass = pass_nr;
+            if let Some(item) = state.open_item {
-            f(state)
+                if data
+                    .get_temp(item.with(Self::ID))
+                    .is_none_or(|item: Self| item.last_visible_pass + 1 < pass_nr)
+                {
+                    // If the open item wasn't shown for at least a frame, reset the open item
+                    state.open_item = None;
+                }
+            }
+            let r = f(&mut state);
+            data.insert_temp(state_id, state);
+            r
         })
     }
 
+    pub fn mark_shown(ctx: &Context, id: Id) {
+        let pass_nr = ctx.cumulative_pass_nr();
+        Self::from_id(ctx, id, |state| {
+            state.last_visible_pass = pass_nr;
+        });
+    }
+
     /// Is the menu with this id the deepest sub menu? (-> no child sub menu is open)
-    pub fn is_deepest_sub_menu(ctx: &Context, id: Id) -> bool {
+    ///
-        Self::from_id(ctx, id, |state| state.open_item.is_none())
+    /// Note: This only returns correct results if called after the menu contents were shown.
+    pub fn is_deepest_open_sub_menu(ctx: &Context, id: Id) -> bool {
+        let pass_nr = ctx.cumulative_pass_nr();
+        let open_item = Self::from_id(ctx, id, |state| state.open_item);
+        // If we have some open item, check if that was actually shown this frame
+        open_item.is_none_or(|submenu_id| {
+            Self::from_id(ctx, submenu_id, |state| state.last_visible_pass != pass_nr)
+        })
     }
 }
 
@@ -399,6 +424,9 @@ impl SubMenu {
     }
 
     /// Show the submenu.
+    ///
+    /// This does some heuristics to check if the `button_response` was the last thing in the
+    /// menu that was hovered/clicked, and if so, shows the submenu.
     pub fn show<R>(
         self,
         ui: &Ui,
@@ -409,6 +437,7 @@ impl SubMenu {
 
         let id = Self::id_from_widget_id(button_response.id);
 
+        // Get the state from the parent menu
         let (open_item, menu_id, parent_config) = MenuState::from_ui(ui, |state, stack| {
             (state.open_item, stack.id, MenuConfig::from_stack(stack))
         });
@@ -452,7 +481,7 @@ impl SubMenu {
             set_open = Some(true);
             is_open = true;
             // Ensure that all other sub menus are closed when we open the menu
-            MenuState::from_id(ui.ctx(), id, |state| {
+            MenuState::from_id(ui.ctx(), menu_id, |state| {
                 state.open_item = None;
             });
         }
@@ -488,7 +517,7 @@ impl SubMenu {
 
         if let Some(popup_response) = &popup_response {
             // If no child sub menu is open means we must be the deepest child sub menu.
-            let is_deepest_submenu = MenuState::is_deepest_sub_menu(ui.ctx(), id);
+            let is_deepest_submenu = MenuState::is_deepest_open_sub_menu(ui.ctx(), id);
 
             // If the user clicks and the cursor is not hovering over our menu rect, it's
             // safe to assume they clicked outside the menu, so we close everything.

crates/egui/src/containers/popup.rs

@@ -604,8 +604,11 @@ impl<'a> Popup<'a> {
             PopupCloseBehavior::IgnoreClicks => false,
         };
 
+        // Mark the menu as shown, so the sub menu open state is not reset
+        MenuState::mark_shown(&ctx, id);
+
         // If a submenu is open, the CloseBehavior is handled there
-        let is_any_submenu_open = !MenuState::is_deepest_sub_menu(&response.response.ctx, id);
+        let is_any_submenu_open = !MenuState::is_deepest_open_sub_menu(&response.response.ctx, id);
 
         let should_close = (!is_any_submenu_open && closed_by_click)
             || ctx.input(|i| i.key_pressed(Key::Escape))

crates/egui_demo_lib/src/demo/popups.rs

@@ -20,6 +20,19 @@ pub struct PopupsDemo {
     color: egui::Color32,
 }
 
+impl Default for PopupsDemo {
+    fn default() -> Self {
+        Self {
+            align4: RectAlign::default(),
+            gap: 4.0,
+            close_behavior: PopupCloseBehavior::CloseOnClick,
+            popup_open: false,
+            checked: true,
+            color: egui::Color32::RED,
+        }
+    }
+}
+
 impl PopupsDemo {
     fn apply_options<'a>(&self, popup: Popup<'a>) -> Popup<'a> {
         popup
@@ -95,6 +108,14 @@ impl PopupsDemo {
                     color_picker_color32(ui, &mut self.color, Alpha::Opaque);
                 });
 
+                if self.checked {
+                    ui.menu_button("Only visible when checked", |ui| {
+                        if ui.button("Remove myself").clicked() {
+                            self.checked = false;
+                        }
+                    });
+                }
+
                 if ui.button("Open…").clicked() {
                     ui.close();
                 }
@@ -102,19 +123,6 @@ impl PopupsDemo {
     }
 }
 
-impl Default for PopupsDemo {
-    fn default() -> Self {
-        Self {
-            align4: RectAlign::default(),
-            gap: 4.0,
-            close_behavior: PopupCloseBehavior::CloseOnClick,
-            popup_open: false,
-            checked: false,
-            color: egui::Color32::RED,
-        }
-    }
-}
-
 impl crate::Demo for PopupsDemo {
     fn name(&self) -> &'static str {
         "\u{2755} Popups"

crates/egui_kittest/tests/regression_tests.rs

@@ -1,5 +1,5 @@
 use egui::accesskit::{self, Role};
-use egui::{Button, ComboBox, Image, Vec2, Widget as _};
+use egui::{Button, ComboBox, Image, Modifiers, Popup, Vec2, Widget as _};
 #[cfg(all(feature = "wgpu", feature = "snapshot"))]
 use egui_kittest::SnapshotResults;
 use egui_kittest::{Harness, kittest::Queryable as _};
@@ -187,3 +187,78 @@ pub fn override_text_color_affects_interactive_widgets() {
     #[cfg(all(feature = "wgpu", feature = "snapshot"))]
     results.add(harness.try_snapshot("override_text_color_interactive"));
 }
+
+/// <https://github.com/rerun-io/rerun/issues/11301>
+#[test]
+pub fn menus_should_close_even_if_submenu_disappears() {
+    const OTHER_BUTTON: &str = "Other button";
+    const MENU_BUTTON: &str = "Menu";
+    const SUB_MENU_BUTTON: &str = "Always here";
+    const TOGGLABLE_SUB_MENU_BUTTON: &str = "Maybe here";
+    const INSIDE_SUB_MENU_BUTTON: &str = "Inside submenu";
+
+    for frame_delay in (0..3).rev() {
+        let mut harness = Harness::builder().build_ui_state(
+            |ui, state| {
+                let _ = ui.button(OTHER_BUTTON).clicked();
+                let response = ui.button(MENU_BUTTON);
+
+                Popup::menu(&response).show(|ui| {
+                    let _ = ui.button(SUB_MENU_BUTTON);
+                    if *state {
+                        ui.menu_button(TOGGLABLE_SUB_MENU_BUTTON, |ui| {
+                            let _ = ui.button(INSIDE_SUB_MENU_BUTTON);
+                        });
+                    }
+                });
+            },
+            true,
+        );
+
+        // Open the main menu
+        harness.get_by_label(MENU_BUTTON).click();
+        harness.run();
+
+        // Open the sub menu
+        harness
+            .get_by_label_contains(TOGGLABLE_SUB_MENU_BUTTON)
+            .hover();
+        harness.run();
+
+        // Have we opened the submenu successfully?
+        harness.get_by_label(INSIDE_SUB_MENU_BUTTON).hover();
+        harness.run();
+
+        // We click manually, since we want to precisely time that the sub menu disappears when the
+        // button is released
+        let center = harness.get_by_label(OTHER_BUTTON).rect().center();
+        harness.input_mut().events.push(egui::Event::PointerButton {
+            pos: center,
+            button: egui::PointerButton::Primary,
+            pressed: true,
+            modifiers: Modifiers::default(),
+        });
+        harness.step();
+
+        // Yank the sub menu from under the pointer
+        *harness.state_mut() = false;
+
+        // See if we handle it with or without a frame delay
+        harness.run_steps(frame_delay);
+
+        // Actually close the menu by clicking somewhere outside
+        harness.input_mut().events.push(egui::Event::PointerButton {
+            pos: center,
+            button: egui::PointerButton::Primary,
+            pressed: false,
+            modifiers: Modifiers::default(),
+        });
+
+        harness.run();
+
+        assert!(
+            harness.query_by_label_contains(SUB_MENU_BUTTON).is_none(),
+            "Menu failed to close. frame_delay = {frame_delay}"
+        );
+    }
+}