[GR-79186] Avoid temporary locals if loop or comprehension target is fast local.

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2026, 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
@@ -97,6 +97,10 @@ public static void checkForbiddenName(ParserCallbacks parserCallbacks, SourceRan
         }
     }
 
+    public static boolean mayBeForbiddenName(String id) {
+        return id.equals("__debug__");
+    }
+
     public static void checkSubscripter(ParserCallbacks parserCallbacks, ExprTy e) {
         if (e instanceof ExprTy.Constant) {
             switch (((ExprTy.Constant) e).value.kind) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/bytecode_dsl/RootNodeCompiler.java

@@ -48,6 +48,7 @@
 import static com.oracle.graal.python.compiler.SSTUtils.checkForbiddenArgs;
 import static com.oracle.graal.python.compiler.SSTUtils.checkIndex;
 import static com.oracle.graal.python.compiler.SSTUtils.checkSubscripter;
+import static com.oracle.graal.python.compiler.SSTUtils.mayBeForbiddenName;
 import static com.oracle.graal.python.compiler.bytecode_dsl.BytecodeDSLCompilerUtils.COMPREHENSION_ARGS;
 import static com.oracle.graal.python.compiler.bytecode_dsl.BytecodeDSLCompilerUtils.NO_ARGS;
 import static com.oracle.graal.python.compiler.bytecode_dsl.BytecodeDSLCompilerUtils.TYPE_PARAMS_DEFAULTS;
@@ -762,8 +763,9 @@ public void reset() {
         prevSaveExceptionLocal = null;
         lastTracedLine = -1;
         this.inExceptStar = false;
+        temporaryLocals.clear();
         if (ASSERTIONS_ENABLED) {
-            temporaryLocals.clear();
+            temporaryLocalsTraces.clear();
         }
     }
 
@@ -795,14 +797,14 @@ void endRootNode(Builder b) {
         b.endRoot();
         endRootSourceSection(b);
         b.endSource();
-        if (ASSERTIONS_ENABLED && !temporaryLocals.isEmpty()) {
+        if (ASSERTIONS_ENABLED && !temporaryLocalsTraces.isEmpty()) {
             throw new AssertionError(this.qualName + "\n\n" + formatTempLocalsStackTraces());
         }
     }
 
     private String formatTempLocalsStackTraces() {
         StringBuilder sb = new StringBuilder();
-        for (Object v : temporaryLocals.values()) {
+        for (Object v : temporaryLocalsTraces.values()) {
             if (v instanceof RuntimeException re) {
                 sb.append("\n==================\n");
                 StringWriter sw = new StringWriter();
@@ -825,44 +827,59 @@ private static boolean assertionsEnabled() {
         return enabled;
     }
 
-    private HashMap<BytecodeLocal, Object> temporaryLocals = ASSERTIONS_ENABLED ? new HashMap<>() : null;
+    private HashMap<BytecodeLocal, Object> temporaryLocalsTraces = ASSERTIONS_ENABLED ? new HashMap<>() : null;
+    private HashSet<BytecodeLocal> temporaryLocals = new HashSet<>();
 
     private BytecodeLocal beginTemporaryLocal(Builder b) {
         BytecodeLocal local = b.createLocal();
+        temporaryLocals.add(local);
+        assert isTemporaryLocal(local);
         if (ASSERTIONS_ENABLED) {
-            Object previous = temporaryLocals.put(local, TRACK_TEMP_LOCALS ? new RuntimeException() : local);
+            Object previous = temporaryLocalsTraces.put(local, TRACK_TEMP_LOCALS ? new RuntimeException() : local);
             if (previous != null) {
                 throw new AssertionError();
             }
         }
         return local;
     }
 
+    public BytecodeLocal beginTemporaryLocalOrGetLocal(ExprTy target, Builder b) {
+        if (target instanceof ExprTy.Name nameExpr) {
+            BytecodeLocal l = getFastLocal(nameExpr.id);
+            if (l != null) {
+                return l;
+            }
+        }
+        return beginTemporaryLocal(b);
+    }
+
+    public boolean isTemporaryLocal(BytecodeLocal local) {
+        return temporaryLocals.contains(local);
+    }
+
     private void endTemporaryLocal(BytecodeLocal local, Builder b) {
-        markTemporaryLocalCleared(local);
-        b.emitClearLocal(local);
+        if (isTemporaryLocal(local)) {
+            markTemporaryLocalCleared(local);
+            b.emitClearLocal(local);
+        }
     }
 
     private void loadAndEndTemporaryLocal(BytecodeLocal local, Builder b) {
-        markTemporaryLocalCleared(local);
-        b.emitLoadAndClearTempLocal(local);
+        if (isTemporaryLocal(local)) {
+            markTemporaryLocalCleared(local);
+            b.emitLoadAndClearTempLocal(local);
+        }
     }
 
     private void markTemporaryLocalCleared(BytecodeLocal local) {
-        if (ASSERTIONS_ENABLED) {
-            if (temporaryLocals.remove(local) == null) {
-                throw new AssertionError();
-            }
+        if (!temporaryLocals.remove(local)) {
+            throw new AssertionError();
         }
-    }
-
-    private BytecodeLocal checkTemporaryLocal(BytecodeLocal local) {
         if (ASSERTIONS_ENABLED) {
-            if (local != null && !temporaryLocals.containsKey(local)) {
-                throw new AssertionError("Temporary local was already cleared");
+            if (temporaryLocalsTraces.remove(local) == null) {
+                throw new AssertionError();
             }
         }
-        return local;
     }
 
     void emitTraceLineChecked(SSTNode node, Builder b) {
@@ -1382,7 +1399,7 @@ private void emitComprehension(ComprehensionTy[] generators, int index, Builder
             statementCompiler.emitAsyncFor(iter, comp.target, null, true, index,
                             (stmtComp, idx) -> emitComprehensionBody(generators, idx, type, collection, accumulateProducer, stmtComp));
         } else {
-            BytecodeLocal localValue = beginTemporaryLocal(b);
+            BytecodeLocal localValue = beginTemporaryLocalOrGetLocal(comp.target, b);
 
             b.beginBlock();
             b.beginBindStackValue();
@@ -1411,7 +1428,7 @@ private void emitComprehension(ComprehensionTy[] generators, int index, Builder
 
             b.beginBlock();
 
-            comp.target.accept(statementCompiler.new StoreVisitor(() -> b.emitLoadLocal(localValue)));
+            statementCompiler.storeTemporaryLocalToTarget(localValue, comp.target, b);
             emitComprehensionBody(generators, index, type, collection, accumulateProducer, statementCompiler);
 
             b.endBlock();
@@ -1851,6 +1868,25 @@ private void emitNameOperation(String name, NameOperation op, Builder b) {
         emitNameSlowOperation(mangled, op, b);
     }
 
+    private BytecodeLocal getFastLocal(String name) {
+        if (mayBeForbiddenName(name)) {
+            return null;
+        }
+        String mangled = maybeMangle(name);
+        EnumSet<DefUse> uses = scope.getUseOfName(mangled);
+        if (uses != null) {
+            if (uses.contains(DefUse.Free) || uses.contains(DefUse.Cell)) {
+                return null;
+            } else if (uses.contains(DefUse.Local)) {
+                if (scope.isFunction()) {
+                    assert varnames.containsKey(mangled) : String.format("scope analysis did not mark %s as a regular variable", mangled);
+                    return locals.get(mangled);
+                }
+            }
+        }
+        return null;
+    }
+
     private void emitReadLocal(String name, Builder b) {
         emitNameOperation(name, NameOperation.Read, b);
     }
@@ -2016,6 +2052,23 @@ private BytecodeLocal beginTemporaryLocal() {
             return RootNodeCompiler.this.beginTemporaryLocal(b);
         }
 
+        /**
+         * If the target expression is simple local variable, we can often just directly store into it.
+         * Otherwise, we create temporary local. {@link BytecodeLocal} instances returned by this should
+         * be passed to {@link #storeTemporaryLocalToTarget(BytecodeLocal, ExprTy, Builder)}.
+         */
+        public BytecodeLocal beginTemporaryLocalOrGetLocal(ExprTy target, Builder b) {
+            return RootNodeCompiler.this.beginTemporaryLocalOrGetLocal(target, b);
+        }
+
+        private void storeTemporaryLocalToTarget(BytecodeLocal temporaryLocal, ExprTy target, Builder b) {
+            if (RootNodeCompiler.this.isTemporaryLocal(temporaryLocal)) {
+                target.accept(new StoreVisitor(() -> {
+                    b.emitLoadLocal(temporaryLocal);
+                }));
+            }
+        }
+
         private void endTemporaryLocal(BytecodeLocal local) {
             RootNodeCompiler.this.endTemporaryLocal(local, b);
         }
@@ -4328,7 +4381,7 @@ public Void visit(StmtTy.For node) {
             inExceptStar = false;
             b.beginBlock();
 
-            BytecodeLocal value = beginTemporaryLocal();
+            BytecodeLocal value = beginTemporaryLocalOrGetLocal(node.target, b);
 
             b.beginBindStackValue();
             b.beginGetIter();
@@ -4355,9 +4408,7 @@ public Void visit(StmtTy.For node) {
             // body
             b.beginBlock();
             continueLabel = b.createLabel();
-            node.target.accept(new StoreVisitor(() -> {
-                b.emitLoadLocal(value);
-            }));
+            storeTemporaryLocalToTarget(value, node.target, b);
 
             visitSequence(node.body);
             b.emitLabel(continueLabel);

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode_dsl/PBytecodeDSLRootNode.java

@@ -1874,7 +1874,6 @@ public static boolean doObjectIterator(VirtualFrame frame, LocalAccessor output,
                         @Bind BytecodeNode bytecode) {
             if (!iterator.hasNext()) {
                 iterator.setExhausted();
-                output.setObject(bytecode, frame, null);
                 return false;
             }
             Object value = iterator.next();
@@ -1917,7 +1916,6 @@ public static boolean doIterator(VirtualFrame frame, LocalAccessor output, Objec
                 output.setObject(bytecode, frame, value);
                 return true;
             } catch (IteratorExhausted e) {
-                output.setObject(bytecode, frame, null);
                 return false;
             }
         }
@@ -3920,38 +3918,6 @@ public static Object perform(
         }
     }
 
-    @Operation(storeBytecodeIndex = false)
-    @ConstantOperand(type = LocalAccessor.class)
-    public static final class TeeLocal {
-        @Specialization
-        public static int doInt(VirtualFrame frame, LocalAccessor local, int value,
-                        @Bind BytecodeNode bytecode) {
-            local.setInt(bytecode, frame, value);
-            return value;
-        }
-
-        @Specialization
-        public static double doDouble(VirtualFrame frame, LocalAccessor local, double value,
-                        @Bind BytecodeNode bytecode) {
-            local.setDouble(bytecode, frame, value);
-            return value;
-        }
-
-        @Specialization
-        public static long doLong(VirtualFrame frame, LocalAccessor local, long value,
-                        @Bind BytecodeNode bytecode) {
-            local.setLong(bytecode, frame, value);
-            return value;
-        }
-
-        @Specialization(replaces = {"doInt", "doDouble", "doLong"})
-        public static Object doObject(VirtualFrame frame, LocalAccessor local, Object value,
-                        @Bind BytecodeNode bytecode) {
-            local.setObject(bytecode, frame, value);
-            return value;
-        }
-    }
-
     @Operation(storeBytecodeIndex = true)
     public static final class GetLen {
         @Specialization