[PATCH 3/3] riched20: Add IID_ITextDocument2 support for ITextServices::QueryInterface().

[Jactry Zeng]

Wine-Bug: https://bugs.winehq.org/show_bug.cgi?id=20613
Signed-off-by: Jactry Zeng <jzeng(a)codeweavers.com>
---
 dlls/riched20/editor.h       |  1 +
 dlls/riched20/richole.c      |  6 ++++++
 dlls/riched20/tests/txtsrv.c | 34 ++++++++++++++++++++++++++++++++++
 dlls/riched20/txtsrv.c       |  5 ++++-
 4 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/dlls/riched20/editor.h b/dlls/riched20/editor.h
index 141c63aca9..4c2d71242b 100644
--- a/dlls/riched20/editor.h
+++ b/dlls/riched20/editor.h
@@ -234,6 +234,7 @@ void ME_GetOLEObjectSize(const ME_Context *c, ME_Run *run, SIZE *pSize) DECLSPEC
 void ME_CopyReObject(REOBJECT *dst, const REOBJECT *src, DWORD flags) DECLSPEC_HIDDEN;
 void ME_DeleteReObject(struct re_object *re_object) DECLSPEC_HIDDEN;
 void ME_GetITextDocumentInterface(IRichEditOle *iface, LPVOID *ppvObj) DECLSPEC_HIDDEN;
+void ME_GetITextDocument2Interface(IRichEditOle *iface, LPVOID *ppvObj) DECLSPEC_HIDDEN;
 
 /* editor.c */
 ME_TextEditor *ME_MakeEditor(ITextHost *texthost, BOOL bEmulateVersion10) DECLSPEC_HIDDEN;
diff --git a/dlls/riched20/richole.c b/dlls/riched20/richole.c
index 6c21a8721c..57ab3d0b31 100644
--- a/dlls/riched20/richole.c
+++ b/dlls/riched20/richole.c
@@ -5987,3 +5987,9 @@ void ME_GetITextDocumentInterface(IRichEditOle *iface, LPVOID *ppvObj)
     IRichEditOleImpl *This = impl_from_IRichEditOle(iface);
     *ppvObj = &This->ITextDocument_iface;
 }
+
+void ME_GetITextDocument2Interface(IRichEditOle *iface, LPVOID *ppvObj)
+{
+    IRichEditOleImpl *This = impl_from_IRichEditOle(iface);
+    *ppvObj = &This->ITextDocument2_iface;
+}
diff --git a/dlls/riched20/tests/txtsrv.c b/dlls/riched20/tests/txtsrv.c
index a8cf6ff1ad..77b5857165 100644
--- a/dlls/riched20/tests/txtsrv.c
+++ b/dlls/riched20/tests/txtsrv.c
@@ -904,6 +904,7 @@ static void test_QueryInterface(void)
     HRESULT hres;
     IRichEditOle *reole, *txtsrv_reole;
     ITextDocument *txtdoc, *txtsrv_txtdoc;
+    ITextDocument2 *txtdoc2, *txtsrv_txtdoc2;
     ULONG refcount;
 
     if(!init_texthost(&txtserv, &host))
@@ -930,6 +931,17 @@ static void test_QueryInterface(void)
     ITextDocument_Release(txtdoc);
     refcount = get_refcount((IUnknown *)txtserv);
     ok(refcount == 2, "got wrong ref count: %d\n", refcount);
+
+    hres = IRichEditOle_QueryInterface(txtsrv_reole, &IID_ITextDocument2, (void **)&txtdoc2);
+    ok(hres == S_OK, "IRichEditOle_QueryInterface: 0x%08x\n", hres);
+    refcount = get_refcount((IUnknown *)txtserv);
+    ok(refcount == 3, "got wrong ref count: %d\n", refcount);
+    refcount = get_refcount((IUnknown *)txtsrv_reole);
+    ok(refcount == 3, "got wrong ref count: %d\n", refcount);
+
+    ITextDocument2_Release(txtdoc2);
+    refcount = get_refcount((IUnknown *)txtserv);
+    ok(refcount == 2, "got wrong ref count: %d\n", refcount);
     IRichEditOle_Release(txtsrv_reole);
     refcount = get_refcount((IUnknown *)txtserv);
     ok(refcount == 1, "got wrong ref count: %d\n", refcount);
@@ -956,6 +968,28 @@ static void test_QueryInterface(void)
     refcount = get_refcount((IUnknown *)txtserv);
     ok(refcount == 1, "got wrong ref count: %d\n", refcount);
 
+    /* ITextDocument2 */
+    hres = ITextServices_QueryInterface(txtserv, &IID_ITextDocument2, (void **)&txtsrv_txtdoc2);
+    ok(hres == S_OK, "ITextServices_QueryInterface: 0x%08x\n", hres);
+    refcount = get_refcount((IUnknown *)txtserv);
+    ok(refcount == 2, "got wrong ref count: %d\n", refcount);
+    refcount = get_refcount((IUnknown *)txtsrv_txtdoc2);
+    ok(refcount == 2, "got wrong ref count: %d\n", refcount);
+
+    hres = ITextDocument2_QueryInterface(txtsrv_txtdoc2, &IID_IRichEditOle, (void **)&reole);
+    ok(hres == S_OK, "ITextDocument2_QueryInterface: 0x%08x\n", hres);
+    refcount = get_refcount((IUnknown *)txtserv);
+    ok(refcount == 3, "got wrong ref count: %d\n", refcount);
+    refcount = get_refcount((IUnknown *)txtsrv_txtdoc2);
+    ok(refcount == 3, "got wrong ref count: %d\n", refcount);
+
+    IRichEditOle_Release(reole);
+    refcount = get_refcount((IUnknown *)txtserv);
+    ok(refcount == 2, "got wrong ref count: %d\n", refcount);
+    ITextDocument2_Release(txtsrv_txtdoc2);
+    refcount = get_refcount((IUnknown *)txtserv);
+    ok(refcount == 1, "got wrong ref count: %d\n", refcount);
+
     ITextServices_Release(txtserv);
     ITextHost_Release(host);
 }
diff --git a/dlls/riched20/txtsrv.c b/dlls/riched20/txtsrv.c
index ee65621b1d..46c3a3c5e4 100644
--- a/dlls/riched20/txtsrv.c
+++ b/dlls/riched20/txtsrv.c
@@ -78,12 +78,15 @@ static HRESULT WINAPI ITextServicesImpl_QueryInterface(IUnknown *iface, REFIID r
       *ppv = &This->IUnknown_inner;
    else if (IsEqualIID(riid, &IID_ITextServices))
       *ppv = &This->ITextServices_iface;
-   else if (IsEqualIID(riid, &IID_IRichEditOle) || IsEqualIID(riid, &IID_ITextDocument)) {
+   else if (IsEqualIID(riid, &IID_IRichEditOle) || IsEqualIID(riid, &IID_ITextDocument) ||
+            IsEqualIID(riid, &IID_ITextDocument2)) {
       if (!This->editor->reOle)
          if (!CreateIRichEditOle(This->outer_unk, This->editor, (void **)(&This->editor->reOle)))
             return E_OUTOFMEMORY;
       if (IsEqualIID(riid, &IID_ITextDocument))
          ME_GetITextDocumentInterface(This->editor->reOle, ppv);
+      else if (IsEqualIID(riid, &IID_ITextDocument2))
+         ME_GetITextDocument2Interface(This->editor->reOle, ppv);
       else
          *ppv = This->editor->reOle;
    } else {
-- 
2.18.0