WIP

Author: Bush2021

src/tabbookmark.cc

@@ -340,24 +340,10 @@ bool HandleKeepTab(WPARAM wParam) {
     return false;
   }
 
-  HWND hwnd = GetFocus();
-  if (GetChromeWidgetWin(hwnd) == nullptr) {
-    return false;
-  }
-
-  if (IsFullScreen(hwnd)) {
-    // Have to exit full screen to find the tab.
-    ExecuteCommand(IDC_FULLSCREEN, hwnd);
-  }
-
-  HWND tmp_hwnd = hwnd;
-  hwnd = GetAncestor(tmp_hwnd, GA_ROOTOWNER);
-  ExecuteCommand(IDC_CLOSE_FIND_OR_STOP, tmp_hwnd);
-
-  NodePtr top_container_view = GetTopContainerView(hwnd);
-  // Use `GetTabCount` directly since we only need tab count here (no mouse pos)
-  int tab_count = GetTabCount(top_container_view);
-  if (!IsNeedKeep(tab_count, KeepTabTrigger::kKeyboardShortcut)) {
+  HWND hwnd = GetForegroundWindow();
+  const auto tab_count = FindTabCount(hwnd);
+  if (tab_count.has_value() &&
+      !IsNeedKeep(tab_count.value(), KeepTabTrigger::kKeyboardShortcut)) {
     return false;
   }
 

src/uia.cc

@@ -7,6 +7,7 @@
 #include <initializer_list>
 #include <optional>
 #include <string_view>
+#include <vector>
 
 #include "utils.h"
 
@@ -208,6 +209,103 @@ ComPtr<IUIAutomationElement> FindFirstInSubtree(
   return hit;
 }
 
+ComPtr<IUIAutomationElement> FindFirstInSubtreeRaw(
+    const ComPtr<IUIAutomation>& automation,
+    const ComPtr<IUIAutomationElement>& root,
+    std::wstring_view class_name) {
+  ComPtr<IUIAutomationTreeWalker> walker;
+  if (FAILED(automation->get_RawViewWalker(&walker)) || !walker) {
+    return nullptr;
+  }
+
+  std::vector<ComPtr<IUIAutomationElement>> stack = {root};
+  while (!stack.empty()) {
+    auto node = std::move(stack.back());
+    stack.pop_back();
+    if (IsClassName(node, class_name)) {
+      return node;
+    }
+    ComPtr<IUIAutomationElement> child;
+    if (FAILED(walker->GetFirstChildElement(node.Get(),
+                                            child.ReleaseAndGetAddressOf()))) {
+      continue;
+    }
+    while (child) {
+      stack.push_back(child);
+      ComPtr<IUIAutomationElement> next;
+      if (FAILED(walker->GetNextSiblingElement(
+              child.Get(), next.ReleaseAndGetAddressOf()))) {
+        break;
+      }
+      child = std::move(next);
+    }
+  }
+  return nullptr;
+}
+
+std::optional<int> CountClassInSubtreeRaw(
+    const ComPtr<IUIAutomation>& automation,
+    const ComPtr<IUIAutomationElement>& root,
+    std::wstring_view class_name) {
+  ComPtr<IUIAutomationTreeWalker> walker;
+  if (FAILED(automation->get_RawViewWalker(&walker)) || !walker) {
+    return std::nullopt;
+  }
+
+  int count = 0;
+  std::vector<ComPtr<IUIAutomationElement>> stack = {root};
+  while (!stack.empty()) {
+    auto node = std::move(stack.back());
+    stack.pop_back();
+    if (IsClassName(node, class_name)) {
+      ++count;
+    }
+    ComPtr<IUIAutomationElement> child;
+    if (FAILED(walker->GetFirstChildElement(node.Get(),
+                                            child.ReleaseAndGetAddressOf()))) {
+      continue;
+    }
+    while (child) {
+      stack.push_back(child);
+      ComPtr<IUIAutomationElement> next;
+      if (FAILED(walker->GetNextSiblingElement(
+              child.Get(), next.ReleaseAndGetAddressOf()))) {
+        break;
+      }
+      child = std::move(next);
+    }
+  }
+  return count;
+}
+
+std::optional<TabContainer> FindTabContainerFromHwnd(
+    const ComPtr<IUIAutomation>& automation,
+    HWND hwnd) {
+  auto root = GetElementFromHandle(automation, hwnd);
+  if (!root) {
+    return std::nullopt;
+  }
+
+  if (auto c = FindFirstInSubtree(
+          root, CreateClassCondition(automation, L"TabContainerImpl")))
+    return TabContainer{c, L"Tab"};
+  if (auto c = FindFirstInSubtree(
+          root, CreateClassCondition(automation,
+                                     L"VerticalUnpinnedTabContainerView"))) {
+    return TabContainer{c, L"VerticalTabView"};
+  }
+
+  if (auto c = FindFirstInSubtreeRaw(automation, root, L"TabContainerImpl")) {
+    return TabContainer{c, L"Tab"};
+  }
+  if (auto c = FindFirstInSubtreeRaw(automation, root,
+                                     L"VerticalUnpinnedTabContainerView")) {
+    return TabContainer{c, L"VerticalTabView"};
+  }
+
+  return std::nullopt;
+}
+
 ComPtr<IUIAutomationElement> FindSiblingByClass(
     const ComPtr<IUIAutomation>& automation,
     const ComPtr<IUIAutomationElement>& element,
@@ -400,6 +498,37 @@ std::optional<TabHitResult> FindTabHitResult(POINT pt,
   return std::nullopt;
 }
 
+std::optional<int> FindTabCount(HWND hwnd) {
+  if (!EnsureAutomation()) {
+    return std::nullopt;
+  }
+  const auto& automation = GetUiaThreadState().automation;
+
+  const auto tab_container = FindTabContainerFromHwnd(automation, hwnd);
+  if (!tab_container) {
+    return std::nullopt;
+  }
+
+  const auto condition =
+      CreateClassCondition(automation, tab_container->tab_class);
+  if (!condition) {
+    return std::nullopt;
+  }
+
+  ComPtr<IUIAutomationElementArray> arr;
+  if (SUCCEEDED(tab_container->container->FindAll(
+          TreeScope_Children, condition.Get(), arr.ReleaseAndGetAddressOf())) &&
+      arr) {
+    int count = 0;
+    if (SUCCEEDED(arr->get_Length(&count)) && count > 0) {
+      return count;
+    }
+  }
+  // Fullscreen
+  return CountClassInSubtreeRaw(automation, tab_container->container,
+                                tab_container->tab_class);
+}
+
 bool IsOnTabBar(POINT pt) {
   if (!EnsureAutomation()) {
     return false;

src/uia.h

@@ -15,6 +15,7 @@ struct TabHitResult {
 std::optional<TabHitResult> FindTabHitResult(POINT pt,
                                              bool need_count,
                                              bool need_close_button);
+std::optional<int> FindTabCount(HWND hwnd);
 bool IsOnTabBar(POINT pt);
 bool IsOnBookmark(POINT pt);
 bool IsOmniboxFocused();