Resend layouts from Rust when editor is re-mounted on HMR

Author: Keavon

editor/src/messages/layout/layout_message.rs

@@ -8,6 +8,7 @@ pub enum LayoutMessage {
 		layout_target: LayoutTarget,
 		widget_id: WidgetId,
 	},
+	ResendAllLayouts,
 	SendLayout {
 		layout: Layout,
 		layout_target: LayoutTarget,

editor/src/messages/layout/layout_message_handler.rs

@@ -33,6 +33,20 @@ impl MessageHandler<LayoutMessage, LayoutMessageContext<'_>> for LayoutMessageHa
 				// Resend that diff
 				self.send_diff(vec![diff], layout_target, responses, action_input_mapping);
 			}
+			LayoutMessage::ResendAllLayouts => {
+				// Collect non-empty layouts and their indices, then clear the stored copies so diffs compute as full re-sends
+				let layouts_to_resend: Vec<_> = self
+					.layouts
+					.iter_mut()
+					.enumerate()
+					.filter(|(_, layout)| !layout.0.is_empty())
+					.map(|(i, layout)| (LayoutTarget::from(i as u8), std::mem::take(layout)))
+					.collect();
+
+				for (layout_target, layout) in layouts_to_resend {
+					self.diff_and_send_layout_to_frontend(layout_target, layout, responses, action_input_mapping);
+				}
+			}
 			LayoutMessage::SendLayout { layout, layout_target } => {
 				self.diff_and_send_layout_to_frontend(layout_target, layout, responses, action_input_mapping);
 			}

editor/src/messages/layout/utility_types/layout_widget.rs

@@ -20,10 +20,30 @@ impl core::fmt::Display for WidgetId {
 	}
 }
 
-#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
-#[derive(PartialEq, Clone, Debug, Hash, Eq, Copy, serde::Serialize, serde::Deserialize)]
-#[repr(u8)]
-pub enum LayoutTarget {
+macro_rules! define_layout_target {
+	($($(#[$attr:meta])* $variant:ident),* $(,)?) => {
+		#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
+		#[derive(PartialEq, Clone, Debug, Hash, Eq, Copy, serde::Serialize, serde::Deserialize)]
+		#[repr(u8)]
+		pub enum LayoutTarget {
+			$($(#[$attr])* $variant,)*
+			// KEEP THIS ENUM LAST
+			// This is a marker that is used to define an array that is used to hold widgets
+			#[serde(skip)]
+			_LayoutTargetLength,
+		}
+
+		impl From<u8> for LayoutTarget {
+			fn from(value: u8) -> Self {
+				match value {
+					$(x if x == Self::$variant as u8 => Self::$variant,)*
+					_ => panic!("Invalid LayoutTarget discriminant: {value}"),
+				}
+			}
+		}
+	};
+}
+define_layout_target!(
 	/// The spreadsheet panel allows for the visualisation of data in the graph.
 	DataPanel,
 	/// Contains the action buttons at the bottom of the dialog. Must be shown with the `FrontendMessage::DisplayDialog` message.
@@ -58,12 +78,7 @@ pub enum LayoutTarget {
 	WelcomeScreenButtons,
 	/// The color swatch for the working colors and a flip and reset button found at the bottom of the tool shelf.
 	WorkingColors,
-
-	// KEEP THIS ENUM LAST
-	// This is a marker that is used to define an array that is used to hold widgets
-	#[serde(skip)]
-	_LayoutTargetLength,
-}
+);
 
 /// For use by structs that define a UI widget layout by implementing the layout() function belonging to this trait.
 /// The send_layout() function can then be called by other code which is a part of the same struct so as to send the layout to the frontend.

frontend/src/components/Editor.svelte

@@ -46,6 +46,9 @@
 		// Initialize certain setup tasks required by the editor backend to be ready for the user now that the frontend is ready.
 		// The backend handles idempotency, so this is safe to call again during HMR re-mounts.
 		editor.handle.initAfterFrontendReady();
+
+		// Re-send all UI layouts from Rust so the frontend has them after an HMR re-mount
+		editor.handle.resendAllLayouts();
 	});
 
 	onDestroy(() => {

frontend/wasm/src/editor_api.rs

@@ -187,6 +187,12 @@ impl EditorHandle {
 	// the backend from the web frontend.
 	// ========================================================================
 
+	/// Re-sends all UI layouts to the frontend. Called during HMR re-mounts when the frontend has lost its layout state.
+	#[wasm_bindgen(js_name = resendAllLayouts)]
+	pub fn resend_all_layouts(&self) {
+		self.dispatch(LayoutMessage::ResendAllLayouts);
+	}
+
 	#[wasm_bindgen(js_name = initAfterFrontendReady)]
 	pub fn init_after_frontend_ready(&self) {
 		// Enforce idempotency, so if this is called again during an HMR re-mount, we don't initialize the editor backend twice