From 6ba5f572f3bed501aea97fd3d8ef780fccc70032 Mon Sep 17 00:00:00 2001
From: Aaron Klotz <aklotz@mozilla.com>
Date: Mon, 10 Sep 2018 20:03:10 +0000
Subject: [PATCH] Bug 1489317: Part 2 - Improvements to mscom::AgileReference;
 r=froydnj

This patch adds the definitions of the RefPtr constructor and operator=.
It also refactors some stuff in AgileReference to make these objects easier
to use. Since it's just a bunch of C++ goop, I figured that you'd be fine to
review this. Let me know if you want to add a reviewer who is more familiar
with the COM nuances.

Depends on D5317

Differential Revision: https://phabricator.services.mozilla.com/D5318

--HG--
extra : moz-landing-system : lando
---
 ipc/mscom/AgileReference.cpp | 54 ++++++++++++++++++++++++++++++++---
 ipc/mscom/AgileReference.h   | 55 +++++++++++++++++++++++++++++++++---
 2 files changed, 101 insertions(+), 8 deletions(-)

diff --git a/ipc/mscom/AgileReference.cpp b/ipc/mscom/AgileReference.cpp
index c0fb458e5836..90d62824bbc3 100644
--- a/ipc/mscom/AgileReference.cpp
+++ b/ipc/mscom/AgileReference.cpp
@@ -30,10 +30,34 @@ HRESULT WINAPI RoGetAgileReference(AgileReferenceOptions options,
 namespace mozilla {
 namespace mscom {
 
+AgileReference::AgileReference()
+  : mIid()
+  , mGitCookie(0)
+{
+}
+
 AgileReference::AgileReference(REFIID aIid, IUnknown* aObject)
   : mIid(aIid)
   , mGitCookie(0)
 {
+  AssignInternal(aObject);
+}
+
+void
+AgileReference::Assign(REFIID aIid, IUnknown* aObject)
+{
+  Clear();
+  mIid = aIid;
+  AssignInternal(aObject);
+}
+
+void
+AgileReference::AssignInternal(IUnknown* aObject)
+{
+  // We expect mIid to already be set
+  DebugOnly<IID> zeroIid = {};
+  MOZ_ASSERT(mIid != zeroIid);
+
   /*
    * There are two possible techniques for creating agile references. Starting
    * with Windows 8.1, we may use the RoGetAgileReference API, which is faster.
@@ -46,7 +70,7 @@ AgileReference::AgileReference(REFIID aIid, IUnknown* aObject)
   MOZ_ASSERT(aObject);
 
   if (pRoGetAgileReference &&
-      SUCCEEDED(pRoGetAgileReference(AGILEREFERENCE_DEFAULT, aIid, aObject,
+      SUCCEEDED(pRoGetAgileReference(AGILEREFERENCE_DEFAULT, mIid, aObject,
                                      getter_AddRefs(mAgileRef)))) {
     return;
   }
@@ -57,7 +81,7 @@ AgileReference::AgileReference(REFIID aIid, IUnknown* aObject)
     return;
   }
 
-  DebugOnly<HRESULT> hr = git->RegisterInterfaceInGlobal(aObject, aIid,
+  DebugOnly<HRESULT> hr = git->RegisterInterfaceInGlobal(aObject, mIid,
                                                          &mGitCookie);
   MOZ_ASSERT(SUCCEEDED(hr));
 }
@@ -72,7 +96,16 @@ AgileReference::AgileReference(AgileReference&& aOther)
 
 AgileReference::~AgileReference()
 {
+  Clear();
+}
+
+void
+AgileReference::Clear()
+{
+  mIid = {};
+
   if (!mGitCookie) {
+    mAgileRef = nullptr;
     return;
   }
 
@@ -84,10 +117,23 @@ AgileReference::~AgileReference()
 
   DebugOnly<HRESULT> hr = git->RevokeInterfaceFromGlobal(mGitCookie);
   MOZ_ASSERT(SUCCEEDED(hr));
+  mGitCookie = 0;
+}
+
+AgileReference&
+AgileReference::operator=(AgileReference&& aOther)
+{
+  Clear();
+  mIid = aOther.mIid;
+  aOther.mIid = {};
+  mAgileRef = std::move(aOther.mAgileRef);
+  mGitCookie = aOther.mGitCookie;
+  aOther.mGitCookie = 0;
+  return *this;
 }
 
 HRESULT
-AgileReference::Resolve(REFIID aIid, void** aOutInterface)
+AgileReference::Resolve(REFIID aIid, void** aOutInterface) const
 {
   MOZ_ASSERT(aOutInterface);
   MOZ_ASSERT(mAgileRef || mGitCookie);
@@ -130,7 +176,7 @@ AgileReference::Resolve(REFIID aIid, void** aOutInterface)
   return originalInterface->QueryInterface(aIid, aOutInterface);
 }
 
-IGlobalInterfaceTable*
+/* static */ IGlobalInterfaceTable*
 AgileReference::ObtainGit()
 {
   // Internally to COM, the Global Interface Table is a singleton, therefore we
diff --git a/ipc/mscom/AgileReference.h b/ipc/mscom/AgileReference.h
index e1ef9b3ed350..9b89914433d3 100644
--- a/ipc/mscom/AgileReference.h
+++ b/ipc/mscom/AgileReference.h
@@ -33,9 +33,17 @@ namespace mscom {
  * HRESULT hr = myAgileRef->Resolve(IID_IFoo, getter_AddRefs(foo));
  * // Now foo may be called from the main thread
  */
-class AgileReference
+class AgileReference final
 {
 public:
+  AgileReference();
+
+  template <typename InterfaceT>
+  explicit AgileReference(RefPtr<InterfaceT>& aObject)
+    : AgileReference(__uuidof(InterfaceT), aObject)
+  {
+  }
+
   AgileReference(REFIID aIid, IUnknown* aObject);
   AgileReference(AgileReference&& aOther);
 
@@ -46,14 +54,31 @@ public:
     return mAgileRef || mGitCookie;
   }
 
-  HRESULT Resolve(REFIID aIid, void** aOutInterface);
+  template <typename T>
+  void Assign(const RefPtr<T>& aOther)
+  {
+    Assign(__uuidof(T), aOther);
+  }
+
+  template <typename T>
+  AgileReference& operator=(const RefPtr<T>& aOther)
+  {
+    Assign(aOther);
+    return *this;
+  }
+
+  HRESULT Resolve(REFIID aIid, void** aOutInterface) const;
 
   AgileReference(const AgileReference& aOther) = delete;
   AgileReference& operator=(const AgileReference& aOther) = delete;
-  AgileReference& operator=(AgileReference&& aOther) = delete;
+
+  AgileReference& operator=(AgileReference&& aOther);
 
 private:
-  IGlobalInterfaceTable* ObtainGit();
+  void Assign(REFIID aIid, IUnknown* aObject);
+  void AssignInternal(IUnknown* aObject);
+  void Clear();
+  static IGlobalInterfaceTable* ObtainGit();
 
 private:
   IID                     mIid;
@@ -64,4 +89,26 @@ private:
 } // namespace mscom
 } // namespace mozilla
 
+template <typename T>
+RefPtr<T>::RefPtr(const mozilla::mscom::AgileReference& aAgileRef)
+{
+  void* newRawPtr;
+  if (FAILED(aAgileRef.Resolve(__uuidof(T), &newRawPtr))) {
+    newRawPtr = nullptr;
+  }
+  assign_assuming_AddRef(static_cast<T*>(newRawPtr));
+}
+
+template <typename T>
+RefPtr<T>&
+RefPtr<T>::operator=(const mozilla::mscom::AgileReference& aAgileRef)
+{
+  void* newRawPtr;
+  if (FAILED(aAgileRef.Resolve(__uuidof(T), &newRawPtr))) {
+    newRawPtr = nullptr;
+  }
+  assign_assuming_AddRef(static_cast<T*>(newRawPtr));
+  return *this;
+}
+
 #endif // mozilla_mscom_AgileReference_h