Add and use primitive function to change arity

Author: wildmaples

lib/truffle/truffle/cext_ruby.rb

@@ -43,7 +43,12 @@ def rb_define_method(mod, name, function, argc)
       res
     end
 
-    mod.define_method(name, method_body)
+    # Even if the argc is -2, the arity number
+    # is still any number of arguments, -1
+    arity = argc == -2 ? -1 : argc
+
+    method_body_with_arity = Primitive.proc_specify_arity(method_body, arity)
+    mod.define_method(name, method_body_with_arity)
   end
 
   private

src/main/java/org/truffleruby/core/proc/ProcNodes.java

@@ -34,6 +34,8 @@
 import org.truffleruby.language.control.RaiseException;
 import org.truffleruby.language.dispatch.DispatchNode;
 import org.truffleruby.language.locals.FindDeclarationVariableNodes.FindAndReadDeclarationVariableNode;
+import org.truffleruby.language.methods.Arity;
+import org.truffleruby.language.methods.SharedMethodInfo;
 import org.truffleruby.language.objects.AllocationTracing;
 import org.truffleruby.language.yield.CallBlockNode;
 import org.truffleruby.parser.ArgumentDescriptor;
@@ -256,19 +258,32 @@ public abstract static class ProcCreateSameArityNode extends PrimitiveArrayArgum
 
         @Specialization
         protected RubyProc createSameArityProc(RubyProc userProc, RubyProc block) {
-            final RubyProc composedProc = new RubyProc(
-                    coreLibrary().procClass,
-                    getLanguage().procShape,
-                    block.type,
+            final RubyProc composedProc = block.withSharedMethodInfo(
                     userProc.sharedMethodInfo,
-                    block.callTargets,
-                    block.callTarget,
-                    block.declarationFrame,
-                    block.declarationVariables,
-                    block.method,
-                    block.block,
-                    block.frameOnStackMarker,
-                    block.declarationContext);
+                    coreLibrary().procClass,
+                    getLanguage().procShape);
+            AllocationTracing.trace(composedProc, this);
+            return composedProc;
+        }
+    }
+
+    @Primitive(name = "proc_specify_arity", lowerFixnum = 1)
+    public abstract static class ProcSpecifyArityNode extends PrimitiveArrayArgumentsNode {
+        @Specialization
+        protected RubyProc specifyArityProc(RubyProc block, int argc) {
+            Arity oldArity = block.sharedMethodInfo.getArity();
+            final Arity newArity;
+            if (argc <= -1) {
+                newArity = new Arity(-(argc + 1), 0, true);
+            } else {
+                newArity = new Arity(argc, 0, false);
+            }
+
+            SharedMethodInfo newSharedMethodInfo = block.sharedMethodInfo.withArity(newArity);
+            final RubyProc composedProc = block.withSharedMethodInfo(
+                    newSharedMethodInfo,
+                    coreLibrary().procClass,
+                    getLanguage().procShape);
             AllocationTracing.trace(composedProc, this);
             return composedProc;
         }