From 05b4582cd7452e928e9bcd9d78c1faa07748d878 Mon Sep 17 00:00:00 2001
From: zhiyuan <32867472+zhiyuang@users.noreply.github.com>
Date: Mon, 5 Feb 2024 03:23:48 +0800
Subject: [PATCH] Properly handle pen tool undo and redo (#1587)

feat: properly handle pen tool undo and redo
---
 .../document/document_message_handler.rs      |  4 +++-
 editor/src/messages/tool/tool_message.rs      |  2 ++
 .../src/messages/tool/tool_message_handler.rs | 21 +++++++++++++++++++
 .../messages/tool/tool_messages/pen_tool.rs   | 19 +++++++++++++++++
 4 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/editor/src/messages/portfolio/document/document_message_handler.rs b/editor/src/messages/portfolio/document/document_message_handler.rs
index 2233f10f..d7b38b2d 100644
--- a/editor/src/messages/portfolio/document/document_message_handler.rs
+++ b/editor/src/messages/portfolio/document/document_message_handler.rs
@@ -630,6 +630,7 @@ impl MessageHandler<DocumentMessage, DocumentInputs<'_>> for DocumentMessageHand
 			Redo => {
 				responses.add(SelectToolMessage::Abort);
 				responses.add(DocumentMessage::DocumentHistoryForward);
+				responses.add(ToolMessage::Redo);
 				responses.add(OverlaysMessage::Draw);
 			}
 			RenameDocument { new_name } => {
@@ -806,10 +807,11 @@ impl MessageHandler<DocumentMessage, DocumentInputs<'_>> for DocumentMessageHand
 			}
 			Undo => {
 				self.undo_in_progress = true;
-				responses.add(BroadcastEvent::ToolAbort);
+
 				responses.add(DocumentMessage::DocumentHistoryBackward);
 				responses.add(OverlaysMessage::Draw);
 				responses.add(DocumentMessage::UndoFinished);
+				responses.add(ToolMessage::Undo);
 			}
 			UndoFinished => self.undo_in_progress = false,
 			UngroupSelectedLayers => {
diff --git a/editor/src/messages/tool/tool_message.rs b/editor/src/messages/tool/tool_message.rs
index f945de8f..0e0ca8ea 100644
--- a/editor/src/messages/tool/tool_message.rs
+++ b/editor/src/messages/tool/tool_message.rs
@@ -126,6 +126,7 @@ pub enum ToolMessage {
 	},
 	DeactivateTools,
 	InitTools,
+	Redo,
 	RefreshToolOptions,
 	ResetColors,
 	SelectPrimaryColor {
@@ -136,6 +137,7 @@ pub enum ToolMessage {
 		color: Color,
 	},
 	SwapColors,
+	Undo,
 	UpdateCursor,
 	UpdateHints,
 }
diff --git a/editor/src/messages/tool/tool_message_handler.rs b/editor/src/messages/tool/tool_message_handler.rs
index 074284e0..b305ae0f 100644
--- a/editor/src/messages/tool/tool_message_handler.rs
+++ b/editor/src/messages/tool/tool_message_handler.rs
@@ -177,6 +177,15 @@ impl MessageHandler<ToolMessage, (&DocumentMessageHandler, DocumentId, &InputPre
 				tool_data.active_tool_mut().process_message(ToolMessage::UpdateHints, responses, &mut data);
 				tool_data.active_tool_mut().process_message(ToolMessage::UpdateCursor, responses, &mut data);
 			}
+			ToolMessage::Redo => {
+				let tool_data = &mut self.tool_state.tool_data;
+				match tool_data.active_tool_type {
+					ToolType::Pen => {
+						responses.add(PenToolMessage::Redo);
+					}
+					_ => {}
+				}
+			}
 			ToolMessage::RefreshToolOptions => {
 				let tool_data = &mut self.tool_state.tool_data;
 				tool_data.tools.get(&tool_data.active_tool_type).unwrap().send_layout(responses, LayoutTarget::ToolOptions);
@@ -221,6 +230,17 @@ impl MessageHandler<ToolMessage, (&DocumentMessageHandler, DocumentId, &InputPre
 
 				document_data.update_working_colors(responses); // TODO: Make this an event
 			}
+			ToolMessage::Undo => {
+				let tool_data = &mut self.tool_state.tool_data;
+				match tool_data.active_tool_type {
+					ToolType::Pen => {
+						responses.add(PenToolMessage::Undo);
+					}
+					_ => {
+						responses.add(BroadcastEvent::ToolAbort);
+					}
+				}
+			}
 
 			// Sub-messages
 			#[remain::unsorted]
@@ -282,6 +302,7 @@ impl MessageHandler<ToolMessage, (&DocumentMessageHandler, DocumentId, &InputPre
 			SelectRandomPrimaryColor,
 			ResetColors,
 			SwapColors,
+			Undo,
 		);
 		list.extend(self.tool_state.tool_data.active_tool().actions());
 		list.extend(self.transform_layer_handler.actions());
diff --git a/editor/src/messages/tool/tool_messages/pen_tool.rs b/editor/src/messages/tool/tool_messages/pen_tool.rs
index 8c837247..2578285b 100644
--- a/editor/src/messages/tool/tool_messages/pen_tool.rs
+++ b/editor/src/messages/tool/tool_messages/pen_tool.rs
@@ -61,6 +61,7 @@ pub enum PenToolMessage {
 		break_handle: Key,
 		lock_angle: Key,
 	},
+	Redo,
 	Undo,
 	UpdateOptions(PenOptionsUpdate),
 }
@@ -198,6 +199,17 @@ struct ModifierState {
 	lock_angle: bool,
 	break_handle: bool,
 }
+
+impl Default for ModifierState {
+	fn default() -> Self {
+		ModifierState {
+			snap_angle: false,
+			lock_angle: false,
+			break_handle: false,
+		}
+	}
+}
+
 #[derive(Clone, Debug, Default)]
 struct PenToolData {
 	weight: f64,
@@ -665,6 +677,7 @@ impl Fsm for PenToolFsmState {
 				PenToolFsmState::DraggingHandle
 			}
 			(PenToolFsmState::PlacingAnchor, PenToolMessage::DragStart) => {
+				responses.add(DocumentMessage::StartTransaction);
 				tool_data.check_break(document, transform, responses);
 				PenToolFsmState::DraggingHandle
 			}
@@ -713,6 +726,12 @@ impl Fsm for PenToolFsmState {
 
 				self
 			}
+			(PenToolFsmState::DraggingHandle | PenToolFsmState::PlacingAnchor, PenToolMessage::Undo) => tool_data
+				.place_anchor(SnapData::new(document, input), transform, input.mouse.position, ModifierState::default(), responses)
+				.unwrap_or(PenToolFsmState::PlacingAnchor),
+			(_, PenToolMessage::Redo) => tool_data
+				.place_anchor(SnapData::new(document, input), transform, input.mouse.position, ModifierState::default(), responses)
+				.unwrap_or(PenToolFsmState::PlacingAnchor),
 			_ => self,
 		}
 	}