Mojo C++ Bindings: Support custom impl ref types

mojo/public/cpp/bindings/associated_binding.h

Index: mojo/public/cpp/bindings/associated_binding.h
diff --git a/mojo/public/cpp/bindings/associated_binding.h b/mojo/public/cpp/bindings/associated_binding.h
index a7a7c822869b414eb33fd63bde32f5a3e1f89de8..05d80f3c740c65bc3911fba7a6f49f9b1288125e 100644
--- a/mojo/public/cpp/bindings/associated_binding.h
+++ b/mojo/public/cpp/bindings/associated_binding.h
@@ -22,6 +22,7 @@
 #include "mojo/public/cpp/bindings/connection_error_callback.h"
 #include "mojo/public/cpp/bindings/interface_endpoint_client.h"
 #include "mojo/public/cpp/bindings/lib/control_message_proxy.h"
+#include "mojo/public/cpp/bindings/raw_ptr_impl_ref_traits.h"
 #include "mojo/public/cpp/bindings/scoped_interface_endpoint_handle.h"
 
 namespace mojo {
@@ -38,37 +39,38 @@ class MessageReceiver;
 // single thread for the purposes of task scheduling. Please note that incoming
 // synchrounous method calls may not be run from this task runner, when they
 // reenter outgoing synchrounous calls on the same thread.
-template <typename Interface>
+template <typename Interface,
+          typename ImplRefTraits = RawPtrImplRefTraits<Interface>>
 class AssociatedBinding {
  public:
+  using ImplPointerType = typename ImplRefTraits::PointerType;
+
   // Constructs an incomplete associated binding that will use the
   // implementation |impl|. It may be completed with a subsequent call to the
   // |Bind| method. Does not take ownership of |impl|, which must outlive this
   // object.
-  explicit AssociatedBinding(Interface* impl) : impl_(impl) {
-    stub_.set_sink(impl_);
-  }
+  explicit AssociatedBinding(ImplPointerType impl) { stub_.set_sink(impl); }
 
   // Constructs a completed associated binding of |impl|. The output |ptr_info|
   // should be passed through the message pipe endpoint referred to by
   // |associated_group| to setup the corresponding asssociated interface
   // pointer. |impl| must outlive this object.
-  AssociatedBinding(Interface* impl,
+  AssociatedBinding(ImplPointerType impl,
                     AssociatedInterfacePtrInfo<Interface>* ptr_info,
                     AssociatedGroup* associated_group,
                     scoped_refptr<base::SingleThreadTaskRunner> runner =
                         base::ThreadTaskRunnerHandle::Get())
-      : AssociatedBinding(impl) {
+      : AssociatedBinding(std::move(impl)) {
     Bind(ptr_info, associated_group, std::move(runner));
   }
 
   // Constructs a completed associated binding of |impl|. |impl| must outlive
   // the binding.
-  AssociatedBinding(Interface* impl,
+  AssociatedBinding(ImplPointerType impl,
                     AssociatedInterfaceRequest<Interface> request,
                     scoped_refptr<base::SingleThreadTaskRunner> runner =
                         base::ThreadTaskRunnerHandle::Get())
-      : AssociatedBinding(impl) {
+      : AssociatedBinding(std::move(impl)) {
     Bind(std::move(request), std::move(runner));
   }
 
@@ -166,7 +168,7 @@ class AssociatedBinding {
   }
 
   // Returns the interface implementation that was previously specified.
-  Interface* impl() { return impl_; }
+  Interface* impl() { return ImplRefTraits::GetRawPointer(&stub_.sink()); }
 
   // Indicates whether the associated binding has been completed.
   bool is_bound() const { return !!endpoint_client_; }
@@ -187,9 +189,7 @@ class AssociatedBinding {
 
  private:
   std::unique_ptr<InterfaceEndpointClient> endpoint_client_;
-
-  typename Interface::Stub_ stub_;
-  Interface* impl_;
+  typename Interface::template Stub_<ImplRefTraits> stub_;
 
   DISALLOW_COPY_AND_ASSIGN(AssociatedBinding);
 };