From 667de345f71268a77cf5e531cd40d36ca67aa9f4 Mon Sep 17 00:00:00 2001
From: "Mark D. Roth" <roth@google.com>
Date: Fri, 1 Dec 2023 00:33:46 +0000
Subject: [PATCH] [RefCounted] allow RefCounted<> to work for const types

---
 src/core/lib/gprpp/ref_counted.h    | 48 ++++++++++++++++++++++-------
 test/core/gprpp/ref_counted_test.cc | 11 +++++++
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/src/core/lib/gprpp/ref_counted.h b/src/core/lib/gprpp/ref_counted.h
index cdf692c5ce7..294f0230525 100644
--- a/src/core/lib/gprpp/ref_counted.h
+++ b/src/core/lib/gprpp/ref_counted.h
@@ -219,7 +219,7 @@ class NonPolymorphicRefCount {
 // Default behavior: Delete the object.
 struct UnrefDelete {
   template <typename T>
-  void operator()(T* p) {
+  void operator()(T* p) const {
     delete p;
   }
 };
@@ -231,7 +231,7 @@ struct UnrefDelete {
 // later by identifying entries for which RefIfNonZero() returns null.
 struct UnrefNoDelete {
   template <typename T>
-  void operator()(T* /*p*/) {}
+  void operator()(T* /*p*/) const {}
 };
 
 // Call the object's dtor but do not delete it.  This is useful for cases
@@ -239,7 +239,7 @@ struct UnrefNoDelete {
 // arena).
 struct UnrefCallDtor {
   template <typename T>
-  void operator()(T* p) {
+  void operator()(T* p) const {
     p->~T();
   }
 };
@@ -279,32 +279,44 @@ class RefCounted : public Impl {
   // Note: Depending on the Impl used, this dtor can be implicitly virtual.
   ~RefCounted() = default;
 
+  // Ref() for mutable types.
   GRPC_MUST_USE_RESULT RefCountedPtr<Child> Ref() {
     IncrementRefCount();
     return RefCountedPtr<Child>(static_cast<Child*>(this));
   }
-
   GRPC_MUST_USE_RESULT RefCountedPtr<Child> Ref(const DebugLocation& location,
                                                 const char* reason) {
     IncrementRefCount(location, reason);
     return RefCountedPtr<Child>(static_cast<Child*>(this));
   }
 
+  // Ref() for const types.
+  GRPC_MUST_USE_RESULT RefCountedPtr<const Child> Ref() const {
+    IncrementRefCount();
+    return RefCountedPtr<const Child>(static_cast<const Child*>(this));
+  }
+  GRPC_MUST_USE_RESULT RefCountedPtr<const Child> Ref(
+      const DebugLocation& location, const char* reason) const {
+    IncrementRefCount(location, reason);
+    return RefCountedPtr<const Child>(static_cast<const Child*>(this));
+  }
+
   // TODO(roth): Once all of our code is converted to C++ and can use
   // RefCountedPtr<> instead of manual ref-counting, make this method
   // private, since it will only be used by RefCountedPtr<>, which is a
   // friend of this class.
-  void Unref() {
+  void Unref() const {
     if (GPR_UNLIKELY(refs_.Unref())) {
-      unref_behavior_(static_cast<Child*>(this));
+      unref_behavior_(static_cast<const Child*>(this));
     }
   }
-  void Unref(const DebugLocation& location, const char* reason) {
+  void Unref(const DebugLocation& location, const char* reason) const {
     if (GPR_UNLIKELY(refs_.Unref(location, reason))) {
-      unref_behavior_(static_cast<Child*>(this));
+      unref_behavior_(static_cast<const Child*>(this));
     }
   }
 
+  // RefIfNonZero() for mutable types.
   GRPC_MUST_USE_RESULT RefCountedPtr<Child> RefIfNonZero() {
     return RefCountedPtr<Child>(refs_.RefIfNonZero() ? static_cast<Child*>(this)
                                                      : nullptr);
@@ -316,6 +328,19 @@ class RefCounted : public Impl {
                                     : nullptr);
   }
 
+  // RefIfNonZero() for const types.
+  GRPC_MUST_USE_RESULT RefCountedPtr<const Child> RefIfNonZero() const {
+    return RefCountedPtr<const Child>(refs_.RefIfNonZero()
+                                          ? static_cast<const Child*>(this)
+                                          : nullptr);
+  }
+  GRPC_MUST_USE_RESULT RefCountedPtr<const Child> RefIfNonZero(
+      const DebugLocation& location, const char* reason) const {
+    return RefCountedPtr<const Child>(refs_.RefIfNonZero(location, reason)
+                                          ? static_cast<const Child*>(this)
+                                          : nullptr);
+  }
+
   // Not copyable nor movable.
   RefCounted(const RefCounted&) = delete;
   RefCounted& operator=(const RefCounted&) = delete;
@@ -336,12 +361,13 @@ class RefCounted : public Impl {
   template <typename T>
   friend class RefCountedPtr;
 
-  void IncrementRefCount() { refs_.Ref(); }
-  void IncrementRefCount(const DebugLocation& location, const char* reason) {
+  void IncrementRefCount() const { refs_.Ref(); }
+  void IncrementRefCount(const DebugLocation& location, const char* reason)
+      const {
     refs_.Ref(location, reason);
   }
 
-  RefCount refs_;
+  mutable RefCount refs_;
   GPR_NO_UNIQUE_ADDRESS UnrefBehavior unref_behavior_;
 };
 
diff --git a/test/core/gprpp/ref_counted_test.cc b/test/core/gprpp/ref_counted_test.cc
index 4d8761ecb1d..7c28cddc6f6 100644
--- a/test/core/gprpp/ref_counted_test.cc
+++ b/test/core/gprpp/ref_counted_test.cc
@@ -53,6 +53,17 @@ TEST(RefCounted, ExtraRef) {
   foo->Unref();
 }
 
+TEST(RefCounted, Const) {
+  const Foo* foo = new Foo();
+  RefCountedPtr<const Foo> foop = foo->Ref();
+  foop.release();
+  foop = foo->RefIfNonZero();
+  foop.release();
+  foo->Unref();
+  foo->Unref();
+  foo->Unref();
+}
+
 class Value : public RefCounted<Value, PolymorphicRefCount, UnrefNoDelete> {
  public:
   Value(int value, std::set<std::unique_ptr<Value>>* registry) : value_(value) {