From 4ac78850e94639c40a8ca7bf3b8da8cd2d811755 Mon Sep 17 00:00:00 2001
From: Aaron Patterson <tenderlove@ruby-lang.org>
Date: Thu, 25 Mar 2021 16:24:27 -0700
Subject: [PATCH] Add constant location information to classes

This patch teaches classes and modules about their own source location.
So for example, with this patch you can do the following:

```ruby
module A
  class B
  end
end

p A::B.const_source_location
```

This is equivalent to `A.const_source_location(:B)` except that you
don't need a reference to `A` and you don't need to know the name of
`B`.
---
 constant.h               |  1 +
 internal/class.h         |  4 ++++
 object.c                 | 13 ++++++++++++-
 test/ruby/test_module.rb |  4 ++++
 variable.c               | 25 ++++++++++---------------
 5 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/constant.h b/constant.h
index e0d36909e1..8bc68640c2 100644
--- a/constant.h
+++ b/constant.h
@@ -42,6 +42,7 @@ VALUE rb_mod_public_constant(int argc, const VALUE *argv, VALUE obj);
 VALUE rb_mod_deprecate_constant(int argc, const VALUE *argv, VALUE obj);
 void rb_free_const_table(struct rb_id_table *tbl);
 VALUE rb_const_source_location(VALUE, ID);
+VALUE rb_const_source_location_of(VALUE);
 
 MJIT_SYMBOL_EXPORT_BEGIN
 int rb_autoloading_value(VALUE mod, ID id, VALUE *value, rb_const_flag_t *flag);
diff --git a/internal/class.h b/internal/class.h
index 6c03a31a4e..fd1102201f 100644
--- a/internal/class.h
+++ b/internal/class.h
@@ -14,6 +14,7 @@
 #include "internal/serial.h"    /* for rb_serial_t */
 #include "ruby/internal/stdbool.h"     /* for bool */
 #include "ruby/intern.h"        /* for rb_alloc_func_t */
+#include "constant.h"      /* for rb_const_entry_t */
 #include "ruby/ruby.h"          /* for struct RBasic */
 
 #ifdef RCLASS_SUPER
@@ -42,6 +43,9 @@ struct rb_classext_struct {
     struct rb_id_table *cc_tbl; /* ID -> [[ci, cc1], cc2, ...] */
     struct rb_subclass_entry *subclasses;
     struct rb_subclass_entry **parent_subclasses;
+
+    rb_const_entry_t *location;
+
     /**
      * In the case that this is an `ICLASS`, `module_subclasses` points to the link
      * in the module's `subclasses` list that indicates that the klass has been
diff --git a/object.c b/object.c
index 8e4c89b974..35e1887995 100644
--- a/object.c
+++ b/object.c
@@ -2788,7 +2788,18 @@ rb_mod_const_source_location(int argc, VALUE *argv, VALUE mod)
     const char *pbeg, *p, *path, *pend;
     ID id;
 
-    rb_check_arity(argc, 1, 2);
+    rb_check_arity(argc, 0, 2);
+
+    if (argc == 0) {
+        rb_const_entry_t *ce = RCLASS_EXT(mod)->location;
+        if (ce) {
+            return rb_assoc_new(ce->file, INT2NUM(ce->line));
+        }
+        else {
+            return Qnil;
+        }
+    }
+
     name = argv[0];
     recur = (argc == 1) ? Qtrue : argv[1];
 
diff --git a/test/ruby/test_module.rb b/test/ruby/test_module.rb
index 3cd2f04eff..4b091b45fe 100644
--- a/test/ruby/test_module.rb
+++ b/test/ruby/test_module.rb
@@ -2869,6 +2869,8 @@ def foo; bar; end
   end
 
   ConstLocation = [__FILE__, __LINE__]
+  class ConstLocationClass; end
+  module ConstLocationModule; end
 
   def test_const_source_location
     assert_equal(ConstLocation, self.class.const_source_location(:ConstLocation))
@@ -2883,6 +2885,8 @@ def test_const_source_location
     assert_raise_with_message(TypeError, %r'does not refer to class/module') {
       self.class.const_source_location("ConstLocation::FILE")
     }
+    assert_equal(self.class.const_source_location(:ConstLocationClass), ConstLocationClass.const_source_location)
+    assert_equal(self.class.const_source_location(:ConstLocationModule), ConstLocationModule.const_source_location)
   end
 
   module CloneTestM_simple
diff --git a/variable.c b/variable.c
index 85ff35ba8c..2f31f1cd1c 100644
--- a/variable.c
+++ b/variable.c
@@ -2998,8 +2998,6 @@ set_namespace_path(VALUE named_namespace, VALUE namespace_path)
 void
 rb_const_set(VALUE klass, ID id, VALUE val)
 {
-    rb_const_entry_t *ce;
-
     if (NIL_P(klass)) {
 	rb_raise(rb_eTypeError, "no class/module to define constant %"PRIsVALUE"",
 		 QUOTE_ID(id));
@@ -3015,20 +3013,14 @@ rb_const_set(VALUE klass, ID id, VALUE val)
     {
         struct rb_id_table *tbl = RCLASS_CONST_TBL(klass);
         if (!tbl) {
-            RCLASS_CONST_TBL(klass) = tbl = rb_id_table_create(0);
-            rb_clear_constant_cache();
-            ce = ZALLOC(rb_const_entry_t);
-            rb_id_table_insert(tbl, id, (VALUE)ce);
-            setup_const_entry(ce, klass, val, CONST_PUBLIC);
-        }
-        else {
-            struct autoload_const ac = {
-                .mod = klass, .id = id,
-                .value = val, .flag = CONST_PUBLIC,
-                /* fill the rest with 0 */
-            };
-            const_tbl_update(&ac);
+            RCLASS_CONST_TBL(klass) = rb_id_table_create(0);
         }
+        struct autoload_const ac = {
+            .mod = klass, .id = id,
+            .value = val, .flag = CONST_PUBLIC,
+            /* fill the rest with 0 */
+        };
+        const_tbl_update(&ac);
     }
     RB_VM_LOCK_LEAVE();
 
@@ -3129,6 +3121,9 @@ const_tbl_update(struct autoload_const *ac)
 
 	ce = ZALLOC(rb_const_entry_t);
 	rb_id_table_insert(tbl, id, (VALUE)ce);
+	if (RB_TYPE_P(val, T_MODULE) || RB_TYPE_P(val, T_CLASS)) {
+            RCLASS_EXT(val)->location = ce;
+        }
 	setup_const_entry(ce, klass, val, visibility);
     }
 }
-- 
2.30.0