Skip to content

Commit 8ab7211

Browse files
Bush2021Bush2021
committed
WIP
1 parent ab5a2f1 commit 8ab7211

3 files changed

Lines changed: 336 additions & 18 deletions

File tree

‎src/tabbookmark.cc‎

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -340,24 +340,10 @@ bool HandleKeepTab(WPARAM wParam) {
340340
return false;
341341
}
342342

343-
HWND hwnd = GetFocus();
344-
if (GetChromeWidgetWin(hwnd) == nullptr) {
345-
return false;
346-
}
347-
348-
if (IsFullScreen(hwnd)) {
349-
// Have to exit full screen to find the tab.
350-
ExecuteCommand(IDC_FULLSCREEN, hwnd);
351-
}
352-
353-
HWND tmp_hwnd = hwnd;
354-
hwnd = GetAncestor(tmp_hwnd, GA_ROOTOWNER);
355-
ExecuteCommand(IDC_CLOSE_FIND_OR_STOP, tmp_hwnd);
356-
357-
NodePtr top_container_view = GetTopContainerView(hwnd);
358-
// Use `GetTabCount` directly since we only need tab count here (no mouse pos)
359-
int tab_count = GetTabCount(top_container_view);
360-
if (!IsNeedKeep(tab_count, KeepTabTrigger::kKeyboardShortcut)) {
343+
HWND hwnd = GetForegroundWindow();
344+
const auto tab_count = FindTabCount(hwnd);
345+
if (tab_count.has_value() &&
346+
!IsNeedKeep(tab_count.value(), KeepTabTrigger::kKeyboardShortcut)) {
361347
return false;
362348
}
363349

‎src/uia.cc‎

Lines changed: 331 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
#include <wrl/client.h>
55

66
#include <algorithm>
7+
#include <functional>
78
#include <initializer_list>
89
#include <optional>
910
#include <string_view>
11+
#include <vector>
1012

1113
#include "utils.h"
1214

@@ -69,6 +71,9 @@ UiaThreadState& GetUiaThreadState() {
6971
struct TabContainer {
7072
ComPtr<IUIAutomationElement> container;
7173
std::wstring_view tab_class;
74+
bool found_via_raw = false;
75+
std::optional<int> raw_tab_count;
76+
int raw_tab_visited = 0;
7277
};
7378

7479
bool EnsureAutomation() {
@@ -208,6 +213,216 @@ ComPtr<IUIAutomationElement> FindFirstInSubtree(
208213
return hit;
209214
}
210215

216+
std::optional<int> CountClassInSubtreeRaw(
217+
const ComPtr<IUIAutomation>& automation,
218+
const ComPtr<IUIAutomationElement>& root,
219+
std::wstring_view class_name,
220+
int visit_limit,
221+
int* visited_out = nullptr) {
222+
if (!automation || !root || class_name.empty() || visit_limit <= 0)
223+
return std::nullopt;
224+
225+
ComPtr<IUIAutomationTreeWalker> walker;
226+
if (FAILED(automation->get_RawViewWalker(&walker)) || !walker)
227+
return std::nullopt;
228+
229+
int count = 0;
230+
int visited = 0;
231+
std::function<void(IUIAutomationElement*)> dfs;
232+
dfs = [&](IUIAutomationElement* node) {
233+
if (!node || visited >= visit_limit)
234+
return;
235+
236+
++visited;
237+
if (IsClassName(ComPtr<IUIAutomationElement>(node), class_name)) {
238+
++count;
239+
}
240+
241+
ComPtr<IUIAutomationElement> child;
242+
if (FAILED(
243+
walker->GetFirstChildElement(node, child.ReleaseAndGetAddressOf())))
244+
return;
245+
246+
while (child && visited < visit_limit) {
247+
dfs(child.Get());
248+
ComPtr<IUIAutomationElement> next;
249+
if (FAILED(walker->GetNextSiblingElement(child.Get(),
250+
next.ReleaseAndGetAddressOf())))
251+
break;
252+
child = std::move(next);
253+
}
254+
};
255+
256+
dfs(root.Get());
257+
if (visited_out)
258+
*visited_out = visited;
259+
return count;
260+
}
261+
262+
ComPtr<IUIAutomationElement> FindRawContainerPreferNonEmpty(
263+
const ComPtr<IUIAutomation>& automation,
264+
const ComPtr<IUIAutomationElement>& root,
265+
std::wstring_view container_class,
266+
std::wstring_view tab_class,
267+
std::optional<int>* selected_raw_count_out = nullptr,
268+
int* selected_raw_visited_out = nullptr) {
269+
if (!automation || !root)
270+
return nullptr;
271+
272+
if (selected_raw_count_out)
273+
*selected_raw_count_out = std::nullopt;
274+
if (selected_raw_visited_out)
275+
*selected_raw_visited_out = 0;
276+
277+
ComPtr<IUIAutomationTreeWalker> walker;
278+
if (FAILED(automation->get_RawViewWalker(&walker)) || !walker)
279+
return nullptr;
280+
281+
ComPtr<IUIAutomationElement> first_candidate;
282+
std::optional<int> first_candidate_raw_count;
283+
int first_candidate_raw_visited = 0;
284+
int candidate_index = 0;
285+
286+
std::function<ComPtr<IUIAutomationElement>(IUIAutomationElement*)> dfs;
287+
dfs = [&](IUIAutomationElement* node) -> ComPtr<IUIAutomationElement> {
288+
if (!node)
289+
return nullptr;
290+
291+
ComPtr<IUIAutomationElement> current(node);
292+
if (IsClassName(current, container_class)) {
293+
++candidate_index;
294+
int visited = 0;
295+
const auto raw_count = CountClassInSubtreeRaw(automation, current,
296+
tab_class, 50000, &visited);
297+
298+
if (!first_candidate) {
299+
first_candidate = current;
300+
first_candidate_raw_count = raw_count;
301+
first_candidate_raw_visited = visited;
302+
}
303+
304+
if (raw_count.has_value()) {
305+
DebugLog(
306+
L"FindTabCount: RawView candidate #{} container='{}' "
307+
L"tab_class='{}' "
308+
L"raw_count={} visited={}",
309+
candidate_index, container_class, tab_class, raw_count.value(),
310+
visited);
311+
if (raw_count.value() > 0) {
312+
if (selected_raw_count_out)
313+
*selected_raw_count_out = raw_count;
314+
if (selected_raw_visited_out)
315+
*selected_raw_visited_out = visited;
316+
return current;
317+
}
318+
} else {
319+
DebugLog(
320+
L"FindTabCount: RawView candidate #{} container='{}' "
321+
L"tab_class='{}' "
322+
L"raw_count=<failed>",
323+
candidate_index, container_class, tab_class);
324+
}
325+
}
326+
327+
ComPtr<IUIAutomationElement> child;
328+
if (FAILED(
329+
walker->GetFirstChildElement(node, child.ReleaseAndGetAddressOf())))
330+
return nullptr;
331+
332+
while (child) {
333+
if (auto hit = dfs(child.Get()))
334+
return hit;
335+
ComPtr<IUIAutomationElement> next;
336+
if (FAILED(walker->GetNextSiblingElement(child.Get(),
337+
next.ReleaseAndGetAddressOf())))
338+
break;
339+
child = std::move(next);
340+
}
341+
return nullptr;
342+
};
343+
344+
if (auto hit = dfs(root.Get())) {
345+
return hit;
346+
}
347+
348+
if (first_candidate) {
349+
DebugLog(
350+
L"FindTabCount: RawView fallback uses first '{}' candidate (all "
351+
L"candidates had raw_count=0)",
352+
container_class);
353+
354+
if (selected_raw_count_out)
355+
*selected_raw_count_out = first_candidate_raw_count;
356+
if (selected_raw_visited_out)
357+
*selected_raw_visited_out = first_candidate_raw_visited;
358+
}
359+
return first_candidate;
360+
}
361+
362+
// 从 HWND 找 tab container，带 fullscreen fallback
363+
std::optional<TabContainer> FindTabContainerFromHwnd(
364+
const ComPtr<IUIAutomation>& automation,
365+
HWND hwnd) {
366+
DebugLog(L"FindTabCount: FindTabContainerFromHwnd start hwnd={:p}",
367+
reinterpret_cast<void*>(hwnd));
368+
369+
auto root = GetElementFromHandle(automation, hwnd);
370+
if (!root) {
371+
DebugLog(L"FindTabCount: ElementFromHandle failed hwnd={:p}",
372+
reinterpret_cast<void*>(hwnd));
373+
return std::nullopt;
374+
}
375+
376+
if (const auto root_class =
377+
GetCurrentStringProperty(root, UIA_ClassNamePropertyId);
378+
root_class.has_value()) {
379+
DebugLog(L"FindTabCount: root class='{}'", root_class.value());
380+
} else {
381+
DebugLog(L"FindTabCount: root class=<unknown>");
382+
}
383+
384+
// horizontal
385+
if (auto c = FindFirstInSubtree(
386+
root, CreateClassCondition(automation, L"TabContainerImpl"))) {
387+
DebugLog(L"FindTabCount: container hit ControlView class=TabContainerImpl");
388+
return TabContainer{c, L"Tab", false};
389+
}
390+
391+
// vertical
392+
if (auto c = FindFirstInSubtree(
393+
root, CreateClassCondition(automation,
394+
L"VerticalUnpinnedTabContainerView"))) {
395+
DebugLog(
396+
L"FindTabCount: container hit ControlView "
397+
L"class=VerticalUnpinnedTabContainerView");
398+
return TabContainer{c, L"VerticalTabView", false};
399+
}
400+
401+
// Fullscreen
402+
DebugLog(L"FindTabCount: ControlView miss, trying RawView fallback");
403+
std::optional<int> raw_tab_count;
404+
int raw_tab_visited = 0;
405+
if (auto c = FindRawContainerPreferNonEmpty(
406+
automation, root, L"TabContainerImpl", L"Tab", &raw_tab_count,
407+
&raw_tab_visited)) {
408+
DebugLog(L"FindTabCount: container hit RawView class=TabContainerImpl");
409+
return TabContainer{c, L"Tab", true, raw_tab_count, raw_tab_visited};
410+
}
411+
if (auto c = FindRawContainerPreferNonEmpty(
412+
automation, root, L"VerticalUnpinnedTabContainerView",
413+
L"VerticalTabView", &raw_tab_count, &raw_tab_visited)) {
414+
DebugLog(
415+
L"FindTabCount: container hit RawView "
416+
L"class=VerticalUnpinnedTabContainerView");
417+
return TabContainer{c, L"VerticalTabView", true, raw_tab_count,
418+
raw_tab_visited};
419+
}
420+
421+
DebugLog(L"FindTabCount: container not found in ControlView or RawView");
422+
423+
return std::nullopt;
424+
}
425+
211426
ComPtr<IUIAutomationElement> FindSiblingByClass(
212427
const ComPtr<IUIAutomation>& automation,
213428
const ComPtr<IUIAutomationElement>& element,
@@ -400,6 +615,122 @@ std::optional<TabHitResult> FindTabHitResult(POINT pt,
400615
return std::nullopt;
401616
}
402617

618+
std::optional<int> FindTabCount(HWND hwnd) {
619+
DebugLog(L"FindTabCount: start hwnd={:p}", reinterpret_cast<void*>(hwnd));
620+
621+
if (!EnsureAutomation()) {
622+
DebugLog(L"FindTabCount: EnsureAutomation failed");
623+
return std::nullopt;
624+
}
625+
const auto& automation = GetUiaThreadState().automation;
626+
627+
const auto tab_container = FindTabContainerFromHwnd(automation, hwnd);
628+
if (!tab_container) {
629+
DebugLog(L"FindTabCount: no tab container");
630+
return std::nullopt;
631+
}
632+
633+
DebugLog(L"FindTabCount: tab class='{}' found_via_raw={}",
634+
tab_container->tab_class,
635+
tab_container->found_via_raw ? L"true" : L"false");
636+
637+
const auto condition =
638+
CreateClassCondition(automation, tab_container->tab_class);
639+
if (!condition) {
640+
DebugLog(L"FindTabCount: CreateClassCondition failed class='{}'",
641+
tab_container->tab_class);
642+
return std::nullopt;
643+
}
644+
645+
ComPtr<IUIAutomationElementArray> arr;
646+
const HRESULT children_hr = tab_container->container->FindAll(
647+
TreeScope_Children, condition.Get(), arr.ReleaseAndGetAddressOf());
648+
if (FAILED(children_hr) || !arr) {
649+
DebugLog(L"FindTabCount: FindAll(TreeScope_Children) failed hr=0x{:08X}",
650+
static_cast<unsigned long>(children_hr));
651+
652+
ComPtr<IUIAutomationElementArray> subtree_arr;
653+
const HRESULT subtree_hr =
654+
tab_container->container->FindAll(TreeScope_Subtree, condition.Get(),
655+
subtree_arr.ReleaseAndGetAddressOf());
656+
if (FAILED(subtree_hr) || !subtree_arr) {
657+
DebugLog(
658+
L"FindTabCount: FindAll(TreeScope_Subtree) also failed hr=0x{:08X}",
659+
static_cast<unsigned long>(subtree_hr));
660+
} else {
661+
int subtree_count = 0;
662+
if (FAILED(subtree_arr->get_Length(&subtree_count))) {
663+
DebugLog(L"FindTabCount: TreeScope_Subtree get_Length failed");
664+
} else {
665+
DebugLog(L"FindTabCount: TreeScope_Subtree count={}", subtree_count);
666+
}
667+
}
668+
return std::nullopt;
669+
}
670+
671+
int count = 0;
672+
if (FAILED(arr->get_Length(&count))) {
673+
DebugLog(L"FindTabCount: TreeScope_Children get_Length failed");
674+
return std::nullopt;
675+
}
676+
677+
DebugLog(L"FindTabCount: TreeScope_Children count={}", count);
678+
679+
ComPtr<IUIAutomationElementArray> subtree_arr;
680+
const HRESULT subtree_hr = tab_container->container->FindAll(
681+
TreeScope_Subtree, condition.Get(), subtree_arr.ReleaseAndGetAddressOf());
682+
if (SUCCEEDED(subtree_hr) && subtree_arr) {
683+
int subtree_count = 0;
684+
if (SUCCEEDED(subtree_arr->get_Length(&subtree_count))) {
685+
DebugLog(L"FindTabCount: TreeScope_Subtree count={}", subtree_count);
686+
}
687+
}
688+
689+
if (count == 0) {
690+
std::optional<int> raw_count = tab_container->raw_tab_count;
691+
int raw_visited = tab_container->raw_tab_visited;
692+
if (!raw_count.has_value()) {
693+
raw_count =
694+
CountClassInSubtreeRaw(automation, tab_container->container,
695+
tab_container->tab_class, 50000, &raw_visited);
696+
}
697+
698+
if (raw_count.has_value()) {
699+
DebugLog(L"FindTabCount: RawWalker count(class='{}')={} visited={}",
700+
tab_container->tab_class, raw_count.value(), raw_visited);
701+
if (raw_count.value() > 0) {
702+
DebugLog(L"FindTabCount: use RawWalker fallback count={}",
703+
raw_count.value());
704+
return raw_count.value();
705+
}
706+
} else {
707+
DebugLog(L"FindTabCount: RawWalker count(class='{}') failed",
708+
tab_container->tab_class);
709+
}
710+
711+
for (const auto probe_class :
712+
{std::wstring_view(L"Tab"), std::wstring_view(L"VerticalTabView"),
713+
std::wstring_view(L"TabSlotView"),
714+
std::wstring_view(L"TabGroupHeader")}) {
715+
if (probe_class == tab_container->tab_class) {
716+
continue;
717+
}
718+
719+
int probe_visited = 0;
720+
const auto probe_count =
721+
CountClassInSubtreeRaw(automation, tab_container->container,
722+
probe_class, 30000, &probe_visited);
723+
if (probe_count.has_value()) {
724+
DebugLog(
725+
L"FindTabCount: RawWalker probe class='{}' count={} visited={}",
726+
probe_class, probe_count.value(), probe_visited);
727+
}
728+
}
729+
}
730+
731+
return count;
732+
}
733+
403734
bool IsOnTabBar(POINT pt) {
404735
if (!EnsureAutomation()) {
405736
return false;

‎src/uia.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ struct TabHitResult {
1515
std::optional<TabHitResult> FindTabHitResult(POINT pt,
1616
bool need_count,
1717
bool need_close_button);
18+
std::optional<int> FindTabCount(HWND hwnd);
1819
bool IsOnTabBar(POINT pt);
1920
bool IsOnBookmark(POINT pt);
2021
bool IsOmniboxFocused();

0 commit comments

Comments
 (0)