Retain widget state against incoming panes

We can associate each state with a `Pane` and compare
that against the new panes to remove states w/ respective
panes which no longer exist.

Because we always increment `Pane`, new states are always
added to the end, so this retain + add new state approach
will ensure continuity when panes are added & removed

2 files changed, 135 insertions, 125 deletions

diff --git a/widget/src/pane_grid.rs b/widget/src/pane_grid.rs
index e6fda660..2644986f 100644
--- a/widget/src/pane_grid.rs
+++ b/widget/src/pane_grid.rs
@@ -92,6 +92,8 @@ use crate::core::{
     Pixels, Point, Rectangle, Shell, Size, Theme, Vector, Widget,
 };
 
+use std::borrow::Cow;
+
 const DRAG_DEADBAND_DISTANCE: f32 = 10.0;
 const THICKNESS_RATIO: f32 = 25.0;
 
@@ -157,7 +159,10 @@ pub struct PaneGrid<
     Theme: Catalog,
     Renderer: core::Renderer,
 {
-    contents: Contents<'a, Content<'a, Message, Theme, Renderer>>,
+    internal: &'a state::Internal,
+    panes: Vec<Pane>,
+    contents: Vec<Content<'a, Message, Theme, Renderer>>,
+    maximized: Option<Pane>,
     width: Length,
     height: Length,
     spacing: f32,
@@ -180,30 +185,21 @@ where
         state: &'a State<T>,
         view: impl Fn(Pane, &'a T, bool) -> Content<'a, Message, Theme, Renderer>,
     ) -> Self {
-        let contents = if let Some((pane, pane_state)) =
-            state.maximized.and_then(|pane| {
-                state.panes.get(&pane).map(|pane_state| (pane, pane_state))
-            }) {
-            Contents::Maximized(
-                pane,
-                view(pane, pane_state, true),
-                Node::Pane(pane),
-            )
-        } else {
-            Contents::All(
-                state
-                    .panes
-                    .iter()
-                    .map(|(pane, pane_state)| {
-                        (*pane, view(*pane, pane_state, false))
-                    })
-                    .collect(),
-                &state.internal,
-            )
-        };
+        let panes = state.panes.keys().copied().collect();
+        let contents = state
+            .panes
+            .iter()
+            .map(|(pane, pane_state)| match &state.maximized {
+                Some(p) if pane == p => view(*pane, pane_state, true),
+                _ => view(*pane, pane_state, false),
+            })
+            .collect();
 
         Self {
+            internal: &state.internal,
+            panes,
             contents,
+            maximized: state.maximized,
             width: Length::Fill,
             height: Length::Fill,
             spacing: 0.0,
@@ -248,7 +244,9 @@ where
     where
         F: 'a + Fn(DragEvent) -> Message,
     {
-        self.on_drag = Some(Box::new(f));
+        if self.maximized.is_none() {
+            self.on_drag = Some(Box::new(f));
+        }
         self
     }
 
@@ -265,7 +263,9 @@ where
     where
         F: 'a + Fn(ResizeEvent) -> Message,
     {
-        self.on_resize = Some((leeway.into().0, Box::new(f)));
+        if self.maximized.is_none() {
+            self.on_resize = Some((leeway.into().0, Box::new(f)));
+        }
         self
     }
 
@@ -291,10 +291,17 @@ where
     }
 
     fn drag_enabled(&self) -> bool {
-        (!self.contents.is_maximized())
+        (self.maximized.is_none())
             .then(|| self.on_drag.is_some())
             .unwrap_or_default()
     }
+
+    fn node(&self) -> Cow<'_, Node> {
+        match self.maximized {
+            Some(pane) => Cow::Owned(Node::Pane(pane)),
+            None => Cow::Borrowed(&self.internal.layout),
+        }
+    }
 }
 
 impl<'a, Message, Theme, Renderer> Widget<Message, Theme, Renderer>
@@ -304,33 +311,48 @@ where
     Renderer: core::Renderer,
 {
     fn tag(&self) -> tree::Tag {
-        tree::Tag::of::<state::Action>()
+        tree::Tag::of::<state::Widget>()
     }
 
     fn state(&self) -> tree::State {
-        tree::State::new(state::Action::Idle)
+        tree::State::new(state::Widget::default())
     }
 
     fn children(&self) -> Vec<Tree> {
-        self.contents
-            .iter()
-            .map(|(_, content)| content.state())
-            .collect()
+        self.contents.iter().map(Content::state).collect()
     }
 
     fn diff(&self, tree: &mut Tree) {
-        match &self.contents {
-            Contents::All(contents, _) => tree.diff_children_custom(
-                contents,
-                |state, (_, content)| content.diff(state),
-                |(_, content)| content.state(),
-            ),
-            Contents::Maximized(_, content, _) => tree.diff_children_custom(
-                &[content],
-                |state, content| content.diff(state),
-                |content| content.state(),
-            ),
-        }
+        let state::Widget { panes, .. } = tree.state.downcast_ref();
+
+        // `Pane` always increments and is iterated by Ord so new
+        // states are always added at the end. We can simply remove
+        // states which no longer exist and `diff_children` will
+        // diff the remaining values in the correct order and
+        // add new states at the end
+
+        let mut i = 0;
+        let mut j = 0;
+        tree.children.retain(|_| {
+            let retain = self.panes.get(i) == panes.get(j);
+
+            if retain {
+                i += 1;
+            }
+            j += 1;
+
+            retain
+        });
+
+        tree.diff_children_custom(
+            &self.contents,
+            |state, content| content.diff(state),
+            Content::state,
+        );
+
+        let state::Widget { panes, .. } = tree.state.downcast_mut();
+
+        panes.clone_from(&self.panes);
     }
 
     fn size(&self) -> Size<Length> {
@@ -347,14 +369,19 @@ where
         limits: &layout::Limits,
     ) -> layout::Node {
         let size = limits.resolve(self.width, self.height, Size::ZERO);
-        let node = self.contents.layout();
-        let regions = node.pane_regions(self.spacing, size);
+        let regions = self.node().pane_regions(self.spacing, size);
 
         let children = self
-            .contents
+            .panes
             .iter()
+            .copied()
+            .zip(&self.contents)
             .zip(tree.children.iter_mut())
             .filter_map(|((pane, content), tree)| {
+                if self.maximized.is_some() && Some(pane) != self.maximized {
+                    return Some(layout::Node::new(Size::ZERO));
+                }
+
                 let region = regions.get(&pane)?;
                 let size = Size::new(region.width, region.height);
 
@@ -379,11 +406,16 @@ where
         operation: &mut dyn widget::Operation,
     ) {
         operation.container(None, layout.bounds(), &mut |operation| {
-            self.contents
+            self.panes
                 .iter()
+                .copied()
+                .zip(&self.contents)
                 .zip(&mut tree.children)
                 .zip(layout.children())
-                .for_each(|(((_pane, content), state), layout)| {
+                .filter(|(((pane, _), _), _)| {
+                    self.maximized.map_or(true, |maximized| *pane == maximized)
+                })
+                .for_each(|(((_, content), state), layout)| {
                     content.operate(state, layout, renderer, operation);
                 });
         });
@@ -402,8 +434,8 @@ where
     ) -> event::Status {
         let mut event_status = event::Status::Ignored;
 
-        let action = tree.state.downcast_mut::<state::Action>();
-        let node = self.contents.layout();
+        let state::Widget { action, .. } = tree.state.downcast_mut();
+        let node = self.node();
 
         let on_drag = if self.drag_enabled() {
             &self.on_drag
@@ -448,7 +480,10 @@ where
                                     layout,
                                     cursor_position,
                                     shell,
-                                    self.contents.iter(),
+                                    self.panes
+                                        .iter()
+                                        .copied()
+                                        .zip(&self.contents),
                                     &self.on_click,
                                     on_drag,
                                 );
@@ -460,7 +495,7 @@ where
                                 layout,
                                 cursor_position,
                                 shell,
-                                self.contents.iter(),
+                                self.panes.iter().copied().zip(&self.contents),
                                 &self.on_click,
                                 on_drag,
                             );
@@ -486,8 +521,10 @@ where
                                     }
                                 } else {
                                     let dropped_region = self
-                                        .contents
+                                        .panes
                                         .iter()
+                                        .copied()
+                                        .zip(&self.contents)
                                         .zip(layout.children())
                                         .find_map(|(target, layout)| {
                                             layout_region(
@@ -572,10 +609,15 @@ where
 
         let picked_pane = action.picked_pane().map(|(pane, _)| pane);
 
-        self.contents
-            .iter_mut()
+        self.panes
+            .iter()
+            .copied()
+            .zip(&mut self.contents)
             .zip(&mut tree.children)
             .zip(layout.children())
+            .filter(|(((pane, _), _), _)| {
+                self.maximized.map_or(true, |maximized| *pane == maximized)
+            })
             .map(|(((pane, content), tree), layout)| {
                 let is_picked = picked_pane == Some(pane);
 
@@ -602,14 +644,14 @@ where
         viewport: &Rectangle,
         renderer: &Renderer,
     ) -> mouse::Interaction {
-        let action = tree.state.downcast_ref::<state::Action>();
+        let state::Widget { action, .. } = tree.state.downcast_ref();
 
         if action.picked_pane().is_some() {
             return mouse::Interaction::Grabbing;
         }
 
         let resize_leeway = self.on_resize.as_ref().map(|(leeway, _)| *leeway);
-        let node = self.contents.layout();
+        let node = self.node();
 
         let resize_axis =
             action.picked_split().map(|(_, axis)| axis).or_else(|| {
@@ -641,11 +683,16 @@ where
             };
         }
 
-        self.contents
+        self.panes
             .iter()
+            .copied()
+            .zip(&self.contents)
             .zip(&tree.children)
             .zip(layout.children())
-            .map(|(((_pane, content), tree), layout)| {
+            .filter(|(((pane, _), _), _)| {
+                self.maximized.map_or(true, |maximized| *pane == maximized)
+            })
+            .map(|(((_, content), tree), layout)| {
                 content.mouse_interaction(
                     tree,
                     layout,
@@ -669,16 +716,11 @@ where
         cursor: mouse::Cursor,
         viewport: &Rectangle,
     ) {
-        let action = tree.state.downcast_ref::<state::Action>();
-        let node = self.contents.layout();
+        let state::Widget { action, .. } =
+            tree.state.downcast_ref::<state::Widget>();
+        let node = self.node();
         let resize_leeway = self.on_resize.as_ref().map(|(leeway, _)| *leeway);
 
-        let contents = self
-            .contents
-            .iter()
-            .zip(&tree.children)
-            .map(|((pane, content), tree)| (pane, (content, tree)));
-
         let picked_pane = action.picked_pane().filter(|(_, origin)| {
             cursor
                 .position()
@@ -747,8 +789,16 @@ where
 
         let style = Catalog::style(theme, &self.class);
 
-        for ((id, (content, tree)), pane_layout) in
-            contents.zip(layout.children())
+        for (((id, content), tree), pane_layout) in self
+            .panes
+            .iter()
+            .copied()
+            .zip(&self.contents)
+            .zip(&tree.children)
+            .zip(layout.children())
+            .filter(|(((pane, _), _), _)| {
+                self.maximized.map_or(true, |maximized| maximized == *pane)
+            })
         {
             match picked_pane {
                 Some((dragging, origin)) if id == dragging => {
@@ -883,11 +933,17 @@ where
         translation: Vector,
     ) -> Option<overlay::Element<'b, Message, Theme, Renderer>> {
         let children = self
-            .contents
-            .iter_mut()
+            .panes
+            .iter()
+            .copied()
+            .zip(&mut self.contents)
             .zip(&mut tree.children)
             .zip(layout.children())
-            .filter_map(|(((_, content), state), layout)| {
+            .filter_map(|(((pane, content), state), layout)| {
+                if self.maximized.is_some() && Some(pane) != self.maximized {
+                    return None;
+                }
+
                 content.overlay(state, layout, renderer, translation)
             })
             .collect::<Vec<_>>();
@@ -1136,52 +1192,6 @@ fn hovered_split<'a>(
     })
 }
 
-/// The visible contents of the [`PaneGrid`]
-#[derive(Debug)]
-pub enum Contents<'a, T> {
-    /// All panes are visible
-    All(Vec<(Pane, T)>, &'a state::Internal),
-    /// A maximized pane is visible
-    Maximized(Pane, T, Node),
-}
-
-impl<'a, T> Contents<'a, T> {
-    /// Returns the layout [`Node`] of the [`Contents`]
-    pub fn layout(&self) -> &Node {
-        match self {
-            Contents::All(_, state) => state.layout(),
-            Contents::Maximized(_, _, layout) => layout,
-        }
-    }
-
-    /// Returns an iterator over the values of the [`Contents`]
-    pub fn iter(&self) -> Box<dyn Iterator<Item = (Pane, &T)> + '_> {
-        match self {
-            Contents::All(contents, _) => Box::new(
-                contents.iter().map(|(pane, content)| (*pane, content)),
-            ),
-            Contents::Maximized(pane, content, _) => {
-                Box::new(std::iter::once((*pane, content)))
-            }
-        }
-    }
-
-    fn iter_mut(&mut self) -> Box<dyn Iterator<Item = (Pane, &mut T)> + '_> {
-        match self {
-            Contents::All(contents, _) => Box::new(
-                contents.iter_mut().map(|(pane, content)| (*pane, content)),
-            ),
-            Contents::Maximized(pane, content, _) => {
-                Box::new(std::iter::once((*pane, content)))
-            }
-        }
-    }
-
-    fn is_maximized(&self) -> bool {
-        matches!(self, Self::Maximized(..))
-    }
-}
-
 /// The appearance of a [`PaneGrid`].
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub struct Style {
diff --git a/widget/src/pane_grid/state.rs b/widget/src/pane_grid/state.rs
index f5d2bb02..e1934930 100644
--- a/widget/src/pane_grid/state.rs
+++ b/widget/src/pane_grid/state.rs
@@ -343,7 +343,7 @@ impl<T> State<T> {
 /// [`PaneGrid`]: super::PaneGrid
 #[derive(Debug, Clone)]
 pub struct Internal {
-    layout: Node,
+    pub(super) layout: Node,
     last_id: usize,
 }
 
@@ -397,11 +397,12 @@ impl Internal {
 /// The current action of a [`PaneGrid`].
 ///
 /// [`PaneGrid`]: super::PaneGrid
-#[derive(Debug, Clone, Copy, PartialEq)]
+#[derive(Debug, Clone, Copy, PartialEq, Default)]
 pub enum Action {
     /// The [`PaneGrid`] is idle.
     ///
     /// [`PaneGrid`]: super::PaneGrid
+    #[default]
     Idle,
     /// A [`Pane`] in the [`PaneGrid`] is being dragged.
     ///
@@ -441,9 +442,8 @@ impl Action {
     }
 }
 
-impl Internal {
-    /// The layout [`Node`] of the [`Internal`] state
-    pub fn layout(&self) -> &Node {
-        &self.layout
-    }
+#[derive(Default)]
+pub(super) struct Widget {
+    pub action: Action,
+    pub panes: Vec<Pane>,
 }