From 5460d5b1193c9e899db9bbe455a2042c57ab4a09 Mon Sep 17 00:00:00 2001
From: Peter Zhu <peter@peterzhu.ca>
Date: Tue, 22 Oct 2024 11:05:51 -0400
Subject: [PATCH] Move error handling for GC.stat_heap to gc.c

---
 gc.c         | 24 +++++++++++++++++++++++-
 gc/default.c | 11 ++++++-----
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/gc.c b/gc.c
index 435e4d1bf4..0adef25312 100644
--- a/gc.c
+++ b/gc.c
@@ -3461,7 +3461,29 @@ gc_stat_heap(rb_execution_context_t *ec, VALUE self, VALUE heap_name, VALUE arg)
         arg = rb_hash_new();
     }
 
-    return rb_gc_impl_stat_heap(rb_gc_get_objspace(), heap_name, arg);
+    if (NIL_P(heap_name)) {
+        if (!RB_TYPE_P(arg, T_HASH)) {
+            rb_raise(rb_eTypeError, "non-hash given");
+        }
+    }
+    else if (FIXNUM_P(heap_name)) {
+        if (!SYMBOL_P(arg) && !RB_TYPE_P(arg, T_HASH)) {
+            rb_raise(rb_eTypeError, "non-hash or symbol given");
+        }
+    }
+    else {
+        rb_raise(rb_eTypeError, "heap_name must be nil or an Integer");
+    }
+
+    VALUE ret = rb_gc_impl_stat_heap(rb_gc_get_objspace(), heap_name, arg);
+
+    if (ret == Qundef) {
+        GC_ASSERT(SYMBOL_P(arg));
+
+        rb_raise(rb_eArgError, "unknown key: %"PRIsVALUE, rb_sym2str(arg));
+    }
+
+    return ret;
 }
 
 static VALUE
diff --git a/gc/default.c b/gc/default.c
index 6b555fbf08..640065b983 100644
--- a/gc/default.c
+++ b/gc/default.c
@@ -7635,8 +7635,9 @@ stat_one_heap(rb_heap_t *heap, VALUE hash, VALUE key)
     SET(total_freed_objects, heap->total_freed_objects);
 #undef SET
 
-    if (!NIL_P(key)) { /* matched key should return above */
-        rb_raise(rb_eArgError, "unknown key: %"PRIsVALUE, rb_sym2str(key));
+    if (!NIL_P(key)) {
+        // Matched key should return above
+        return Qundef;
     }
 
     return hash;
@@ -7651,7 +7652,7 @@ rb_gc_impl_stat_heap(void *objspace_ptr, VALUE heap_name, VALUE hash_or_sym)
 
     if (NIL_P(heap_name)) {
         if (!RB_TYPE_P(hash_or_sym, T_HASH)) {
-            rb_raise(rb_eTypeError, "non-hash given");
+            rb_bug("non-hash given");
         }
 
         for (int i = 0; i < HEAP_COUNT; i++) {
@@ -7678,11 +7679,11 @@ rb_gc_impl_stat_heap(void *objspace_ptr, VALUE heap_name, VALUE hash_or_sym)
             return stat_one_heap(&heaps[heap_idx], hash_or_sym, Qnil);
         }
         else {
-            rb_raise(rb_eTypeError, "non-hash or symbol given");
+            rb_bug("non-hash or symbol given");
         }
     }
     else {
-        rb_raise(rb_eTypeError, "heap_name must be nil or an Integer");
+        rb_bug("heap_name must be nil or an Integer");
     }
 
     return hash_or_sym;