From 403618888c84e6f9c7d6eb34522f6f722f62e4e0 Mon Sep 17 00:00:00 2001
From: Alexander Rosenberg <zanderpkg@pm.me>
Date: Tue, 9 Sep 2025 05:45:57 -0700
Subject: [PATCH] Add debug functions

---
 include/refcount/refcount.h | 18 ++++++++++++++
 src/refcount.c              | 47 ++++++++++++++++++++++++++++++++++++-
 test/test_refcount.c        |  5 ++++
 3 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/include/refcount/refcount.h b/include/refcount/refcount.h
index 698778f..e19e031 100644
--- a/include/refcount/refcount.h
+++ b/include/refcount/refcount.h
@@ -454,6 +454,24 @@ static inline bool refcount_remove_destructor(void *obj, void *key) {
                                               key);
 }
 
+// Debug Functions
+
+uint64_t refcount_debug_context_count_object(const RefcountContext *ctx,
+                                             void *obj, void *target);
+
+/**
+ * Same as #refcount_debug_context_count_object, but only operates on the global
+ * context.
+ * @param obj The root object
+ * @param target The object to look for
+ * @return The number of times the target appeared in the reference tree of the
+ * root
+ */
+static inline uint64_t refcount_debug_count_object(void *obj, void *target) {
+    return refcount_debug_context_count_object(refcount_default_context, obj,
+                                               target);
+}
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/refcount.c b/src/refcount.c
index 8588649..2292ae9 100644
--- a/src/refcount.c
+++ b/src/refcount.c
@@ -250,7 +250,7 @@ end:
  * @param refs Where to store the refs
  * @return True on success
  */
-static inline bool obj_held_refs(RefcountContext *ctx, void *obj,
+static inline bool obj_held_refs(const RefcountContext *ctx, void *obj,
                                  RefcountList **refs) {
     if (ctx->held_refs_callback) {
         return ctx->held_refs_callback(obj, refs, ctx->user_data);
@@ -843,3 +843,48 @@ bool refcount_context_remove_destructor(const RefcountContext *ctx, void *obj,
     unlock_entry_mtx(ENTRY);
     return success;
 }
+
+// Debug Functions
+
+/**
+ * Count all instances of a target object by walking the references of some root
+ * object. This is for debug purposes only. The root is not included in the
+ * count (as in, if `obj == target`, it will not be counted).
+ * @param ctx The #RefcountContext
+ * @param obj The root object
+ * @param target The object to look for
+ * @return The number of times the target appeared in the reference tree of the
+ * root
+ */
+uint64_t refcount_debug_context_count_object(const RefcountContext *ctx,
+                                             void *obj, void *target) {
+    static const HTTableFunctions SEEN_FNS = {
+        .destroy_key = NULL,
+        .destroy_value = NULL,
+        .equal = ht_intptr_equal_callback,
+        .hash = ht_intptr_hash_callback,
+        .user_data = NULL,
+    };
+    if (!obj) {
+        return 0;
+    }
+    RefcountList *queue = NULL;
+    obj_held_refs(ctx, obj, &queue);
+    uint64_t total_count = 0;
+    HTTable *seen = ht_new(&SEEN_FNS, &ctx->ht_alloc, NULL);
+    while (queue) {
+        void *cur = queue->data;
+        queue = refcount_list_pop_full(queue, NULL, &ctx->alloc);
+        // count NULL
+        if (cur == target) {
+            ++total_count;
+        }
+        // but don't try to descend into it
+        if (cur && !ht_has(seen, cur)) {
+            ht_insert(seen, cur, NULL);
+            obj_held_refs(ctx, cur, &queue);
+        }
+    }
+    ht_free(seen);
+    return total_count;
+}
diff --git a/test/test_refcount.c b/test/test_refcount.c
index d27c383..fcbf921 100644
--- a/test/test_refcount.c
+++ b/test/test_refcount.c
@@ -65,6 +65,10 @@ int main(int argc, const char **argv) {
     STATIC_A(1);
     STATIC_A(2);
     STATIC_A(3);
+    assert(refcount_debug_context_count_object(c, &static_a_1, static_a_1.next)
+           == 1);
+    assert(refcount_debug_context_count_object(c, &static_a_1, static_a_2.next)
+           == 0);
 
     A *a = make_a(c, 10, "Hello world\n");
     assert(!refcount_context_is_static(c, a));
@@ -124,6 +128,7 @@ int main(int argc, const char **argv) {
     first->next = a;
     a = first;
     assert(refcount_context_num_refs(c, a) == 1);
+    assert(refcount_debug_context_count_object(c, a, a) == 1);
 
     refcount_context_ref(c, a);
     refcount_context_unref(c, a);