Remove loop from special variable lookup

Author: horakivo

src/main/java/org/truffleruby/language/SpecialVariablesSendingNode.java

@@ -16,8 +16,8 @@
 import org.truffleruby.core.kernel.TruffleKernelNodes.GetSpecialVariableStorage;
 import org.truffleruby.language.arguments.ReadCallerVariablesNode;
 
-import org.truffleruby.language.locals.FindDeclarationVariableNodes;
 import org.truffleruby.language.threadlocal.SpecialVariableStorage;
+import org.truffleruby.parser.BlockFrameDescriptorInfo;
 
 /** Some Ruby methods need access to the caller special variables: see usages of {@link ReadCallerVariablesNode}. This
  * is used for methods which need to access the last regexp MatchData or the last IO line.
@@ -45,10 +45,12 @@ protected Assumption getSpecialVariableAssumption(Frame frame) {
             return Assumption.ALWAYS_VALID;
         }
 
-        var outerFrameDescriptor = FindDeclarationVariableNodes.getOuterFrameDescriptor(frame.getFrameDescriptor());
+        var descriptor = frame.getFrameDescriptor();
 
-        if (SpecialVariableStorage.hasSpecialVariableAssumption(outerFrameDescriptor)) {
-            return SpecialVariableStorage.getAssumption(outerFrameDescriptor);
+        if (SpecialVariableStorage.hasSpecialVariableAssumption(descriptor)) {
+            return SpecialVariableStorage.getAssumption(descriptor);
+        } else if (descriptor.getInfo() instanceof BlockFrameDescriptorInfo blockFrameDescriptorInfo) {
+            return blockFrameDescriptorInfo.getSpecialVariableAssumption();
         } else {
             return Assumption.ALWAYS_VALID;
         }

src/main/java/org/truffleruby/parser/BlockFrameDescriptorInfo.java

@@ -9,9 +9,13 @@
  */
 package org.truffleruby.parser;
 
+import com.oracle.truffle.api.Assumption;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
+import org.truffleruby.language.threadlocal.SpecialVariableStorage;
+
+import static org.truffleruby.language.threadlocal.SpecialVariableStorage.ASSUMPTION_NAME;
 
 public final class BlockFrameDescriptorInfo {
 
@@ -26,13 +30,25 @@ public static FrameDescriptor getDeclarationFrameDescriptor(FrameDescriptor topD
     }
 
     @CompilationFinal private FrameDescriptor descriptor;
+    private final Assumption specialVariableAssumption;
 
-    public BlockFrameDescriptorInfo() {
+    public BlockFrameDescriptorInfo(Assumption specialVariableAssumption) {
+        assert specialVariableAssumption != null && specialVariableAssumption.getName() == ASSUMPTION_NAME;
+        this.specialVariableAssumption = specialVariableAssumption;
     }
 
     public BlockFrameDescriptorInfo(FrameDescriptor descriptor) {
-        assert descriptor != null;
+        assert descriptor != null && descriptor.getInfo() != null;
         this.descriptor = descriptor;
+        this.specialVariableAssumption = getSpecialVariableAssumptionFromDescriptor(descriptor);
+    }
+
+    private Assumption getSpecialVariableAssumptionFromDescriptor(FrameDescriptor descriptor) {
+        if (descriptor.getInfo() instanceof BlockFrameDescriptorInfo blockFrameDescriptorInfo) {
+            return blockFrameDescriptorInfo.getSpecialVariableAssumption();
+        } else {
+            return SpecialVariableStorage.getAssumption(descriptor);
+        }
     }
 
     public FrameDescriptor getDescriptor() {
@@ -44,4 +60,9 @@ void set(FrameDescriptor descriptor) {
         assert this.descriptor == null;
         this.descriptor = descriptor;
     }
+
+    public Assumption getSpecialVariableAssumption() {
+        assert specialVariableAssumption != null;
+        return specialVariableAssumption;
+    }
 }

src/main/java/org/truffleruby/parser/TranslatorEnvironment.java

@@ -85,10 +85,13 @@ public TranslatorEnvironment(
                 BlockFrameDescriptorInfo parentBlockDescriptor = Objects
                         .requireNonNull(parent.blockFrameDescriptorInfo);
                 this.frameDescriptorBuilder = newFrameDescriptorBuilderForBlock(parentBlockDescriptor);
+                this.blockFrameDescriptorInfo = new BlockFrameDescriptorInfo(
+                        parentBlockDescriptor.getSpecialVariableAssumption());
             } else {
-                this.frameDescriptorBuilder = newFrameDescriptorBuilderForMethod();
+                var specialVariableAssumption = createSpecialVariableAssumption();
+                this.frameDescriptorBuilder = newFrameDescriptorBuilderForMethod(specialVariableAssumption);
+                this.blockFrameDescriptorInfo = new BlockFrameDescriptorInfo(specialVariableAssumption);
             }
-            this.blockFrameDescriptorInfo = new BlockFrameDescriptorInfo();
         } else {
             this.frameDescriptor = descriptor;
             this.blockFrameDescriptorInfo = new BlockFrameDescriptorInfo(descriptor);
@@ -163,15 +166,16 @@ public static FrameDescriptor.Builder newFrameDescriptorBuilderForBlock(BlockFra
         return builder;
     }
 
-    public static FrameDescriptor.Builder newFrameDescriptorBuilderForMethod() {
+    private static Assumption createSpecialVariableAssumption() {
+        return Assumption.create(SpecialVariableStorage.ASSUMPTION_NAME);
+    }
 
+    private static FrameDescriptor.Builder newFrameDescriptorBuilderForMethod(Assumption specialVariableAssumption) {
         var builder = FrameDescriptor.newBuilder().defaultValue(Nil.INSTANCE);
         // We need to access this Assumption from the FrameDescriptor,
         // and there is no way to get a RootNode from a FrameDescriptor, so we store it in the descriptor info.
         // We do not store it as slot info for footprint, to avoid needing an info array per FrameDescriptor.
-        final Assumption doesNotNeedSpecialVariableStorageAssumption = Assumption
-                .create(SpecialVariableStorage.ASSUMPTION_NAME);
-        builder.info(doesNotNeedSpecialVariableStorageAssumption);
+        builder.info(specialVariableAssumption);
 
 
         int selfIndex = builder.addSlot(FrameSlotKind.Illegal, SelfNode.SELF_IDENTIFIER, null);
@@ -187,6 +191,12 @@ public static FrameDescriptor.Builder newFrameDescriptorBuilderForMethod() {
         return builder;
     }
 
+
+    public static FrameDescriptor.Builder newFrameDescriptorBuilderForMethod() {
+        var specialVariableAssumption = createSpecialVariableAssumption();
+        return newFrameDescriptorBuilderForMethod(specialVariableAssumption);
+    }
+
     public int declareVar(Object name) {
         assert name != null && !(name instanceof String && ((String) name).isEmpty());