Make Module#{public,private,protected,module_function} return arguments

Previously, each of these methods returned self, but it is
more useful to return arguments, to allow for simpler method
decorators, such as:

```ruby
cached private def foo; some_long_calculation; end
```

Where cached sets up caching for the method.

For each of these methods, the following behavior is used:

1) No arguments returns nil
2) Single argument is returned
3) Multiple arguments are returned as an array

The single argument case is really the case we are trying to
optimize for, for the same reason that def was changed to return
a symbol for the method.

Idea and initial patch from Herwin Quarantainenet.

Implements [Feature #12495]

8 files changed, 139 insertions, 37 deletions

diff --git a/spec/ruby/core/main/private_spec.rb b/spec/ruby/core/main/private_spec.rb
index 78c5d287d4..cac0645b40 100644
--- a/spec/ruby/core/main/private_spec.rb
+++ b/spec/ruby/core/main/private_spec.rb
@@ -32,8 +32,16 @@ describe "main#private" do
     end
   end
 
-  it "returns Object" do
-    eval("private :main_public_method", TOPLEVEL_BINDING).should equal(Object)
+  ruby_version_is ''...'3.1' do
+    it "returns Object" do
+      eval("private :main_public_method", TOPLEVEL_BINDING).should equal(Object)
+    end
+  end
+
+  ruby_version_is '3.1' do
+    it "returns argument" do
+      eval("private :main_public_method", TOPLEVEL_BINDING).should equal(:main_public_method)
+    end
   end
 
   it "raises a NameError when at least one of given method names is undefined" do
diff --git a/spec/ruby/core/main/public_spec.rb b/spec/ruby/core/main/public_spec.rb
index bfc27a9e80..91f045dbab 100644
--- a/spec/ruby/core/main/public_spec.rb
+++ b/spec/ruby/core/main/public_spec.rb
@@ -32,10 +32,19 @@ describe "main#public" do
     end
   end
 
-  it "returns Object" do
-    eval("public :main_private_method", TOPLEVEL_BINDING).should equal(Object)
+  ruby_version_is ''...'3.1' do
+    it "returns Object" do
+      eval("public :main_private_method", TOPLEVEL_BINDING).should equal(Object)
+    end
   end
 
+  ruby_version_is '3.1' do
+    it "returns argument" do
+      eval("public :main_private_method", TOPLEVEL_BINDING).should equal(:main_private_method)
+    end
+  end
+
+
   it "raises a NameError when given an undefined name" do
     -> do
       eval "public :main_undefined_method", TOPLEVEL_BINDING
diff --git a/spec/ruby/core/module/module_function_spec.rb b/spec/ruby/core/module/module_function_spec.rb
index 407237d48f..0602e95ca9 100644
--- a/spec/ruby/core/module/module_function_spec.rb
+++ b/spec/ruby/core/module/module_function_spec.rb
@@ -38,14 +38,23 @@ describe "Module#module_function with specific method names" do
     m.respond_to?(:test3).should == false
   end
 
-  it "returns the current module" do
-    x = nil
-    m = Module.new do
-      def test()  end
-      x = module_function :test
+  ruby_version_is ""..."3.1" do
+    it "returns self" do
+      Module.new do
+        def foo; end
+        module_function(:foo).should equal(self)
+      end
     end
+  end
 
-    x.should == m
+  ruby_version_is "3.1" do
+    it "returns argument or arguments if given" do
+      Module.new do
+        def foo; end
+        module_function(:foo).should equal(:foo)
+        module_function(:foo, :foo).should == [:foo, :foo]
+      end
+    end
   end
 
   it "creates an independent copy of the method, not a redirect" do
@@ -160,13 +169,20 @@ describe "Module#module_function as a toggle (no arguments) in a Module body" do
     m.respond_to?(:test2).should == true
   end
 
-  it "returns the current module" do
-    x = nil
-    m = Module.new {
-      x = module_function
-    }
+  ruby_version_is ""..."3.1" do
+    it "returns self" do
+      Module.new do
+        module_function.should equal(self)
+      end
+    end
+  end
 
-    x.should == m
+  ruby_version_is "3.1" do
+    it "returns nil" do
+      Module.new do
+        module_function.should equal(nil)
+      end
+    end
   end
 
   it "stops creating module functions if the body encounters another toggle " \
diff --git a/spec/ruby/core/module/private_spec.rb b/spec/ruby/core/module/private_spec.rb
index e893c24f38..ead806637c 100644
--- a/spec/ruby/core/module/private_spec.rb
+++ b/spec/ruby/core/module/private_spec.rb
@@ -38,11 +38,25 @@ describe "Module#private" do
                   :module_specs_public_method_on_object_for_kernel_private)
   end
 
-  it "returns self" do
-    (class << Object.new; self; end).class_eval do
-      def foo; end
-      private(:foo).should equal(self)
-      private.should equal(self)
+  ruby_version_is ""..."3.1" do
+    it "returns self" do
+      (class << Object.new; self; end).class_eval do
+        def foo; end
+        private(:foo).should equal(self)
+        private.should equal(self)
+      end
+    end
+  end
+
+  ruby_version_is "3.1" do
+    it "returns argument or arguments if given" do
+      (class << Object.new; self; end).class_eval do
+        def foo; end
+        private(:foo).should equal(:foo)
+        private([:foo, :foo]).should == [:foo, :foo]
+        private(:foo, :foo).should == [:foo, :foo]
+        private.should equal(nil)
+      end
     end
   end
 
diff --git a/spec/ruby/core/module/protected_spec.rb b/spec/ruby/core/module/protected_spec.rb
index aa04a42fb8..058d49d751 100644
--- a/spec/ruby/core/module/protected_spec.rb
+++ b/spec/ruby/core/module/protected_spec.rb
@@ -39,11 +39,25 @@ describe "Module#protected" do
                   :module_specs_public_method_on_object_for_kernel_protected)
   end
 
-  it "returns self" do
-    (class << Object.new; self; end).class_eval do
-      def foo; end
-      protected(:foo).should equal(self)
-      protected.should equal(self)
+  ruby_version_is ""..."3.1" do
+    it "returns self" do
+      (class << Object.new; self; end).class_eval do
+        def foo; end
+        protected(:foo).should equal(self)
+        protected.should equal(self)
+      end
+    end
+  end
+
+  ruby_version_is "3.1" do
+    it "returns argument or arguments if given" do
+      (class << Object.new; self; end).class_eval do
+        def foo; end
+        protected(:foo).should equal(:foo)
+        protected([:foo, :foo]).should == [:foo, :foo]
+        protected(:foo, :foo).should == [:foo, :foo]
+        protected.should equal(nil)
+      end
     end
   end
 
diff --git a/spec/ruby/core/module/public_spec.rb b/spec/ruby/core/module/public_spec.rb
index e7059f6aa6..e3b183f228 100644
--- a/spec/ruby/core/module/public_spec.rb
+++ b/spec/ruby/core/module/public_spec.rb
@@ -27,12 +27,25 @@ describe "Module#public" do
                   :module_specs_private_method_on_object_for_kernel_public)
   end
 
-  it "returns self" do
-    (class << Object.new; self; end).class_eval do
-      def foo; end
-      private :foo
-      public(:foo).should equal(self)
-      public.should equal(self)
+  ruby_version_is ""..."3.1" do
+    it "returns self" do
+      (class << Object.new; self; end).class_eval do
+        def foo; end
+        public(:foo).should equal(self)
+        public.should equal(self)
+      end
+    end
+  end
+
+  ruby_version_is "3.1" do
+    it "returns argument or arguments if given" do
+      (class << Object.new; self; end).class_eval do
+        def foo; end
+        public(:foo).should equal(:foo)
+        public([:foo, :foo]).should == [:foo, :foo]
+        public(:foo, :foo).should == [:foo, :foo]
+        public.should equal(nil)
+      end
     end
   end
 
diff --git a/test/ruby/test_module.rb b/test/ruby/test_module.rb
index be23b84c46..e1524a5d81 100644
--- a/test/ruby/test_module.rb
+++ b/test/ruby/test_module.rb
@@ -1014,6 +1014,28 @@ class TestModule < Test::Unit::TestCase
     assert_raise(NoMethodError, /protected method/) {o.aClass2}
   end
 
+  def test_visibility_method_return_value
+    no_arg_results = nil
+    c = Module.new do
+      singleton_class.send(:public, :public, :private, :protected, :module_function)
+      def foo; end
+      def bar; end
+      no_arg_results = [public, private, protected, module_function]
+    end
+
+    assert_equal([nil]*4, no_arg_results)
+
+    assert_equal(:foo, c.private(:foo))
+    assert_equal(:foo, c.public(:foo))
+    assert_equal(:foo, c.protected(:foo))
+    assert_equal(:foo, c.module_function(:foo))
+
+    assert_equal([:foo, :bar], c.private(:foo, :bar))
+    assert_equal([:foo, :bar], c.public(:foo, :bar))
+    assert_equal([:foo, :bar], c.protected(:foo, :bar))
+    assert_equal([:foo, :bar], c.module_function(:foo, :bar))
+  end
+
   def test_s_constants
     c1 = Module.constants
     Object.module_eval "WALTER = 99"
diff --git a/vm_method.c b/vm_method.c
index 337d326092..9d0c59c01e 100644
--- a/vm_method.c
+++ b/vm_method.c
@@ -2118,11 +2118,14 @@ set_visibility(int argc, const VALUE *argv, VALUE module, rb_method_visibility_t
     if (argc == 0) {
         scope_visibility_check();
         rb_scope_visibility_set(visi);
+        return Qnil;
     }
-    else {
-	set_method_visibility(module, argc, argv, visi);
+
+    set_method_visibility(module, argc, argv, visi);
+    if (argc == 1) {
+        return argv[0];
     }
-    return module;
+    return rb_ary_new_from_values(argc, argv);
 }
 
 /*
@@ -2469,7 +2472,7 @@ rb_mod_modfunc(int argc, VALUE *argv, VALUE module)
 
     if (argc == 0) {
 	rb_scope_module_func_set();
-	return module;
+        return Qnil;
     }
 
     set_method_visibility(module, argc, argv, METHOD_VISI_PRIVATE);
@@ -2495,7 +2498,10 @@ rb_mod_modfunc(int argc, VALUE *argv, VALUE module)
 	}
 	rb_method_entry_set(rb_singleton_class(module), id, me, METHOD_VISI_PUBLIC);
     }
-    return module;
+    if (argc == 1) {
+        return argv[0];
+    }
+    return rb_ary_new_from_values(argc, argv);
 }
 
 #ifdef __GNUC__