fix(segmented_button): diff the context menu

src/widget/menu.rs

@@ -72,5 +72,6 @@ pub use menu_tree::{
 };
 
 pub use crate::style::menu_bar::{Appearance, StyleSheet};
+pub(crate) use menu_bar::{menu_roots_children, menu_roots_diff};
 pub(crate) use menu_inner::Menu;
 pub use menu_inner::{CloseCondition, ItemHeight, ItemWidth, PathHighlight};

src/widget/menu/menu_bar.rs

@@ -63,6 +63,82 @@ impl Default for MenuBarState {
     }
 }
 
+pub(crate) fn menu_roots_children<'a, Message, Renderer>(
+    menu_roots: &Vec<MenuTree<'a, Message, Renderer>>,
+) -> Vec<Tree>
+where
+    Renderer: renderer::Renderer,
+{
+    /*
+    menu bar
+        menu root 1 (stateless)
+            flat tree
+        menu root 2 (stateless)
+            flat tree
+        ...
+    */
+
+    menu_roots
+        .iter()
+        .map(|root| {
+            let mut tree = Tree::empty();
+            let flat = root
+                .flattern()
+                .iter()
+                .map(|mt| Tree::new(mt.item.as_widget()))
+                .collect();
+            tree.children = flat;
+            tree
+        })
+        .collect()
+}
+
+#[allow(invalid_reference_casting)]
+pub(crate) fn menu_roots_diff<'a, Message, Renderer>(
+    menu_roots: &mut Vec<MenuTree<'a, Message, Renderer>>,
+    tree: &mut Tree,
+) where
+    Renderer: renderer::Renderer,
+{
+    if tree.children.len() > menu_roots.len() {
+        tree.children.truncate(menu_roots.len());
+    }
+
+    tree.children
+        .iter_mut()
+        .zip(menu_roots.iter())
+        .for_each(|(t, root)| {
+            let mut flat = root
+                .flattern()
+                .iter()
+                .map(|mt| {
+                    let widget = mt.item.as_widget();
+                    let widget_ptr = widget as *const dyn Widget<Message, crate::Theme, Renderer>;
+                    let widget_ptr_mut =
+                        widget_ptr as *mut dyn Widget<Message, crate::Theme, Renderer>;
+                    //TODO: find a way to diff_children without unsafe code
+                    unsafe { &mut *widget_ptr_mut }
+                })
+                .collect::<Vec<_>>();
+
+            t.diff_children(flat.as_mut_slice());
+        });
+
+    if tree.children.len() < menu_roots.len() {
+        let extended = menu_roots[tree.children.len()..].iter().map(|root| {
+            let mut tree = Tree::empty();
+            let flat = root
+                .flattern()
+                .iter()
+                .map(|mt| Tree::new(mt.item.as_widget()))
+                .collect();
+            tree.children = flat;
+            tree
+        });
+        tree.children.extend(extended);
+    }
+}
+
 /// A `MenuBar` collects `MenuTree`s and handles all the layout, event processing, and drawing.
 #[allow(missing_debug_implementations)]
 pub struct MenuBar<'a, Message, Renderer = crate::Renderer>
@@ -212,46 +288,8 @@ where
         iced_core::Size::new(self.width, self.height)
     }
 
-    #[allow(invalid_reference_casting)]
     fn diff(&mut self, tree: &mut Tree) {
-        if tree.children.len() > self.menu_roots.len() {
+        menu_roots_diff(&mut self.menu_roots, tree);
-            tree.children.truncate(self.menu_roots.len());
-        }
-
-        tree.children
-            .iter_mut()
-            .zip(self.menu_roots.iter())
-            .for_each(|(t, root)| {
-                let mut flat = root
-                    .flattern()
-                    .iter()
-                    .map(|mt| {
-                        let widget = mt.item.as_widget();
-                        let widget_ptr =
-                            widget as *const dyn Widget<Message, crate::Theme, Renderer>;
-                        let widget_ptr_mut =
-                            widget_ptr as *mut dyn Widget<Message, crate::Theme, Renderer>;
-                        //TODO: find a way to diff_children without unsafe code
-                        unsafe { &mut *widget_ptr_mut }
-                    })
-                    .collect::<Vec<_>>();
-
-                t.diff_children(flat.as_mut_slice());
-            });
-
-        if tree.children.len() < self.menu_roots.len() {
-            let extended = self.menu_roots[tree.children.len()..].iter().map(|root| {
-                let mut tree = Tree::empty();
-                let flat = root
-                    .flattern()
-                    .iter()
-                    .map(|mt| Tree::new(mt.item.as_widget()))
-                    .collect();
-                tree.children = flat;
-                tree
-            });
-            tree.children.extend(extended);
-        }
     }
 
     fn tag(&self) -> tree::Tag {
@@ -263,28 +301,7 @@ where
     }
 
     fn children(&self) -> Vec<Tree> {
-        /*
+        menu_roots_children(&self.menu_roots)
-        menu bar
-            menu root 1 (stateless)
-                flat tree
-            menu root 2 (stateless)
-                flat tree
-            ...
-        */
-
-        self.menu_roots
-            .iter()
-            .map(|root| {
-                let mut tree = Tree::empty();
-                let flat = root
-                    .flattern()
-                    .iter()
-                    .map(|mt| Tree::new(mt.item.as_widget()))
-                    .collect();
-                tree.children = flat;
-                tree
-            })
-            .collect()
     }
 
     fn layout(&self, tree: &mut Tree, renderer: &Renderer, limits: &Limits) -> Node {

src/widget/segmented_button/widget.rs

@@ -6,7 +6,8 @@ use crate::iced_core::id::Internal;
 use crate::theme::{SegmentedButton as Style, THEME};
 use crate::widget::dnd_destination::DragId;
 use crate::widget::menu::{
-    self, CloseCondition, ItemHeight, ItemWidth, MenuBarState, PathHighlight,
+    self, menu_roots_children, menu_roots_diff, CloseCondition, ItemHeight, ItemWidth,
+    MenuBarState, PathHighlight,
 };
 use crate::widget::{icon, Icon};
 use crate::{Element, Renderer};
@@ -560,20 +561,7 @@ where
         if let Some(ref context_menu) = self.context_menu {
             let mut tree = Tree::empty();
             tree.state = tree::State::new(MenuBarState::default());
-            tree.children = context_menu
+            tree.children = menu_roots_children(&context_menu);
-                .iter()
-                .map(|root| {
-                    let mut tree = Tree::empty();
-                    let flat = root
-                        .flattern()
-                        .iter()
-                        .map(|mt| Tree::new(mt.item.as_widget()))
-                        .collect();
-                    tree.children = flat;
-                    tree
-                })
-                .collect();
-
             children.push(tree);
         }
 
@@ -655,7 +643,19 @@ where
             }
         }
 
-        // TODO: diff the context menu
+        // Diff the context menu
+        if let Some(context_menu) = &mut self.context_menu {
+            if tree.children.is_empty() {
+                let mut child_tree = Tree::empty();
+                child_tree.state = tree::State::new(MenuBarState::default());
+                tree.children.push(child_tree);
+            } else {
+                tree.children.truncate(1);
+            }
+            menu_roots_diff(context_menu, &mut tree.children[0]);
+        } else {
+            tree.children.clear();
+        }
     }
 
     fn size(&self) -> Size<Length> {