fix set flags for *args and **kwargs when dealing with generators

Author: cosminbasca

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/code/PCode.java

@@ -214,10 +214,10 @@ private static int extractStackSize(RootNode rootNode) {
     }
 
     private void extractArgStats() {
-        String[] freevars = getFreeVars(rootNode);
-        String[] cellvars = getCellVars(rootNode);
-        Set<String> freeVarsSet = asSet(freevars);
-        Set<String> cellVarsSet = asSet(cellvars);
+        this.freevars = getFreeVars(rootNode);
+        this.cellvars = getCellVars(rootNode);
+        Set<String> freeVarsSet = asSet((String[])freevars);
+        Set<String> cellVarsSet = asSet((String[])cellvars);
 
         List<ReadKeywordNode> readKeywordNodes = NodeUtil.findAllNodeInstances(rootNode, ReadKeywordNode.class);
         List<ReadIndexedArgumentNode> readIndexedArgumentNodes = NodeUtil.findAllNodeInstances(rootNode, ReadIndexedArgumentNode.class);
@@ -229,52 +229,49 @@ private void extractArgStats() {
         allArgNames.addAll(kwNames);
         allArgNames.addAll(argNames);
 
-        int argcount = readIndexedArgumentNodes.size();
-        int kwonlyargcount = 0;
-        int flags = 0;
+        this.argcount = readIndexedArgumentNodes.size();
+        this.kwonlyargcount = 0;
+        this.flags = 0;
 
         for (ReadKeywordNode kwNode : readKeywordNodes) {
             if (!kwNode.canBePositional()) {
                 kwonlyargcount++;
             }
         }
 
-        Set<String> varnames = new HashSet<>();
+        Set<String> varnamesSet = new HashSet<>();
         for (Object identifier : rootNode.getFrameDescriptor().getIdentifiers()) {
             if (identifier instanceof String) {
                 String varName = (String) identifier;
 
                 if (core.getParser().isIdentifier(core, varName)) {
                     if (allArgNames.contains(varName)) {
-                        varnames.add(varName);
+                        varnamesSet.add(varName);
                     } else if (!freeVarsSet.contains(varName) && !cellVarsSet.contains(varName)) {
-                        varnames.add(varName);
+                        varnamesSet.add(varName);
                     }
                 }
             }
         }
 
-        // set the flags
+        // 0x20 - generator
+        RootNode funcRootNode = rootNode;
+        if (funcRootNode instanceof GeneratorFunctionRootNode) {
+            flags |= (1 << 5);
+            funcRootNode = ((GeneratorFunctionRootNode) funcRootNode).getFunctionRootNode();
+        }
+
         // 0x04 - *arguments
-        if (NodeUtil.findAllNodeInstances(rootNode, ReadVarArgsNode.class).size() == 1) {
+        if (NodeUtil.findAllNodeInstances(funcRootNode, ReadVarArgsNode.class).size() == 1) {
             flags |= (1 << 2);
         }
         // 0x08 - **keywords
-        if (NodeUtil.findAllNodeInstances(rootNode, ReadVarKeywordsNode.class).size() == 1) {
+        if (NodeUtil.findAllNodeInstances(funcRootNode, ReadVarKeywordsNode.class).size() == 1) {
             flags |= (1 << 3);
         }
-        // 0x20 - generator
-        if (rootNode instanceof GeneratorFunctionRootNode) {
-            flags |= (1 << 5);
-        }
 
-        this.argcount = argcount;
-        this.kwonlyargcount = kwonlyargcount;
-        this.freevars = freevars;
-        this.cellvars = cellvars;
-        this.varnames = varnames.toArray();
-        this.nlocals = varnames.size();
-        this.flags = flags;
+        this.varnames = varnamesSet.toArray();
+        this.nlocals = varnamesSet.size();
     }
 
     public RootNode getRootNode() {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/generator/GeneratorFunctionRootNode.java

@@ -48,6 +48,7 @@
 import com.oracle.truffle.api.RootCallTarget;
 import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.RootNode;
 
 public class GeneratorFunctionRootNode extends PClosureFunctionRootNode {
     private final RootCallTarget callTarget;
@@ -76,4 +77,8 @@ public GeneratorFunctionRootNode(PythonLanguage language, RootCallTarget callTar
     public Object execute(VirtualFrame frame) {
         return factory.createGenerator(getName(), callTarget, frameDescriptor, frame.getArguments(), closure, cellSlots, numOfActiveFlags, numOfGeneratorBlockNode, numOfGeneratorForNode);
     }
+
+    public RootNode getFunctionRootNode() {
+        return callTarget.getRootNode();
+    }
 }