From 6e4708733cc4ec08670e9ae40f476562eb80fb47 Mon Sep 17 00:00:00 2001
From: "Unknown W. Brackets" <checkins@unknownbrackets.org>
Date: Sun, 29 Aug 2021 17:08:27 -0700
Subject: [PATCH] UI: Restore focus on PSP view as well.

---
 UI/ControlMappingScreen.cpp | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/UI/ControlMappingScreen.cpp b/UI/ControlMappingScreen.cpp
index 3645747dba..3f09c61bdc 100644
--- a/UI/ControlMappingScreen.cpp
+++ b/UI/ControlMappingScreen.cpp
@@ -1009,6 +1009,7 @@ public:
 
 	MockPSP(UI::LayoutParams *layoutParams = nullptr);
 	void SelectButton(int btn);
+	void FocusButton(int btn);
 	float GetPopupOffset();
 
 	UI::Event ButtonClick;
@@ -1055,6 +1056,12 @@ void MockPSP::SelectButton(int btn) {
 	selectedButton_ = btn;
 }
 
+void MockPSP::FocusButton(int btn) {
+	MockButton *view = buttons_[selectedButton_];
+	if (view)
+		view->SetFocus();
+}
+
 float MockPSP::GetPopupOffset() {
 	MockButton *view = buttons_[selectedButton_];
 	if (!view)
@@ -1164,8 +1171,13 @@ void VisualMappingScreen::HandleKeyMapping(KeyDef key) {
 			nextKey_ = VIRTKEY_AXIS_X_MIN;
 		else if (nextKey_ == VIRTKEY_AXIS_X_MIN)
 			nextKey_ = VIRTKEY_AXIS_X_MAX;
-		else
+		else {
+			if (nextKey_ == VIRTKEY_AXIS_X_MAX)
+				psp_->FocusButton(VIRTKEY_AXIS_Y_MAX);
+			else
+				psp_->FocusButton(nextKey_);
 			nextKey_ = 0;
+		}
 	} else if ((size_t)bindAll_ + 1 < bindAllOrder.size()) {
 		bindAll_++;
 		nextKey_ = bindAllOrder[bindAll_];
@@ -1179,6 +1191,9 @@ void VisualMappingScreen::dialogFinished(const Screen *dialog, DialogResult resu
 	if (result == DR_YES && nextKey_ != 0) {
 		MapNext();
 	} else {
+		// This means they canceled.
+		if (nextKey_ != 0)
+			psp_->FocusButton(nextKey_);
 		nextKey_ = 0;
 		bindAll_ = -1;
 		psp_->SelectButton(0);