Module#constant_source_location [Feature #10771]

4 files changed, 176 insertions, 0 deletions

diff --git a/constant.h b/constant.h
index 6c8cda08db..f1a13feedb 100644
--- a/constant.h
+++ b/constant.h
@@ -47,5 +47,8 @@ int rb_public_const_defined_at(VALUE klass, ID id);
 int rb_public_const_defined_from(VALUE klass, ID id);
 rb_const_entry_t *rb_const_lookup(VALUE klass, ID id);
 int rb_autoloading_value(VALUE mod, ID id, VALUE *value, rb_const_flag_t *flag);
+VALUE rb_const_source_location(VALUE, ID);
+VALUE rb_const_source_location_from(VALUE, ID);
+VALUE rb_const_source_location_at(VALUE, ID);
 
 #endif /* CONSTANT_H */
diff --git a/object.c b/object.c
index 17964bb5c8..9536cd36f2 100644
--- a/object.c
+++ b/object.c
@@ -2720,6 +2720,105 @@ rb_mod_const_defined(int argc, VALUE *argv, VALUE mod)
     return Qtrue;
 }
 
+static VALUE
+rb_mod_const_source_location(int argc, VALUE *argv, VALUE mod)
+{
+    VALUE name, recur, loc = Qnil;
+    rb_encoding *enc;
+    const char *pbeg, *p, *path, *pend;
+    ID id;
+
+    rb_check_arity(argc, 1, 2);
+    name = argv[0];
+    recur = (argc == 1) ? Qtrue : argv[1];
+
+    if (SYMBOL_P(name)) {
+	if (!rb_is_const_sym(name)) goto wrong_name;
+	id = rb_check_id(&name);
+	if (!id) return Qnil;
+	return RTEST(recur) ? rb_const_source_location(mod, id) : rb_const_source_location_at(mod, id);
+    }
+
+    path = StringValuePtr(name);
+    enc = rb_enc_get(name);
+
+    if (!rb_enc_asciicompat(enc)) {
+	rb_raise(rb_eArgError, "invalid class path encoding (non ASCII)");
+    }
+
+    pbeg = p = path;
+    pend = path + RSTRING_LEN(name);
+
+    if (p >= pend || !*p) {
+      wrong_name:
+	rb_name_err_raise(wrong_constant_name, mod, name);
+    }
+
+    if (p + 2 < pend && p[0] == ':' && p[1] == ':') {
+	mod = rb_cObject;
+	p += 2;
+	pbeg = p;
+    }
+
+    while (p < pend) {
+	VALUE part;
+	long len, beglen;
+
+	while (p < pend && *p != ':') p++;
+
+	if (pbeg == p) goto wrong_name;
+
+	id = rb_check_id_cstr(pbeg, len = p-pbeg, enc);
+	beglen = pbeg-path;
+
+	if (p < pend && p[0] == ':') {
+	    if (p + 2 >= pend || p[1] != ':') goto wrong_name;
+	    p += 2;
+	    pbeg = p;
+	}
+
+	if (!id) {
+	    part = rb_str_subseq(name, beglen, len);
+	    OBJ_FREEZE(part);
+	    if (!rb_is_const_name(part)) {
+		name = part;
+		goto wrong_name;
+	    }
+	    else {
+		return Qnil;
+	    }
+	}
+	if (!rb_is_const_id(id)) {
+	    name = ID2SYM(id);
+	    goto wrong_name;
+	}
+	if (p < pend) {
+	    if (RTEST(recur)) {
+		mod = rb_const_get(mod, id);
+	    }
+	    else {
+		mod = rb_const_get_at(mod, id);
+	    }
+	    if (!RB_TYPE_P(mod, T_MODULE) && !RB_TYPE_P(mod, T_CLASS)) {
+		rb_raise(rb_eTypeError, "%"PRIsVALUE" does not refer to class/module",
+			 QUOTE(name));
+	    }
+	}
+	else {
+	    if (RTEST(recur)) {
+		loc = rb_const_source_location(mod, id);
+	    }
+	    else {
+		loc = rb_const_source_location_at(mod, id);
+	    }
+	    break;
+	}
+	recur = Qfalse;
+    }
+
+    return loc;
+}
+
 /*
  *  call-seq:
  *     obj.instance_variable_get(symbol)    -> obj
@@ -4249,6 +4348,7 @@ InitVM_Object(void)
     rb_define_method(rb_cModule, "const_get", rb_mod_const_get, -1);
     rb_define_method(rb_cModule, "const_set", rb_mod_const_set, 2);
     rb_define_method(rb_cModule, "const_defined?", rb_mod_const_defined, -1);
+    rb_define_method(rb_cModule, "const_source_location", rb_mod_const_source_location, -1);
     rb_define_private_method(rb_cModule, "remove_const",
 			     rb_mod_remove_const, 1); /* in variable.c */
     rb_define_method(rb_cModule, "const_missing",
diff --git a/test/ruby/test_module.rb b/test/ruby/test_module.rb
index 3786eed4a1..604edf5f59 100644
--- a/test/ruby/test_module.rb
+++ b/test/ruby/test_module.rb
@@ -2375,6 +2375,23 @@ class TestModule < Test::Unit::TestCase
     }
   end
 
+  ConstLocation = [__FILE__, __LINE__]
+
+  def test_const_source_location
+    assert_equal(ConstLocation, self.class.const_source_location(:ConstLocation))
+    assert_equal(ConstLocation, self.class.const_source_location("ConstLocation"))
+    assert_equal(ConstLocation, Object.const_source_location("#{self.class.name}::ConstLocation"))
+    assert_raise(TypeError) {
+      self.class.const_source_location(nil)
+    }
+    assert_raise_with_message(NameError, /wrong constant name/) {
+      self.class.const_source_location("xxx")
+    }
+    assert_raise_with_message(TypeError, %r'does not refer to class/module') {
+      self.class.const_source_location("ConstLocation::FILE")
+    }
+  end
+
   private
 
   def assert_top_method_is_private(method)
diff --git a/variable.c b/variable.c
index 296e58a742..9dc1a3052c 100644
--- a/variable.c
+++ b/variable.c
@@ -2483,6 +2483,62 @@ undefined_constant(VALUE mod, VALUE name)
                       mod, name);
 }
 
+static VALUE
+rb_const_location_from(VALUE klass, ID id, int exclude, int recurse, int visibility)
+{
+    while (RTEST(klass)) {
+	rb_const_entry_t *ce;
+
+	while ((ce = rb_const_lookup(klass, id))) {
+	    if (visibility && RB_CONST_PRIVATE_P(ce)) {
+		return Qnil;
+	    }
+	    if (exclude && klass == rb_cObject) {
+		goto not_found;
+	    }
+	    if (NIL_P(ce->file)) return rb_ary_new();
+	    return rb_assoc_new(ce->file, INT2NUM(ce->line));
+	}
+	if (!recurse) break;
+	klass = RCLASS_SUPER(klass);
+    }
+
+  not_found:
+    return Qnil;
+}
+
+static VALUE
+rb_const_location(VALUE klass, ID id, int exclude, int recurse, int visibility)
+{
+    VALUE loc;
+
+    if (klass == rb_cObject) exclude = FALSE;
+    loc = rb_const_location_from(klass, id, exclude, recurse, visibility);
+    if (!NIL_P(loc)) return loc;
+    if (exclude) return loc;
+    if (BUILTIN_TYPE(klass) != T_MODULE) return loc;
+    /* search global const too, if klass is a module */
+    return rb_const_location_from(rb_cObject, id, FALSE, recurse, visibility);
+}
+
+VALUE
+rb_const_source_location_from(VALUE klass, ID id)
+{
+    return rb_const_location(klass, id, TRUE, TRUE, FALSE);
+}
+
+VALUE
+rb_const_source_location(VALUE klass, ID id)
+{
+    return rb_const_location(klass, id, FALSE, TRUE, FALSE);
+}
+
+VALUE
+rb_const_source_location_at(VALUE klass, ID id)
+{
+    return rb_const_location(klass, id, TRUE, FALSE, FALSE);
+}
+
 /*
  *  call-seq:
  *     remove_const(sym)   -> obj