Create multiple call targets for generators

msimacek · msimacek · commit 4e3c20e0fa4c · 2022-05-25T16:46:10.000+02:00
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/GeneratorBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/GeneratorBuiltins.java
@@ -142,7 +142,7 @@ abstract static class ResumeGeneratorNode extends Node {
         public abstract Object execute(VirtualFrame frame, PGenerator self, Object sendValue);
 
         @Specialization(guards = "sameCallTarget(self.getCurrentCallTarget(), call.getCallTarget())", limit = "getCallSiteInlineCacheMaxDepth()")
-        static Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
+        Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
                         @Cached("createDirectCall(self.getCurrentCallTarget())") CallTargetInvokeNode call,
                         @Cached PRaiseNode raiseNode) {
             self.setRunning(true);
@@ -153,20 +153,25 @@ static Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
             Object result;
             try {
                 result = call.execute(frame, null, null, null, arguments);
+                if (result == null) {
+                    self.markAsFinished();
+                }
             } catch (PException e) {
                 self.markAsFinished();
                 throw e;
             } finally {
                 self.setRunning(false);
-                self.setNextCallTarget();
+                if (!self.isFinished()) {
+                    self.setNextCallTarget(PythonLanguage.get(this));
+                }
             }
             handleResult(self, result, raiseNode);
             return result;
         }
 
         @Specialization(replaces = "cached")
         @Megamorphic
-        static Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
+        Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
                         @Cached ConditionProfile hasFrameProfile,
                         @Cached GenericInvokeNode call,
                         @Cached PRaiseNode raiseNode) {
@@ -182,12 +187,17 @@ static Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
                 } else {
                     result = call.execute(self.getCurrentCallTarget(), arguments);
                 }
+                if (result == null) {
+                    self.markAsFinished();
+                }
             } catch (PException e) {
                 self.markAsFinished();
                 throw e;
             } finally {
                 self.setRunning(false);
-                self.setNextCallTarget();
+                if (!self.isFinished()) {
+                    self.setNextCallTarget(PythonLanguage.get(this));
+                }
             }
             handleResult(self, result, raiseNode);
             return result;
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/PGenerator.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/PGenerator.java
@@ -34,6 +34,7 @@
 import com.oracle.graal.python.builtins.objects.function.PArguments;
 import com.oracle.graal.python.builtins.objects.iterator.PIntRangeIterator;
 import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
+import com.oracle.graal.python.nodes.bytecode.PBytecodeGeneratorRootNode;
 import com.oracle.graal.python.nodes.bytecode.PBytecodeRootNode;
 import com.oracle.graal.python.nodes.generator.AbstractYieldNode;
 import com.oracle.graal.python.nodes.generator.YieldFromNode;
@@ -42,6 +43,7 @@
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
 import com.oracle.truffle.api.Assumption;
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.RootCallTarget;
 import com.oracle.truffle.api.Truffle;
@@ -54,7 +56,6 @@ public final class PGenerator extends PythonBuiltinObject {
 
     private String name;
     private String qualname;
-    private final boolean usesBytecode;
     /**
      * Call targets with copies of the generator's AST. Each call target corresponds to one possible
      * entry point into the generator: the first call, and continuation for each yield. Each AST can
@@ -71,6 +72,8 @@ public final class PGenerator extends PythonBuiltinObject {
     private final Object iterator;
     private final boolean isPRangeIterator;
     private final GeneratorInfo generatorInfo;
+    private final PBytecodeRootNode bytecodeRootNode;
+    private final PBytecodeRootNode.FrameInfo frameInfo;
     // running means it is currently on the stack, not just started
     private boolean running;
 
@@ -106,12 +109,12 @@ public static PGenerator create(PythonLanguage lang, String name, String qualnam
         }
         assignCells(generatorFrame, cellVarSlots, cellVarAssumptions);
         PArguments.setGeneratorFrameLocals(generatorFrameArguments, factory.createDictLocals(generatorFrame));
-        return new PGenerator(lang, name, qualname, callTargets, generatorInfo, arguments, closure, iterator, false);
+        return new PGenerator(lang, name, qualname, callTargets, generatorInfo, arguments, closure, iterator);
     }
 
-    public static PGenerator create(PythonLanguage lang, String name, String qualname, PBytecodeRootNode rootNode, RootCallTarget bytecodeCallTarget, Object[] arguments) {
+    public static PGenerator create(PythonLanguage lang, String name, String qualname, PBytecodeRootNode rootNode, RootCallTarget[] callTargets, Object[] arguments) {
         rootNode.createGeneratorFrame(arguments);
-        return new PGenerator(lang, name, qualname, new RootCallTarget[]{bytecodeCallTarget}, null, arguments, null, null, true);
+        return new PGenerator(lang, name, qualname, rootNode, callTargets, arguments);
     }
 
     @ExplodeLoop
@@ -131,7 +134,7 @@ private static void assignClosure(PCell[] closure, MaterializedFrame generatorFr
     }
 
     private PGenerator(PythonLanguage lang, String name, String qualname, RootCallTarget[] callTargets, GeneratorInfo generatorInfo, Object[] arguments,
-                    PCell[] closure, Object iterator, boolean usesBytecode) {
+                    PCell[] closure, Object iterator) {
         super(PythonBuiltinClassType.PGenerator, PythonBuiltinClassType.PGenerator.getInstanceShape(lang));
         this.name = name;
         this.qualname = qualname;
@@ -143,11 +146,35 @@ private PGenerator(PythonLanguage lang, String name, String qualname, RootCallTa
         this.finished = false;
         this.iterator = iterator;
         this.isPRangeIterator = iterator instanceof PIntRangeIterator;
-        this.usesBytecode = usesBytecode;
+        this.bytecodeRootNode = null;
+        this.frameInfo = null;
     }
 
-    public void setNextCallTarget() {
-        if (!usesBytecode) {
+    private PGenerator(PythonLanguage lang, String name, String qualname, PBytecodeRootNode rootNode, RootCallTarget[] callTargets, Object[] arguments) {
+        super(PythonBuiltinClassType.PGenerator, PythonBuiltinClassType.PGenerator.getInstanceShape(lang));
+        this.name = name;
+        this.qualname = qualname;
+        this.callTargets = callTargets;
+        this.currentCallTarget = 0;
+        this.arguments = arguments;
+        this.finished = false;
+        this.bytecodeRootNode = rootNode;
+        this.frameInfo = (PBytecodeRootNode.FrameInfo) rootNode.getFrameDescriptor().getInfo();
+        this.iterator = null;
+        this.isPRangeIterator = false;
+        this.closure = null;
+        this.generatorInfo = null;
+    }
+
+    public void setNextCallTarget(PythonLanguage language) {
+        if (usesBytecode()) {
+            currentCallTarget = getBci();
+            if (callTargets[currentCallTarget] == null) {
+                CompilerDirectives.transferToInterpreterAndInvalidate();
+                PBytecodeGeneratorRootNode rootNode = new PBytecodeGeneratorRootNode(language, bytecodeRootNode, currentCallTarget, getStackTop());
+                callTargets[currentCallTarget] = rootNode.getCallTarget();
+            }
+        } else {
             currentCallTarget = PArguments.getControlDataFromGeneratorArguments(getArguments()).getLastYieldIndex();
         }
     }
@@ -162,6 +189,7 @@ public RootCallTarget getCurrentCallTarget() {
     }
 
     public AbstractYieldNode getCurrentYieldNode() {
+        assert !usesBytecode();
         if (currentCallTarget == 0 || running || finished) {
             // Not stopped on a yield
             return null;
@@ -171,38 +199,36 @@ public AbstractYieldNode getCurrentYieldNode() {
     }
 
     public boolean usesBytecode() {
-        return usesBytecode;
+        return bytecodeRootNode != null;
     }
 
     public Object getYieldFrom() {
-        if (!usesBytecode) {
+        if (!usesBytecode()) {
             AbstractYieldNode currentYield = getCurrentYieldNode();
             if (currentYield instanceof YieldFromNode) {
                 int iteratorSlot = ((YieldFromNode) currentYield).getIteratorSlot();
                 return PArguments.getControlDataFromGeneratorArguments(arguments).getIteratorAt(iteratorSlot);
             }
             return null;
         } else {
-            MaterializedFrame generatorFrame = PArguments.getGeneratorFrame(arguments);
-            PBytecodeRootNode.FrameInfo info = (PBytecodeRootNode.FrameInfo) generatorFrame.getFrameDescriptor().getInfo();
-            return info.getYieldFrom(generatorFrame);
+            return frameInfo.getYieldFrom(PArguments.getGeneratorFrame(arguments));
         }
     }
 
     public boolean isStarted() {
-        return (currentCallTarget != 0 || usesBytecode && getBci() > 0) && !running;
+        return currentCallTarget != 0 && !running;
     }
 
     public int getBci() {
-        MaterializedFrame generatorFrame = PArguments.getGeneratorFrame(arguments);
-        PBytecodeRootNode.FrameInfo info = (PBytecodeRootNode.FrameInfo) generatorFrame.getFrameDescriptor().getInfo();
-        return info.getBci(generatorFrame);
+        return frameInfo.getBci(PArguments.getGeneratorFrame(arguments));
+    }
+
+    public int getStackTop() {
+        return frameInfo.getGeneratorStackTop(PArguments.getGeneratorFrame(arguments));
     }
 
     public Object getReturnValue() {
-        MaterializedFrame generatorFrame = PArguments.getGeneratorFrame(arguments);
-        PBytecodeRootNode.FrameInfo info = (PBytecodeRootNode.FrameInfo) generatorFrame.getFrameDescriptor().getInfo();
-        return info.getGeneratorReturnValue(generatorFrame);
+        return frameInfo.getGeneratorReturnValue(PArguments.getGeneratorFrame(arguments));
     }
 
     public Object[] getArguments() {
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorFunctionRootNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorFunctionRootNode.java
@@ -47,6 +47,7 @@
 import com.oracle.graal.python.nodes.PRootNode;
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.RootCallTarget;
 import com.oracle.truffle.api.frame.FrameDescriptor;
@@ -55,17 +56,20 @@
 
 public class PBytecodeGeneratorFunctionRootNode extends PRootNode {
     private final PBytecodeRootNode rootNode;
-    private final RootCallTarget callTarget;
     private final String originalName;
 
+    @CompilationFinal(dimensions = 1) private final RootCallTarget[] callTargets;
+
     @Child private PythonObjectFactory factory = PythonObjectFactory.create();
 
     @TruffleBoundary
     public PBytecodeGeneratorFunctionRootNode(PythonLanguage language, FrameDescriptor frameDescriptor, PBytecodeRootNode rootNode, String originalName) {
         super(language, frameDescriptor);
         this.rootNode = rootNode;
         this.originalName = originalName;
-        this.callTarget = rootNode.getCallTarget();
+        // TODO compress somehow
+        this.callTargets = new RootCallTarget[rootNode.getCodeUnit().code.length];
+        this.callTargets[0] = rootNode.getCallTarget();
     }
 
     @Override
@@ -75,11 +79,7 @@ public Object execute(VirtualFrame frame) {
         // This is passed from InvokeNode node
         PFunction generatorFunction = PArguments.getGeneratorFunction(arguments);
         assert generatorFunction != null;
-        return factory.createGenerator(generatorFunction.getName(), generatorFunction.getQualname(), rootNode, callTarget, arguments);
-    }
-
-    public PBytecodeRootNode getFunctionRootNode() {
-        return rootNode;
+        return factory.createGenerator(generatorFunction.getName(), generatorFunction.getQualname(), rootNode, callTargets, arguments);
     }
 
     @Override
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorRootNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorRootNode.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * The Universal Permissive License (UPL), Version 1.0
+ *
+ * Subject to the condition set forth below, permission is hereby granted to any
+ * person obtaining a copy of this software, associated documentation and/or
+ * data (collectively the "Software"), free of charge and under any and all
+ * copyright rights in the Software, and any and all patent rights owned or
+ * freely licensable by each licensor hereunder covering either (i) the
+ * unmodified Software as contributed to or provided by such licensor, or (ii)
+ * the Larger Works (as defined below), to deal in both
+ *
+ * (a) the Software, and
+ *
+ * (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
+ * one is included with the Software each a "Larger Work" to which the Software
+ * is contributed by such licensors),
+ *
+ * without restriction, including without limitation the rights to copy, create
+ * derivative works of, display, perform, and distribute the Software and make,
+ * use, sell, offer for sale, import, export, have made, and have sold the
+ * Software and the Larger Work(s), and to sublicense the foregoing rights on
+ * either these or other terms.
+ *
+ * This license is subject to the following condition:
+ *
+ * The above copyright notice and either this complete permission notice or at a
+ * minimum a reference to the UPL must be included in all copies or substantial
+ * portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+package com.oracle.graal.python.nodes.bytecode;
+
+import com.oracle.graal.python.PythonLanguage;
+import com.oracle.graal.python.builtins.objects.function.Signature;
+import com.oracle.graal.python.nodes.PRootNode;
+import com.oracle.graal.python.runtime.ExecutionContext;
+import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.source.SourceSection;
+
+public class PBytecodeGeneratorRootNode extends PRootNode {
+    private final PBytecodeRootNode rootNode;
+    private final int resumeBci;
+    private final int resumeStackTop;
+
+    @Child private ExecutionContext.CalleeContext calleeContext = ExecutionContext.CalleeContext.create();
+
+    @TruffleBoundary
+    public PBytecodeGeneratorRootNode(PythonLanguage language, PBytecodeRootNode rootNode, int resumeBci, int resumeStackTop) {
+        super(language, rootNode.getFrameDescriptor());
+        this.rootNode = rootNode;
+        this.resumeBci = resumeBci;
+        this.resumeStackTop = resumeStackTop;
+    }
+
+    @Override
+    public Object execute(VirtualFrame frame) {
+        calleeContext.enter(frame);
+        try {
+            return rootNode.executeFromBci(frame, resumeBci, resumeStackTop);
+        } finally {
+            calleeContext.exit(frame, this);
+        }
+    }
+
+    @Override
+    public String getName() {
+        return rootNode.getName();
+    }
+
+    @Override
+    public String toString() {
+        CompilerAsserts.neverPartOfCompilation();
+        return "<bytecode " + rootNode.getName() + ">";
+    }
+
+    @Override
+    public Signature getSignature() {
+        return rootNode.getSignature();
+    }
+
+    @Override
+    public boolean isPythonInternal() {
+        return rootNode.isPythonInternal();
+    }
+
+    @Override
+    public SourceSection getSourceSection() {
+        return rootNode.getSourceSection();
+    }
+}
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/object/PythonObjectFactory.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/object/PythonObjectFactory.java
