From 3978ff6538839280543ad4523d7ccdca79b86931 Mon Sep 17 00:00:00 2001
From: Connor McAdams <cmcadams@codeweavers.com>
Date: Thu, 9 Mar 2023 13:22:43 -0500
Subject: [PATCH] uiautomationcore: Implement
 IUIAutomationTreeWalker::Get{First/Last}ChildElement{BuildCache}.

Signed-off-by: Connor McAdams <cmcadams@codeweavers.com>
---
 dlls/uiautomationcore/tests/uiautomation.c | 90 ++++++++++++++++++++++
 dlls/uiautomationcore/uia_com_client.c     | 68 ++++++++++++++--
 2 files changed, 150 insertions(+), 8 deletions(-)

diff --git a/dlls/uiautomationcore/tests/uiautomation.c b/dlls/uiautomationcore/tests/uiautomation.c
index 4f4ce10d622..de60a0e5a1b 100644
--- a/dlls/uiautomationcore/tests/uiautomation.c
+++ b/dlls/uiautomationcore/tests/uiautomation.c
@@ -11701,8 +11701,37 @@ static void test_Element_Find(IUIAutomation *uia_iface)
     UnregisterClassA("test_Element_Find class", NULL);
 }
 
+static const struct prov_method_sequence treewalker_seq1[] = {
+    { &Provider, FRAG_NAVIGATE }, /* NavigateDirection_FirstChild */
+    NODE_CREATE_SEQ(&Provider_child),
+    { 0 }
+};
+
+static const struct prov_method_sequence treewalker_seq2[] = {
+    { &Provider, FRAG_NAVIGATE }, /* NavigateDirection_FirstChild */
+    NODE_CREATE_SEQ(&Provider_child),
+    { &Provider_child, FRAG_GET_RUNTIME_ID },
+    { 0 }
+};
+
+static const struct prov_method_sequence treewalker_seq3[] = {
+    { &Provider, FRAG_NAVIGATE }, /* NavigateDirection_LastChild */
+    NODE_CREATE_SEQ(&Provider_child2),
+    { 0 }
+};
+
+static const struct prov_method_sequence treewalker_seq4[] = {
+    { &Provider, FRAG_NAVIGATE }, /* NavigateDirection_LastChild */
+    NODE_CREATE_SEQ(&Provider_child2),
+    { &Provider_child2, FRAG_GET_RUNTIME_ID },
+    { 0 }
+};
+
 static void test_CUIAutomation_TreeWalker_ifaces(IUIAutomation *uia_iface)
 {
+    HWND hwnd = create_test_hwnd("test_CUIAutomation_TreeWalker_ifaces class");
+    IUIAutomationElement *element, *element2;
+    IUIAutomationCacheRequest *cache_req;
     IUIAutomationCondition *cond, *cond2;
     IUIAutomationTreeWalker *walker;
     HRESULT hr;
@@ -11735,7 +11764,68 @@ static void test_CUIAutomation_TreeWalker_ifaces(IUIAutomation *uia_iface)
     IUIAutomationCondition_Release(cond);
     IUIAutomationCondition_Release(cond2);
 
+    cache_req = NULL;
+    hr = IUIAutomation_CreateCacheRequest(uia_iface, &cache_req);
+    ok(hr == S_OK, "Unexpected hr %#lx.\n", hr);
+    ok(!!cache_req, "cache_req == NULL\n");
+
+    element = create_test_element_from_hwnd(uia_iface, hwnd, TRUE);
+    initialize_provider(&Provider_child, ProviderOptions_ServerSideProvider, NULL, TRUE);
+    initialize_provider(&Provider_child2, ProviderOptions_ServerSideProvider, NULL, TRUE);
+    Provider.frag_root = &Provider.IRawElementProviderFragmentRoot_iface;
+    provider_add_child(&Provider, &Provider_child);
+    provider_add_child(&Provider, &Provider_child2);
+
+    /* NavigateDirection_FirstChild. */
+    element2 = NULL;
+    hr = IUIAutomationTreeWalker_GetFirstChildElement(walker, element, &element2);
+    ok(hr == S_OK, "Unexpected hr %#lx.\n", hr);
+    ok(Provider_child.ref == 2, "Unexpected refcnt %ld\n", Provider_child.ref);
+    ok(!!element2, "element2 == NULL\n");
+    ok_method_sequence(treewalker_seq1, "treewalker_seq1");
+
+    IUIAutomationElement_Release(element2);
+    ok(Provider_child.ref == 1, "Unexpected refcnt %ld\n", Provider_child.ref);
+
+    element2 = NULL;
+    hr = IUIAutomationTreeWalker_GetFirstChildElementBuildCache(walker, element, cache_req, &element2);
+    ok(hr == S_OK, "Unexpected hr %#lx.\n", hr);
+    ok(Provider_child.ref == 2, "Unexpected refcnt %ld\n", Provider_child.ref);
+    ok(!!element2, "element2 == NULL\n");
+    ok_method_sequence(treewalker_seq2, "treewalker_seq2");
+
+    IUIAutomationElement_Release(element2);
+    ok(Provider_child.ref == 1, "Unexpected refcnt %ld\n", Provider_child.ref);
+
+    /* NavigateDirection_LastChild. */
+    element2 = NULL;
+    hr = IUIAutomationTreeWalker_GetLastChildElement(walker, element, &element2);
+    ok(hr == S_OK, "Unexpected hr %#lx.\n", hr);
+    ok(Provider_child2.ref == 2, "Unexpected refcnt %ld\n", Provider_child2.ref);
+    ok(!!element2, "element2 == NULL\n");
+    ok_method_sequence(treewalker_seq3, "treewalker_seq3");
+
+    IUIAutomationElement_Release(element2);
+    ok(Provider_child2.ref == 1, "Unexpected refcnt %ld\n", Provider_child2.ref);
+
+    element2 = NULL;
+    hr = IUIAutomationTreeWalker_GetLastChildElementBuildCache(walker, element, cache_req, &element2);
+    ok(hr == S_OK, "Unexpected hr %#lx.\n", hr);
+    ok(Provider_child2.ref == 2, "Unexpected refcnt %ld\n", Provider_child2.ref);
+    ok(!!element2, "element2 == NULL\n");
+    ok_method_sequence(treewalker_seq4, "treewalker_seq4");
+
+    IUIAutomationElement_Release(element2);
+    ok(Provider_child2.ref == 1, "Unexpected refcnt %ld\n", Provider_child2.ref);
+
+    IUIAutomationElement_Release(element);
+    ok(Provider.ref == 1, "Unexpected refcnt %ld\n", Provider.ref);
+
+    IUIAutomationCacheRequest_Release(cache_req);
     IUIAutomationTreeWalker_Release(walker);
+
+    DestroyWindow(hwnd);
+    UnregisterClassA("test_CUIAutomation_TreeWalker_ifaces class", NULL);
 }
 
 struct uia_com_classes {
diff --git a/dlls/uiautomationcore/uia_com_client.c b/dlls/uiautomationcore/uia_com_client.c
index ec634be5e13..0af6b1f1310 100644
--- a/dlls/uiautomationcore/uia_com_client.c
+++ b/dlls/uiautomationcore/uia_com_client.c
@@ -2449,6 +2449,7 @@ struct uia_tree_walker {
     IUIAutomationTreeWalker IUIAutomationTreeWalker_iface;
     LONG ref;
 
+    IUIAutomationCacheRequest *default_cache_req;
     IUIAutomationCondition *nav_cond;
     struct UiaCondition *cond_struct;
 };
@@ -2486,6 +2487,8 @@ static ULONG WINAPI uia_tree_walker_Release(IUIAutomationTreeWalker *iface)
     TRACE("%p, refcount %ld\n", tree_walker, ref);
     if (!ref)
     {
+        if (tree_walker->default_cache_req)
+            IUIAutomationCacheRequest_Release(tree_walker->default_cache_req);
         IUIAutomationCondition_Release(tree_walker->nav_cond);
         heap_free(tree_walker);
     }
@@ -2503,15 +2506,21 @@ static HRESULT WINAPI uia_tree_walker_GetParentElement(IUIAutomationTreeWalker *
 static HRESULT WINAPI uia_tree_walker_GetFirstChildElement(IUIAutomationTreeWalker *iface, IUIAutomationElement *elem,
         IUIAutomationElement **first)
 {
-    FIXME("%p, %p, %p: stub\n", iface, elem, first);
-    return E_NOTIMPL;
+    struct uia_tree_walker *tree_walker = impl_from_IUIAutomationTreeWalker(iface);
+
+    TRACE("%p, %p, %p\n", iface, elem, first);
+
+    return IUIAutomationTreeWalker_GetFirstChildElementBuildCache(iface, elem, tree_walker->default_cache_req, first);
 }
 
 static HRESULT WINAPI uia_tree_walker_GetLastChildElement(IUIAutomationTreeWalker *iface, IUIAutomationElement *elem,
         IUIAutomationElement **last)
 {
-    FIXME("%p, %p, %p: stub\n", iface, elem, last);
-    return E_NOTIMPL;
+    struct uia_tree_walker *tree_walker = impl_from_IUIAutomationTreeWalker(iface);
+
+    TRACE("%p, %p, %p\n", iface, elem, last);
+
+    return IUIAutomationTreeWalker_GetLastChildElementBuildCache(iface, elem, tree_walker->default_cache_req, last);
 }
 
 static HRESULT WINAPI uia_tree_walker_GetNextSiblingElement(IUIAutomationTreeWalker *iface, IUIAutomationElement *elem,
@@ -2535,6 +2544,40 @@ static HRESULT WINAPI uia_tree_walker_NormalizeElement(IUIAutomationTreeWalker *
     return E_NOTIMPL;
 }
 
+static HRESULT uia_tree_walker_navigate(IUIAutomationTreeWalker *walker, IUIAutomationCacheRequest *cache_req,
+        IUIAutomationElement *start_elem, int nav_dir, IUIAutomationElement **out_elem)
+{
+    struct uia_tree_walker *tree_walker = impl_from_IUIAutomationTreeWalker(walker);
+    struct UiaCacheRequest *cache_req_struct;
+    struct uia_element *element;
+    BSTR tree_struct = NULL;
+    SAFEARRAY *sa = NULL;
+    HRESULT hr;
+
+    if (!out_elem)
+        return E_POINTER;
+
+    *out_elem = NULL;
+    if (!start_elem)
+        return E_POINTER;
+
+    hr = get_uia_cache_request_struct_from_iface(cache_req, &cache_req_struct);
+    if (FAILED(hr))
+        return hr;
+
+    element = impl_from_IUIAutomationElement9((IUIAutomationElement9 *)start_elem);
+    hr = UiaNavigate(element->node, nav_dir, tree_walker->cond_struct, cache_req_struct, &sa, &tree_struct);
+    if (SUCCEEDED(hr) && sa)
+    {
+        hr = create_uia_element_from_cache_req(out_elem, element->from_cui8, cache_req_struct, 0, sa, tree_struct);
+        tree_struct = NULL;
+    }
+
+    SysFreeString(tree_struct);
+    SafeArrayDestroy(sa);
+    return hr;
+}
+
 static HRESULT WINAPI uia_tree_walker_GetParentElementBuildCache(IUIAutomationTreeWalker *iface,
         IUIAutomationElement *elem, IUIAutomationCacheRequest *cache_req, IUIAutomationElement **parent)
 {
@@ -2545,15 +2588,17 @@ static HRESULT WINAPI uia_tree_walker_GetParentElementBuildCache(IUIAutomationTr
 static HRESULT WINAPI uia_tree_walker_GetFirstChildElementBuildCache(IUIAutomationTreeWalker *iface,
         IUIAutomationElement *elem, IUIAutomationCacheRequest *cache_req, IUIAutomationElement **first)
 {
-    FIXME("%p, %p, %p, %p: stub\n", iface, elem, cache_req, first);
-    return E_NOTIMPL;
+    TRACE("%p, %p, %p, %p\n", iface, elem, cache_req, first);
+
+    return uia_tree_walker_navigate(iface, cache_req, elem, NavigateDirection_FirstChild, first);
 }
 
 static HRESULT WINAPI uia_tree_walker_GetLastChildElementBuildCache(IUIAutomationTreeWalker *iface,
         IUIAutomationElement *elem, IUIAutomationCacheRequest *cache_req, IUIAutomationElement **last)
 {
-    FIXME("%p, %p, %p, %p: stub\n", iface, elem, cache_req, last);
-    return E_NOTIMPL;
+    TRACE("%p, %p, %p, %p\n", iface, elem, cache_req, last);
+
+    return uia_tree_walker_navigate(iface, cache_req, elem, NavigateDirection_LastChild, last);
 }
 
 static HRESULT WINAPI uia_tree_walker_GetNextSiblingElementBuildCache(IUIAutomationTreeWalker *iface,
@@ -2636,6 +2681,13 @@ static HRESULT create_uia_tree_walker(IUIAutomationTreeWalker **out_tree_walker,
     IUIAutomationCondition_AddRef(nav_cond);
     tree_walker->cond_struct = cond_struct;
 
+    hr = create_uia_cache_request_iface(&tree_walker->default_cache_req);
+    if (FAILED(hr))
+    {
+        IUIAutomationTreeWalker_Release(&tree_walker->IUIAutomationTreeWalker_iface);
+        return hr;
+    }
+
     *out_tree_walker = &tree_walker->IUIAutomationTreeWalker_iface;
     return S_OK;
 }
-- 
GitLab