From 0a0177d0190537f46d39b01b25ebd16c7534c91a Mon Sep 17 00:00:00 2001
From: outfoxxed <outfoxxed@outfoxxed.me>
Date: Sun, 16 Apr 2023 23:53:22 -0700
Subject: [PATCH] Fix wrong window selection on close

---
 src/Hy3Layout.cpp | 29 +++++++++++++++++++++++------
 src/Hy3Layout.hpp |  3 ++-
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/Hy3Layout.cpp b/src/Hy3Layout.cpp
index bbdc85d..de95f6a 100644
--- a/src/Hy3Layout.cpp
+++ b/src/Hy3Layout.cpp
@@ -416,6 +416,7 @@ void Hy3Layout::onWindowCreatedTiling(CWindow* window) {
 	}
 	Debug::log(LOG, "open new window %p(node: %p:%p) on winodow %p in %p", window, &node, node.data.as_window, opening_after, opening_into);
 
+	opening_into->data.as_group.lastFocusedChild = &node;
 	opening_into->recalcSizePosRecursive();
 	Debug::log(LOG, "opening_into (%p) contains new child (%p)? %d", opening_into, &node, opening_into->data.as_group.hasChild(&node));
 }
@@ -440,9 +441,22 @@ void Hy3Layout::onWindowRemovedTiling(CWindow* window) {
 	auto* parent = node->parent;
 	auto* group = &parent->data.as_group;
 
+	if (group->children.size() > 2) {
+		auto iter = std::find(group->children.begin(), group->children.end(), node);
+		if (iter == group->children.begin()) {
+			group->lastFocusedChild = *std::next(iter);
+		} else {
+			group->lastFocusedChild = *std::prev(iter);
+		}
+	}
+
 	group->children.remove(node);
 	this->nodes.remove(*node);
 
+	if (group->children.size() == 1) {
+		group->lastFocusedChild = group->children.front();
+	}
+
 	while (parent->parent != nullptr && group->children.empty()) {
 		auto* child = parent;
 		parent = parent->parent;
@@ -457,23 +471,26 @@ void Hy3Layout::onWindowRemovedTiling(CWindow* window) {
 			}
 		}
 
-		parent->data.as_group.children.remove(child);
+		group->children.remove(child);
+		this->nodes.remove(*child);
 
 		if (group->children.size() == 1) {
 			group->lastFocusedChild = group->children.front();
 		}
-
-		this->nodes.remove(*child);
 	}
 
 	if (parent != nullptr) {
 		parent->recalcSizePosRecursive();
-		auto* focus = group->lastFocusedChild;
-		while (focus->data.type == Hy3NodeData::Group) focus = focus->data.as_group.lastFocusedChild;
-		g_pCompositor->focusWindow(focus->data.as_window);
 	}
 }
 
+CWindow* Hy3Layout::getNextWindowCandidate(CWindow* window) {
+	auto* node = this->getWorkspaceRootGroup(window->m_iWorkspaceID);
+	if (node == nullptr) return nullptr;
+	while (node->data.type == Hy3NodeData::Group) node = node->data.as_group.lastFocusedChild;
+	return node->data.as_window;
+}
+
 void Hy3Layout::onWindowFocusChange(CWindow* window) {
 	Debug::log(LOG, "Switched windows from %p to %p", this->lastActiveWindow, window);
 	auto* node = this->getNodeFromWindow(this->lastActiveWindow);
diff --git a/src/Hy3Layout.hpp b/src/Hy3Layout.hpp
index 6e60653..bbc5c33 100644
--- a/src/Hy3Layout.hpp
+++ b/src/Hy3Layout.hpp
@@ -93,7 +93,8 @@ public:
 	virtual void switchWindows(CWindow*, CWindow*);
 	virtual void alterSplitRatio(CWindow*, float, bool);
 	virtual std::string getLayoutName();
-	virtual void replaceWindowDataWith(CWindow* from, CWindow* to);
+	virtual CWindow* getNextWindowCandidate(CWindow*);
+	virtual void replaceWindowDataWith(CWindow*, CWindow*);
 
 	virtual void onEnable();
 	virtual void onDisable();