Fix argcells in LocalsStorage

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/CodeUnit.java

@@ -82,6 +82,7 @@ public final class CodeUnit {
     public final TruffleString[] cellvars;
     public final TruffleString[] freevars;
     public final int[] cell2arg;
+    public final int[] arg2cell;
 
     public final Object[] constants;
     public final long[] primitiveConstants;
@@ -109,6 +110,7 @@ public CodeUnit(TruffleString name, TruffleString qualname,
                     int[] exceptionHandlerRanges, int conditionProfileCount,
                     int startOffset, int startLine,
                     byte[] outputCanQuicken, byte[] variableShouldUnbox, int[][] generalizeInputsMap, int[][] generalizeVarsMap) {
+        int[] arg2cell;
         this.name = name;
         this.qualname = qualname != null ? qualname : name;
         this.argCount = argCount;
@@ -123,6 +125,17 @@ public CodeUnit(TruffleString name, TruffleString qualname,
         this.cellvars = cellvars;
         this.freevars = freevars;
         this.cell2arg = cell2arg;
+        arg2cell = null;
+        if (cell2arg != null) {
+            arg2cell = new int[getTotalArgCount()];
+            Arrays.fill(arg2cell, -1);
+            for (int i = 0; i < cell2arg.length; i++) {
+                if (cell2arg[i] >= 0) {
+                    arg2cell[cell2arg[i]] = i;
+                }
+            }
+        }
+        this.arg2cell = arg2cell;
         this.constants = constants;
         this.primitiveConstants = primitiveConstants;
         this.exceptionHandlerRanges = exceptionHandlerRanges;
@@ -218,10 +231,21 @@ public boolean isGeneratorOrCoroutine() {
         return (flags & (PCode.CO_GENERATOR | PCode.CO_COROUTINE | PCode.CO_ASYNC_GENERATOR | PCode.CO_ITERABLE_COROUTINE)) != 0;
     }
 
-    public int getTotalArgCount() {
+    public int getRegularArgCount() {
         return argCount + positionalOnlyArgCount + kwOnlyArgCount;
     }
 
+    public int getTotalArgCount() {
+        int count = getRegularArgCount();
+        if (takesVarArgs()) {
+            count++;
+        }
+        if (takesVarKeywordArgs()) {
+            count++;
+        }
+        return count;
+    }
+
     @SuppressWarnings("fallthrough")
     @Override
     public String toString() {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java

@@ -496,7 +496,15 @@ private static FrameDescriptor makeFrameDescriptor(CodeUnit co) {
         newBuilder.info(new FrameInfo());
         // locals
         for (int i = 0; i < co.varnames.length; i++) {
-            newBuilder.addSlot(FrameSlotKind.Illegal, co.varnames[i], null);
+            TruffleString varname = co.varnames[i];
+            if (co.arg2cell != null && i < co.arg2cell.length && co.arg2cell[i] >= 0) {
+                /*
+                 * If an argument is a cell, its slot gets superseded by the cell's slot below. We
+                 * need to hide it from LocalsStorage and other introspection.
+                 */
+                varname = null;
+            }
+            newBuilder.addSlot(FrameSlotKind.Illegal, varname, null);
         }
         // cells
         for (int i = 0; i < co.cellvars.length; i++) {
@@ -776,7 +784,7 @@ private void copyArgs(Object[] args, Frame localFrame) {
             copyArgsFirstTime(args, localFrame);
             return;
         }
-        int argCount = co.getTotalArgCount();
+        int argCount = co.getRegularArgCount();
         for (int i = 0; i < argCount; i++) {
             Object arg = args[i + PArguments.USER_ARGUMENTS_OFFSET];
             if (variableTypes[i] == QuickeningTypes.OBJECT) {
@@ -799,7 +807,7 @@ private void copyArgs(Object[] args, Frame localFrame) {
     private void copyArgsFirstTime(Object[] args, Frame localFrame) {
         CompilerAsserts.neverPartOfCompilation();
         variableTypes = new byte[varnames.length];
-        int argCount = co.getTotalArgCount();
+        int argCount = co.getRegularArgCount();
         for (int i = 0; i < argCount; i++) {
             Object arg = args[i + PArguments.USER_ARGUMENTS_OFFSET];
             if ((variableShouldUnbox[i] & QuickeningTypes.INT) != 0 && arg instanceof Integer) {
@@ -831,7 +839,7 @@ public void createGeneratorFrame(Object[] arguments) {
 
     private void copyArgsAndCells(Frame localFrame, Object[] arguments) {
         copyArgs(arguments, localFrame);
-        int varIdx = co.getTotalArgCount();
+        int varIdx = co.getRegularArgCount();
         if (co.takesVarArgs()) {
             localFrame.setObject(varIdx++, factory.createTuple(PArguments.getVariableArguments(arguments)));
         }