Allow for substring lookups in the frozen string table.

Previously, we always extracted the precise range of bytes we needed when looking up an entry in the frozen string table. If an entry already existed, we'd discard that extracted range in favor of what was already in the cache. This change defers making a copy of the string's bytes until we need to insert into the cache. For situations with many cache hits this approach can be much faster.

Author: nirvdrum

src/main/java/org/truffleruby/core/encoding/TStringUtils.java

@@ -44,9 +44,13 @@ public static TruffleString.Encoding jcodingToTEncoding(Encoding jcoding) {
     }
 
     public static TruffleString fromByteArray(byte[] bytes, TruffleString.Encoding tencoding) {
+        return fromByteArray(bytes, 0, bytes.length, tencoding);
+    }
+
+    public static TruffleString fromByteArray(byte[] bytes, int offset, int length, TruffleString.Encoding tencoding) {
         CompilerAsserts.neverPartOfCompilation(
                 "Use createString(TruffleString.FromByteArrayNode, byte[], RubyEncoding) instead");
-        return TruffleString.fromByteArrayUncached(bytes, 0, bytes.length, tencoding, false);
+        return TruffleString.fromByteArrayUncached(bytes, offset, length, tencoding, false);
     }
 
     public static TruffleString fromByteArray(byte[] bytes, RubyEncoding rubyEncoding) {

src/main/java/org/truffleruby/core/string/FrozenStringLiterals.java

@@ -41,7 +41,16 @@ public ImmutableRubyString getFrozenStringLiteral(TruffleString tstring, RubyEnc
             throw CompilerDirectives.shouldNotReachHere();
         }
 
-        return getFrozenStringLiteral(TStringUtils.getBytesOrCopy(tstring, encoding), encoding);
+        // Ensure all ImmutableRubyString have a TruffleString from the TStringCache
+        var cachedTString = tstringCache.getTString(tstring, encoding);
+        var tstringWithEncoding = new TStringWithEncoding(cachedTString, encoding);
+
+        final ImmutableRubyString string = values.get(tstringWithEncoding);
+        if (string != null) {
+            return string;
+        } else {
+            return getFrozenStringLiteral(TStringUtils.getBytesOrCopy(tstring, encoding), encoding);
+        }
     }
 
     @TruffleBoundary

src/main/java/org/truffleruby/core/string/TBytesKey.java

@@ -12,19 +12,34 @@
 import java.util.Arrays;
 import java.util.Objects;
 
+import com.oracle.truffle.api.strings.InternalByteArray;
 import com.oracle.truffle.api.strings.TruffleString;
+import org.truffleruby.core.array.ArrayUtils;
 import org.truffleruby.core.encoding.RubyEncoding;
+import org.truffleruby.core.encoding.TStringUtils;
 
 public final class TBytesKey {
 
     private final byte[] bytes;
+    private final int offset;
+    private final int length;
     private RubyEncoding encoding;
     private final int bytesHashCode;
 
-    public TBytesKey(byte[] bytes, RubyEncoding encoding) {
+    public TBytesKey(byte[] bytes, int offset, int length, int bytesHashCode, RubyEncoding encoding) {
         this.bytes = bytes;
+        this.offset = offset;
+        this.length = length;
+        this.bytesHashCode = bytesHashCode;
         this.encoding = encoding;
-        this.bytesHashCode = Arrays.hashCode(bytes);
+    }
+
+    public TBytesKey(byte[] bytes, RubyEncoding encoding) {
+        this(bytes, 0, bytes.length, Arrays.hashCode(bytes), encoding);
+    }
+
+    public TBytesKey(InternalByteArray byteArray, RubyEncoding encoding) {
+        this(byteArray.getArray(), byteArray.getOffset(), byteArray.getLength(), hashCode(byteArray), encoding);
     }
 
     @Override
@@ -37,15 +52,15 @@ public boolean equals(Object o) {
         if (o instanceof TBytesKey) {
             final TBytesKey other = (TBytesKey) o;
             if (encoding == null) {
-                if (Arrays.equals(bytes, other.bytes)) {
+                if (equalBytes(this, other)) {
                     // For getMatchedEncoding()
                     this.encoding = Objects.requireNonNull(other.encoding);
                     return true;
                 } else {
                     return false;
                 }
             } else {
-                return encoding == other.encoding && Arrays.equals(bytes, other.bytes);
+                return encoding == other.encoding && equalBytes(this, other);
             }
         }
 
@@ -62,4 +77,48 @@ public String toString() {
         return TruffleString.fromByteArrayUncached(bytes, encoding, false).toString();
     }
 
+    private static int hashCode(InternalByteArray byteArray) {
+        return hashCode(byteArray.getArray(), byteArray.getOffset(), byteArray.getLength());
+    }
+
+    // A variant of <code>Arrays.hashCode</code> that allows for selecting a range within the array.
+    private static int hashCode(byte[] bytes, int offset, int length) {
+        if (bytes == null) {
+            return 0;
+        }
+
+        int result = 1;
+        for (int i = offset; i < offset + length; i++) {
+            result = 31 * result + bytes[i];
+        }
+
+        return result;
+    }
+
+    private boolean equalBytes(TBytesKey a, TBytesKey b) {
+        return Arrays.equals(a.bytes, a.offset, a.offset + a.length, b.bytes, b.offset, b.offset + b.length);
+    }
+
+    private boolean isPerfectFit() {
+        return offset == 0 && length == bytes.length;
+    }
+
+    public TBytesKey makeCacheable() {
+        if (isPerfectFit()) {
+            // TODO (nirvdrum 2023-Jun-17): We can avoid cloning the key if we know the byte array came from an immutable string.
+            return new TBytesKey(bytes.clone(), encoding);
+        }
+
+        var simplified = ArrayUtils.extractRange(this.bytes, this.offset, this.offset + this.length);
+        return new TBytesKey(simplified, encoding);
+    }
+
+    public TBytesKey withNewEncoding(RubyEncoding encoding) {
+        return new TBytesKey(bytes, offset, length, bytesHashCode, encoding);
+    }
+
+    public TruffleString toTruffleString() {
+        return TStringUtils.fromByteArray(bytes, offset, length, encoding.tencoding);
+    }
+
 }

src/main/java/org/truffleruby/core/string/TStringCache.java

@@ -69,20 +69,30 @@ private void register(TruffleString tstring, RubyEncoding encoding) {
         }
     }
 
-    public TruffleString getTString(TruffleString string, RubyEncoding encoding) {
-        return getTString(TStringUtils.getBytesOrCopy(string, encoding), encoding);
+    @TruffleBoundary
+    public TruffleString getTString(TruffleString string, RubyEncoding rubyEncoding) {
+        assert rubyEncoding != null;
+
+        var byteArray = string.getInternalByteArrayUncached(rubyEncoding.tencoding);
+        final TBytesKey key = new TBytesKey(byteArray, rubyEncoding);
+
+        return getTString(key);
     }
 
-    @TruffleBoundary
     public TruffleString getTString(byte[] bytes, RubyEncoding rubyEncoding) {
         assert rubyEncoding != null;
 
-        final TBytesKey key = new TBytesKey(bytes, rubyEncoding);
+        return getTString(new TBytesKey(bytes, rubyEncoding));
+    }
+
+    @TruffleBoundary
+    private TruffleString getTString(TBytesKey lookupKey) {
+        final TruffleString tstring = bytesToTString.get(lookupKey);
+        var rubyEncoding = lookupKey.getMatchedEncoding();
 
-        final TruffleString tstring = bytesToTString.get(key);
         if (tstring != null) {
             ++tstringsReusedCount;
-            tstringBytesSaved += tstring.byteLength(rubyEncoding.tencoding);
+            tstringBytesSaved += tstring.byteLength(lookupKey.getMatchedEncoding().tencoding);
 
             return tstring;
         }
@@ -92,7 +102,7 @@ public TruffleString getTString(byte[] bytes, RubyEncoding rubyEncoding) {
         // reference equality optimizations. So, do another search but with a marker encoding. The only guarantee
         // we can make about the resulting TruffleString is that it would have the same logical byte[], but that's good enough
         // for our purposes.
-        TBytesKey keyNoEncoding = new TBytesKey(bytes, null);
+        TBytesKey keyNoEncoding = lookupKey.withNewEncoding(null);
         final TruffleString tstringWithSameBytesButDifferentEncoding = bytesToTString.get(keyNoEncoding);
 
         final TruffleString newTString;
@@ -104,12 +114,11 @@ public TruffleString getTString(byte[] bytes, RubyEncoding rubyEncoding) {
             ++byteArrayReusedCount;
             tstringBytesSaved += newTString.byteLength(rubyEncoding.tencoding);
         } else {
-            newTString = TStringUtils.fromByteArray(bytes, rubyEncoding);
+            newTString = lookupKey.toTruffleString();
         }
 
         // Use the new TruffleString bytes in the cache, so we do not keep bytes alive unnecessarily.
-        final TBytesKey newKey = new TBytesKey(TStringUtils.getBytesOrCopy(newTString, rubyEncoding), rubyEncoding);
-        return bytesToTString.addInCacheIfAbsent(newKey, newTString);
+        return bytesToTString.addInCacheIfAbsent(lookupKey.makeCacheable(), newTString);
     }
 
     public boolean contains(TruffleString string, RubyEncoding encoding) {