[GR-19220] Fix incorrect arity for C-extension methods (#2275)

PullRequest: truffleruby/2459

Author: eregon

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;
         }

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

@@ -74,6 +74,22 @@ public RubyProc(
         this.declarationContext = declarationContext;
     }
 
+    public RubyProc withSharedMethodInfo(SharedMethodInfo newSharedMethodInfo, RubyClass newRubyClass, Shape newShape) {
+        return new RubyProc(
+                newRubyClass,
+                newShape,
+                type,
+                newSharedMethodInfo,
+                callTargets,
+                callTarget,
+                declarationFrame,
+                declarationVariables,
+                method,
+                block,
+                frameOnStackMarker,
+                declarationContext);
+    }
+
     @Override
     public void getAdjacentObjects(Set<Object> reachable) {
         ObjectGraph.getObjectsInFrame(declarationFrame, reachable);

src/main/java/org/truffleruby/language/methods/SharedMethodInfo.java

@@ -91,6 +91,18 @@ public SharedMethodInfo convertMethodMissingToMethod(RubyModule declaringModule,
                 newArgs);
     }
 
+    public SharedMethodInfo withArity(Arity newArity) {
+        return new SharedMethodInfo(
+                sourceSection,
+                lexicalScope,
+                newArity,
+                backtraceName,
+                blockDepth,
+                parseName,
+                notes,
+                argumentDescriptors);
+    }
+
     public SourceSection getSourceSection() {
         return sourceSection;
     }