Fix RefCountedPtr to handle polymorphism.

3 files changed, 61 insertions, 10 deletions

diff --git a/src/core/lib/gprpp/orphanable.h b/src/core/lib/gprpp/orphanable.h
index d0ec9b6461..3123e3f5a3 100644
--- a/src/core/lib/gprpp/orphanable.h
+++ b/src/core/lib/gprpp/orphanable.h
@@ -86,7 +86,8 @@ class InternallyRefCounted : public Orphanable {
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
 
   // Allow RefCountedPtr<> to access Unref() and IncrementRefCount().
-  friend class RefCountedPtr<Child>;
+  template <typename T>
+  friend class RefCountedPtr;
 
   InternallyRefCounted() { gpr_ref_init(&refs_, 1); }
   virtual ~InternallyRefCounted() {}
@@ -129,7 +130,8 @@ class InternallyRefCountedWithTracing : public Orphanable {
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
 
   // Allow RefCountedPtr<> to access Unref() and IncrementRefCount().
-  friend class RefCountedPtr<Child>;
+  template <typename T>
+  friend class RefCountedPtr;
 
   InternallyRefCountedWithTracing()
       : InternallyRefCountedWithTracing(static_cast<TraceFlag*>(nullptr)) {}
diff --git a/src/core/lib/gprpp/ref_counted.h b/src/core/lib/gprpp/ref_counted.h
index ddac5bd475..03c293f6ed 100644
--- a/src/core/lib/gprpp/ref_counted.h
+++ b/src/core/lib/gprpp/ref_counted.h
@@ -73,7 +73,8 @@ class RefCounted {
 
  private:
   // Allow RefCountedPtr<> to access IncrementRefCount().
-  friend class RefCountedPtr<Child>;
+  template <typename T>
+  friend class RefCountedPtr;
 
   void IncrementRefCount() { gpr_ref(&refs_); }
 
@@ -152,7 +153,8 @@ class RefCountedWithTracing {
 
  private:
   // Allow RefCountedPtr<> to access IncrementRefCount().
-  friend class RefCountedPtr<Child>;
+  template <typename T>
+  friend class RefCountedPtr;
 
   void IncrementRefCount() { gpr_ref(&refs_); }
 
diff --git a/src/core/lib/gprpp/ref_counted_ptr.h b/src/core/lib/gprpp/ref_counted_ptr.h
index 534d3d03cb..8a6615a779 100644
--- a/src/core/lib/gprpp/ref_counted_ptr.h
+++ b/src/core/lib/gprpp/ref_counted_ptr.h
@@ -36,7 +36,8 @@ class RefCountedPtr {
   RefCountedPtr(std::nullptr_t) {}
 
   // If value is non-null, we take ownership of a ref to it.
-  explicit RefCountedPtr(T* value) { value_ = value; }
+  template <typename Y>
+  explicit RefCountedPtr(Y* value) { value_ = value; }
 
   // Move support.
   RefCountedPtr(RefCountedPtr&& other) {
@@ -49,6 +50,18 @@ class RefCountedPtr {
     other.value_ = nullptr;
     return *this;
   }
+  template <typename Y>
+  RefCountedPtr(RefCountedPtr<Y>&& other) {
+    value_ = other.value_;
+    other.value_ = nullptr;
+  }
+  template <typename Y>
+  RefCountedPtr& operator=(RefCountedPtr<Y>&& other) {
+    if (value_ != nullptr) value_->Unref();
+    value_ = other.value_;
+    other.value_ = nullptr;
+    return *this;
+  }
 
   // Copy support.
   RefCountedPtr(const RefCountedPtr& other) {
@@ -63,17 +76,37 @@ class RefCountedPtr {
     value_ = other.value_;
     return *this;
   }
+  template <typename Y>
+  RefCountedPtr(const RefCountedPtr<Y>& other) {
+    if (other.value_ != nullptr) other.value_->IncrementRefCount();
+    value_ = other.value_;
+  }
+  template <typename Y>
+  RefCountedPtr& operator=(const RefCountedPtr<Y>& other) {
+    // Note: Order of reffing and unreffing is important here in case value_
+    // and other.value_ are the same object.
+    if (other.value_ != nullptr) other.value_->IncrementRefCount();
+    if (value_ != nullptr) value_->Unref();
+    value_ = other.value_;
+    return *this;
+  }
 
   ~RefCountedPtr() {
     if (value_ != nullptr) value_->Unref();
   }
 
   // If value is non-null, we take ownership of a ref to it.
-  void reset(T* value = nullptr) {
+  template <typename Y>
+  void reset(Y* value) {
     if (value_ != nullptr) value_->Unref();
     value_ = value;
   }
 
+  void reset() {
+    if (value_ != nullptr) value_->Unref();
+    value_ = nullptr;
+  }
+
   // TODO(roth): This method exists solely as a transition mechanism to allow
   // us to pass a ref to idiomatic C code that does not use RefCountedPtr<>.
   // Once all of our code has been converted to idiomatic C++, this
@@ -89,16 +122,30 @@ class RefCountedPtr {
   T& operator*() const { return *value_; }
   T* operator->() const { return value_; }
 
-  bool operator==(const RefCountedPtr& other) const {
+  template <typename Y>
+  bool operator==(const RefCountedPtr<Y>& other) const {
     return value_ == other.value_;
   }
-  bool operator==(const T* other) const { return value_ == other; }
-  bool operator!=(const RefCountedPtr& other) const {
+
+  template <typename Y>
+  bool operator==(const Y* other) const { return value_ == other; }
+
+  bool operator==(std::nullptr_t) const { return value_ == nullptr; }
+
+  template <typename Y>
+  bool operator!=(const RefCountedPtr<Y>& other) const {
     return value_ != other.value_;
   }
-  bool operator!=(const T* other) const { return value_ != other; }
+
+  template <typename Y>
+  bool operator!=(const Y* other) const { return value_ != other; }
+
+  bool operator!=(std::nullptr_t) const { return value_ != nullptr; }
 
  private:
+  template <typename Y>
+  friend class RefCountedPtr;
+
   T* value_ = nullptr;
 };