jruby · enebo · Feb 11, 2025 · Jan 28, 2025 · Jan 28, 2025
diff --git a/core/src/main/java/org/jruby/ir/persistence/IRReader.java b/core/src/main/java/org/jruby/ir/persistence/IRReader.java
@@ -15,6 +15,7 @@
 import org.jruby.runtime.Signature;
 
 import java.io.IOException;
+import java.util.BitSet;
 import java.util.EnumSet;
 
 import org.jruby.util.ByteList;
@@ -123,10 +124,15 @@ private static StaticScope decodeStaticScope(IRReaderDecoder decoder, StaticScop
         StaticScope.Type type = decoder.decodeStaticScopeType();
         if (RubyInstanceConfig.IR_READING_DEBUG) System.out.println("decodeStaticScope: type = " + type);
         String[] ids = decoder.decodeStringArray();
-        int firstKeywordIndex = decoder.decodeInt();
-        if (RubyInstanceConfig.IR_READING_DEBUG) System.out.println("decodeStaticScope: keyword index = " + firstKeywordIndex);
+        byte[] keywordIndices = decoder.decodeByteArray();
+        if (RubyInstanceConfig.IR_READING_DEBUG) System.out.println("decodeStaticScope: keyword indices count = " + keywordIndices.length);
 
-        StaticScope scope = StaticScopeFactory.newStaticScope(parentScope, type, decoder.getFilename(), ids, firstKeywordIndex);
+        StaticScope scope = StaticScopeFactory.newStaticScope(parentScope, type, decoder.getFilename(), ids, -1);
+
+        // We encode as empty list to not deal with null check here
+        if (keywordIndices.length > 0) {
+            scope.setKeywordIndices(BitSet.valueOf(keywordIndices));
+        }
 
         Signature signature = decoder.decodeSignature();
         scope.setSignature(signature);

diff --git a/core/src/main/java/org/jruby/ir/persistence/IRWriter.java b/core/src/main/java/org/jruby/ir/persistence/IRWriter.java
@@ -9,6 +9,7 @@
 import org.jruby.parser.StaticScope;
 
 import java.io.IOException;
+import java.util.BitSet;
 
 /**
  * Write IR data out to persistent store.  IRReader is capable of re-reading this
@@ -95,7 +96,7 @@ private static void persistStaticScope(IRWriterEncoder file, StaticScope staticS
         // This naively looks like a bug because these represent id's and not properly encoded names BUT all of those
         // symbols for these ids will be created when IRScope loads the LocalVariable versions of these...so this is ok.
         file.encode(staticScope.getVariables());
-        file.encode(staticScope.getFirstKeywordIndex());
+        file.encode(staticScope.getKeywordIndices().toByteArray());
         file.encode(staticScope.getSignature());
     }
 

diff --git a/core/src/main/java/org/jruby/parser/StaticScope.java b/core/src/main/java/org/jruby/parser/StaticScope.java
@@ -34,6 +34,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 
+import java.util.BitSet;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -48,15 +49,13 @@
 import org.jruby.ast.AssignableNode;
 import org.jruby.ast.DAsgnNode;
 import org.jruby.ast.DVarNode;
-import org.jruby.ast.IScopedNode;
 import org.jruby.ast.LocalAsgnNode;
 import org.jruby.ast.LocalVarNode;
 import org.jruby.ast.Node;
 import org.jruby.ast.VCallNode;
 import org.jruby.ir.IRMethod;
 import org.jruby.ir.IRScope;
 import org.jruby.ir.IRScopeType;
-import org.jruby.runtime.Block;
 import org.jruby.runtime.DynamicScope;
 import org.jruby.runtime.Helpers;
 import org.jruby.runtime.Signature;
@@ -112,8 +111,6 @@ public class StaticScope implements Serializable, Cloneable {
     private boolean isBlockOrEval;
     private boolean isArgumentScope; // Is this block and argument scope of a define_method.
 
-    private int firstKeywordIndex;
-
     // Method/Closure that this static scope corresponds to.  This is used to tell whether this
     // scope refers to a method scope or to determined IRScope of the parent of a compiling eval.
     private IRScope irScope;
@@ -124,6 +121,18 @@ public class StaticScope implements Serializable, Cloneable {
 
     private volatile Collection<String> ivarNames;
 
+    private BitSet keywordIndices = null;
+
+    // This is set for Prism since it knows all names for a scope right after parse but does not know which
+    // ones are keywords until we are actually processing those keywords in builder.
+    public void setKeywordIndices(BitSet keywordIndices) {
+        this.keywordIndices = keywordIndices;
+    }
+
+    public BitSet getKeywordIndices() {
+        return keywordIndices == null ? new BitSet() : keywordIndices;
+    }
+
     public enum Type {
         LOCAL, BLOCK, EVAL;
 
@@ -172,7 +181,6 @@ protected StaticScope(Type type, StaticScope enclosingScope, String file, String
         if (enclosingScope != null) this.scopeType = enclosingScope.getScopeType();
         this.isBlockOrEval = (type != Type.LOCAL);
         this.isArgumentScope = !isBlockOrEval;
-        this.firstKeywordIndex = firstKeywordIndex;
         this.file = file;
     }
 
@@ -181,7 +189,7 @@ protected StaticScope(Type type, StaticScope enclosingScope, String[] names) {
     }
 
     public int getFirstKeywordIndex() {
-        return firstKeywordIndex;
+        return -1;
     }
 
     public DynamicScope construct(DynamicScope parent) {
@@ -398,20 +406,19 @@ public AssignableNode assign(int line, RubySymbol name, Node value) {
      * a kwarg or an ordinary variable during live execution (See keywordExists).
      */
     public AssignableNode assignKeyword(int line, RubySymbol symbolID, Node value) {
-        AssignableNode assignment = assign(line, symbolID, value, this, 0);
+        AssignableNode assignable = assign(line, symbolID, value, this, 0);
 
-        // register first keyword index encountered
-        if (firstKeywordIndex == -1) firstKeywordIndex = ((IScopedNode) assignment).getIndex();
+        if (keywordIndices == null) keywordIndices = new BitSet();
+        keywordIndices.set(isDefined(symbolID.idString()));
 
-        return assignment;
+        return assignable;
     }
 
     public boolean keywordExists(String name) {
         if (name.equals("_")) return true;
         int slot = exists(name);
 
-        return slot >= 0 && firstKeywordIndex != -1 &&
-                slot >= firstKeywordIndex  && slot < firstKeywordIndex + signature.kwargs();
+        return slot >= 0 && keywordIndices != null && keywordIndices.get(slot);
     }
 
     /**
@@ -680,7 +687,7 @@ public StaticScope duplicate() {
         dupe.setModule(cref);
         dupe.setFile(file);
         dupe.setSignature(signature);
-        dupe.firstKeywordIndex = firstKeywordIndex;
+        dupe.setKeywordIndices(keywordIndices);
 
         return dupe;
     }

diff --git a/core/src/main/java/org/jruby/runtime/Helpers.java b/core/src/main/java/org/jruby/runtime/Helpers.java
@@ -23,6 +23,7 @@
 import java.nio.file.NotDirectoryException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.BitSet;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -2137,7 +2138,7 @@ public static String describeScope(StaticScope scope) {
                 Integer.toString(scope.getType().ordinal()) + ';'
                 + scope.getFile() + ';'
                 + Arrays.stream(scope.getVariables()).collect(Collectors.joining(",")) + ';'
-                + scope.getFirstKeywordIndex() + ';' +
+                + Arrays.toString(scope.getKeywordIndices().toByteArray()) + ';' +
                 + (signature == null ? Signature.NO_ARGUMENTS.encode() : signature.encode()) + ';'
                 + scope.getIRScope().getScopeType().ordinal() + ';'
                 + (instanceVariableNames.size() > 0
@@ -2154,13 +2155,20 @@ public static StaticScope restoreScope(String descriptor, StaticScope enclosingS
         String file = bits[1];
 
         String[] varNames = bits[2].split(",");
-        int kwIndex = Integer.parseInt(bits[3]);
         Signature signature = Signature.decode(Long.parseLong(bits[4]));
         IRScopeType scopeType = IRScopeType.fromOrdinal(Integer.parseInt(bits[5]));
         String encodedIvars = bits[6];
         Collection<String> ivarNames = encodedIvars.equals("NONE") ? Collections.EMPTY_LIST : Arrays.asList(encodedIvars.split(","));
 
-        StaticScope scope = StaticScopeFactory.newStaticScope(enclosingScope, type, file, varNames, kwIndex);
+        StaticScope scope = StaticScopeFactory.newStaticScope(enclosingScope, type, file, varNames, -1);
+
+        if (bits[3].length() > 0) {
+            BitSet keywordIndices = new BitSet();
+            for (int i = 0; i < bits[3].length(); i++) {
+                keywordIndices.set(i);
+            }
+            scope.setKeywordIndices(keywordIndices);
+        }
 
         scope.setSignature(signature);
         scope.setScopeType(scopeType);