Make Object#singleton_method return methods in modules included in or prepended to singleton class

To simplify the implementation, this makes Object#singleton_method
call the same method called by Object#method (rb_obj_method), then
check that the returned Method is defined before the superclass of the
object's singleton class.  To keep the same error messages, it rescues
exceptions raised by rb_obj_method, and then raises its own exception.

Fixes [Bug #20620]

2 files changed, 65 insertions, 14 deletions

diff --git a/proc.c b/proc.c
index b1956bb346..8e82c13720 100644
--- a/proc.c
+++ b/proc.c
@@ -2046,6 +2046,19 @@ rb_obj_public_method(VALUE obj, VALUE vid)
     return obj_method(obj, vid, TRUE);
 }
 
+static VALUE
+rb_obj_singleton_method_lookup(VALUE arg)
+{
+    VALUE *args = (VALUE *)arg;
+    return rb_obj_method(args[0], args[1]);
+}
+
+static VALUE
+rb_obj_singleton_method_lookup_fail(VALUE arg1, VALUE arg2)
+{
+    return Qfalse;
+}
+
 /*
  *  call-seq:
  *     obj.singleton_method(sym)    -> method
@@ -2073,11 +2086,12 @@ rb_obj_public_method(VALUE obj, VALUE vid)
 VALUE
 rb_obj_singleton_method(VALUE obj, VALUE vid)
 {
-    VALUE klass = rb_singleton_class_get(obj);
+    VALUE sc = rb_singleton_class_get(obj);
+    VALUE klass;
     ID id = rb_check_id(&vid);
 
-    if (NIL_P(klass) ||
-        NIL_P(klass = RCLASS_ORIGIN(klass)) ||
+    if (NIL_P(sc) ||
+        NIL_P(klass = RCLASS_ORIGIN(sc)) ||
         !NIL_P(rb_special_singleton_class(obj))) {
         /* goto undef; */
     }
@@ -2087,21 +2101,26 @@ rb_obj_singleton_method(VALUE obj, VALUE vid)
         /* else goto undef; */
     }
     else {
-        const rb_method_entry_t *me = rb_method_entry_at(klass, id);
-        vid = ID2SYM(id);
-
-        if (UNDEFINED_METHOD_ENTRY_P(me)) {
-            /* goto undef; */
-        }
-        else if (UNDEFINED_REFINED_METHOD_P(me->def)) {
-            /* goto undef; */
-        }
-        else {
-            return mnew_from_me(me, klass, klass, obj, id, rb_cMethod, FALSE);
+        VALUE args[2] = {obj, vid};
+        VALUE ruby_method = rb_rescue(rb_obj_singleton_method_lookup, (VALUE)args, rb_obj_singleton_method_lookup_fail, Qfalse);
+        if (ruby_method) {
+            struct METHOD *method = (struct METHOD *)RTYPEDDATA_GET_DATA(ruby_method);
+            VALUE lookup_class = RBASIC_CLASS(obj);
+            VALUE stop_class = rb_class_superclass(sc);
+            VALUE method_class = method->iclass;
+
+            /* Determine if method is in singleton class, or module included in or prepended to it */
+            do {
+                if (lookup_class == method_class) {
+                    return ruby_method;
+                }
+                lookup_class = RCLASS_SUPER(lookup_class);
+            } while (lookup_class && lookup_class != stop_class);
         }
     }
 
   /* undef: */
+    vid = ID2SYM(id);
     rb_name_err_raise("undefined singleton method '%1$s' for '%2$s'",
                       obj, vid);
     UNREACHABLE_RETURN(Qundef);
diff --git a/test/ruby/test_method.rb b/test/ruby/test_method.rb
index a7945082c2..ebe711ddb4 100644
--- a/test/ruby/test_method.rb
+++ b/test/ruby/test_method.rb
@@ -940,6 +940,38 @@ class TestMethod < Test::Unit::TestCase
     assert_raise(NameError, bug14658) {o.singleton_method(:bar)}
   end
 
+  def test_singleton_method_included_or_prepended_bug_20620
+    m = Module.new do
+      extend self
+      def foo = :foo
+    end
+    assert_equal(:foo, m.singleton_method(:foo).call)
+    assert_raise(NameError) {m.singleton_method(:puts)}
+
+    sc = Class.new do
+      def t = :t
+    end
+    c = Class.new(sc) do
+      singleton_class.prepend(Module.new do
+        def bar = :bar
+      end)
+      extend(Module.new do
+        def quux = :quux
+      end)
+      def self.baz = :baz
+    end
+    assert_equal(:bar, c.singleton_method(:bar).call)
+    assert_equal(:baz, c.singleton_method(:baz).call)
+    assert_equal(:quux, c.singleton_method(:quux).call)
+
+    assert_raise(NameError) {c.singleton_method(:t)}
+
+    c2 = Class.new(c)
+    assert_raise(NameError) {c2.singleton_method(:bar)}
+    assert_raise(NameError) {c2.singleton_method(:baz)}
+    assert_raise(NameError) {c2.singleton_method(:quux)}
+  end
+
   Feature9783 = '[ruby-core:62212] [Feature #9783]'
 
   def assert_curry_three_args(m)