From a86f01b67a150c8b0581f29e90a210eaf1b8b44e Mon Sep 17 00:00:00 2001
From: 0SlowPoke0 <govindpratyush@gmail.com>
Date: Wed, 1 Jan 2025 15:22:45 +0530
Subject: [PATCH] implement restore selection before drag

---
 .../messages/input_mapper/input_mappings.rs   |  1 +
 .../tool/tool_messages/select_tool.rs         | 32 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/editor/src/messages/input_mapper/input_mappings.rs b/editor/src/messages/input_mapper/input_mappings.rs
index 9ff6e41294..38968c1b10 100644
--- a/editor/src/messages/input_mapper/input_mappings.rs
+++ b/editor/src/messages/input_mapper/input_mappings.rs
@@ -96,6 +96,7 @@ pub fn input_mappings() -> Mapping {
 		entry!(PointerMove; refresh_keys=[Control, Shift], action_dispatch=TransformLayerMessage::PointerMove { slow_key: Shift, snap_key: Control }),
 		//
 		// SelectToolMessage
+		entry!(KeyDownNoRepeat(Shift); action_dispatch=SelectToolMessage::RestoreSelection),
 		entry!(PointerMove; refresh_keys=[Control, Alt, Shift], action_dispatch=SelectToolMessage::PointerMove(SelectToolPointerKeys { axis_align: Shift, snap_angle: Control, center: Alt, duplicate: Alt })),
 		entry!(KeyDown(MouseLeft); action_dispatch=SelectToolMessage::DragStart { extend_selection: Shift, select_deepest: Accel }),
 		entry!(KeyUp(MouseLeft); action_dispatch=SelectToolMessage::DragStop { remove_from_selection: Shift, negative_box_selection: Control }),
diff --git a/editor/src/messages/tool/tool_messages/select_tool.rs b/editor/src/messages/tool/tool_messages/select_tool.rs
index 9783aebe69..d38e4f758c 100644
--- a/editor/src/messages/tool/tool_messages/select_tool.rs
+++ b/editor/src/messages/tool/tool_messages/select_tool.rs
@@ -79,6 +79,7 @@ pub enum SelectToolMessage {
 	PointerOutsideViewport(SelectToolPointerKeys),
 	SelectOptions(SelectOptionsUpdate),
 	SetPivot { position: PivotPosition },
+	RestoreSelection,
 }
 
 impl ToolMetadata for SelectTool {
@@ -229,6 +230,7 @@ impl<'a> MessageHandler<ToolMessage, &mut ToolActionHandlerData<'a>> for SelectT
 
 		let additional = match self.fsm_state {
 			SelectToolFsmState::Ready { .. } => actions!(SelectToolMessageDiscriminant; DragStart),
+			SelectToolFsmState::DrawingBox { .. } => actions!(SelectToolMessageDiscriminant;RestoreSelection,DragStop),
 			_ => actions!(SelectToolMessageDiscriminant; DragStop),
 		};
 		common.extend(additional);
@@ -280,6 +282,7 @@ struct SelectToolData {
 	selected_layers_changed: bool,
 	snap_candidates: Vec<SnapCandidatePoint>,
 	auto_panning: AutoPanning,
+	previously_selected_layers: Vec<LayerNodeIdentifier>,
 }
 
 impl SelectToolData {
@@ -547,6 +550,7 @@ impl Fsm for SelectToolFsmState {
 					.collect();
 				let intersection_list = document.click_list(input).collect::<Vec<_>>();
 				let intersection = document.find_deepest(&intersection_list);
+				tool_data.previously_selected_layers = selected.clone();
 
 				// If the user is dragging the bounding box bounds, go into ResizingBounds mode.
 				// If the user is dragging the rotate trigger, go into RotatingBounds mode.
@@ -1096,6 +1100,34 @@ impl Fsm for SelectToolFsmState {
 
 				self
 			}
+			(SelectToolFsmState::DrawingBox { .. }, SelectToolMessage::RestoreSelection) => {
+				if !tool_data.previously_selected_layers.iter().any(|item| tool_data.layers_dragging.contains(item)) {
+					tool_data.layers_dragging.extend(tool_data.previously_selected_layers.clone());
+				} else {
+					let selection_set: HashSet<_> = tool_data.previously_selected_layers.iter().cloned().collect();
+					tool_data.layers_dragging.retain(|layer| !selection_set.contains(layer));
+				}
+				let selected_nodes: Vec<_> = tool_data
+					.layers_dragging
+					.iter()
+					.filter_map(|layer| {
+						if *layer != LayerNodeIdentifier::ROOT_PARENT {
+							Some(layer.to_node())
+						} else {
+							log::error!("ROOT_PARENT found in tool_data.layers_dragging during selection set");
+							None
+						}
+					})
+					.collect();
+
+				// Dispatch the SelectedNodesSet message
+				responses.add(NodeGraphMessage::SelectedNodesSet { nodes: selected_nodes });
+
+				// Dispatch the OverlaysMessage::Draw message to update the UI
+				responses.add(OverlaysMessage::Draw);
+
+				self
+			}
 			_ => self,
 		}
 	}