Remove unnecessary data copies when working with the frozen string table.

By restructuring to work with `InternalByteArray`, we can now support working with native strings as well. The `InternalByteArray` will make a copy of the native memory into a Java `byte[]`, which won't change behind our backs.

Author: nirvdrum

CHANGELOG.md

@@ -7,7 +7,7 @@ New features:
 Bug fixes:
 
 * Fix `Dir.glob` returning blank string entry with leading `**/` in glob and `base:` argument (@rwstauner).
-* Fix class lookup after an object's class has been replaced by `IO#reopen` (@itarato, @eregon).
+* Fix class lookup after an object's class has been replaced by `IO#reopen` (@itarato, @nirvdrum, @eregon).
 * Fix `Marshal.load` and raise `ArgumentError` when dump is broken and is too short (#3108, @andrykonchin).
 * Fix `super` method lookup for unbounded attached methods (#3131,  @itarato).
 * Fix `Module#define_method(name, Method)` to respect `module_function` visibility (#3181, @andrykonchin).
@@ -43,6 +43,7 @@ Compatibility:
 Performance:
 
 * Improve `Truffle::FeatureLoader.loaded_feature_path` by removing expensive string ops from a loop. Speeds up feature lookup time (#3010, @itarato).
+* Improve `String#-@` performance by reducing unnecessary data copying and supporting substring lookups (@nirvdrum)
 
 Changes:
 

src/main/java/org/truffleruby/RubyLanguage.java

@@ -33,6 +33,7 @@
 import com.oracle.truffle.api.source.Source;
 import com.oracle.truffle.api.source.SourceSection;
 import com.oracle.truffle.api.strings.AbstractTruffleString;
+import com.oracle.truffle.api.strings.InternalByteArray;
 import com.oracle.truffle.api.strings.TruffleString;
 import org.graalvm.options.OptionDescriptors;
 import org.truffleruby.annotations.SuppressFBWarnings;
@@ -788,6 +789,11 @@ public ImmutableRubyString getFrozenStringLiteral(TruffleString tstring, RubyEnc
         return frozenStringLiterals.getFrozenStringLiteral(tstring, encoding);
     }
 
+    public ImmutableRubyString getFrozenStringLiteral(InternalByteArray byteArray, boolean isImmutable,
+            RubyEncoding encoding) {
+        return frozenStringLiterals.getFrozenStringLiteral(byteArray, isImmutable, encoding);
+    }
+
     public long getNextObjectID() {
         final long id = nextObjectID.getAndAdd(ObjectSpaceManager.OBJECT_ID_INCREMENT_BY);
 

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

@@ -152,4 +152,10 @@ public static String toJavaStringOrThrow(AbstractTruffleString tstring, RubyEnco
             return tstring.toJavaStringUncached();
         }
     }
+
+    public static boolean hasImmutableInternalByteArray(AbstractTruffleString string) {
+        // Immutable strings trivially have immutable byte arrays.
+        // Native strings also have immutable byte arrays because we need to copy the data into Java.
+        return string.isImmutable() || string.isNative();
+    }
 }

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

@@ -11,6 +11,7 @@
 
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import com.oracle.truffle.api.strings.InternalByteArray;
 import com.oracle.truffle.api.strings.TruffleString;
 import org.truffleruby.collections.WeakValueCache;
 import org.truffleruby.core.encoding.RubyEncoding;
@@ -37,34 +38,23 @@ public FrozenStringLiterals(TStringCache tStringCache) {
 
     @TruffleBoundary
     public ImmutableRubyString getFrozenStringLiteral(TruffleString tstring, RubyEncoding encoding) {
-        if (tstring.isNative()) {
-            throw CompilerDirectives.shouldNotReachHere();
-        }
-
-        // 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);
-        }
+        return getFrozenStringLiteral(tstring.getInternalByteArrayUncached(encoding.tencoding),
+                TStringUtils.hasImmutableInternalByteArray(tstring),
+                encoding);
     }
 
     @TruffleBoundary
-    public ImmutableRubyString getFrozenStringLiteral(byte[] bytes, RubyEncoding encoding) {
+    public ImmutableRubyString getFrozenStringLiteral(InternalByteArray byteArray, boolean isImmutable,
+            RubyEncoding encoding) {
         // Ensure all ImmutableRubyString have a TruffleString from the TStringCache
-        var cachedTString = tstringCache.getTString(bytes, encoding);
+        var cachedTString = tstringCache.getTString(byteArray, isImmutable, encoding);
         var tstringWithEncoding = new TStringWithEncoding(cachedTString, encoding);
 
         final ImmutableRubyString string = values.get(tstringWithEncoding);
         if (string != null) {
             return string;
         } else {
-            return values.addInCacheIfAbsent(tstringWithEncoding,
-                    new ImmutableRubyString(cachedTString, encoding));
+            return values.addInCacheIfAbsent(tstringWithEncoding, new ImmutableRubyString(cachedTString, encoding));
         }
     }
 

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

@@ -4357,10 +4357,11 @@ public abstract static class InternNode extends PrimitiveArrayArgumentsNode {
         @Specialization
         protected ImmutableRubyString internString(RubyString string,
                 @Cached RubyStringLibrary libString,
-                @Cached TruffleString.AsManagedNode asManagedNode) {
+                @Cached TruffleString.GetInternalByteArrayNode getInternalByteArrayNode) {
             var encoding = libString.getEncoding(string);
-            TruffleString immutableManagedString = asManagedNode.execute(string.tstring, encoding.tencoding);
-            return getLanguage().getFrozenStringLiteral(immutableManagedString, encoding);
+            var byteArray = getInternalByteArrayNode.execute(string.tstring, encoding.tencoding);
+            return getLanguage().getFrozenStringLiteral(byteArray,
+                    TStringUtils.hasImmutableInternalByteArray(string.tstring), encoding);
         }
     }
 

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

@@ -23,23 +23,37 @@ public final class TBytesKey {
     private final byte[] bytes;
     private final int offset;
     private final int length;
+    private final boolean isImmutable;
     private RubyEncoding encoding;
     private final int bytesHashCode;
 
-    public TBytesKey(byte[] bytes, int offset, int length, int bytesHashCode, RubyEncoding encoding) {
+    public TBytesKey(
+            byte[] bytes,
+            int offset,
+            int length,
+            boolean isImmutable,
+            int bytesHashCode,
+            RubyEncoding encoding) {
         this.bytes = bytes;
         this.offset = offset;
         this.length = length;
+        this.isImmutable = isImmutable;
         this.bytesHashCode = bytesHashCode;
         this.encoding = encoding;
     }
 
     public TBytesKey(byte[] bytes, RubyEncoding encoding) {
-        this(bytes, 0, bytes.length, Arrays.hashCode(bytes), encoding);
+        this(bytes, 0, bytes.length, true, Arrays.hashCode(bytes), encoding);
     }
 
-    public TBytesKey(InternalByteArray byteArray, RubyEncoding encoding) {
-        this(byteArray.getArray(), byteArray.getOffset(), byteArray.getLength(), hashCode(byteArray), encoding);
+    public TBytesKey(InternalByteArray byteArray, boolean isImmutable, RubyEncoding encoding) {
+        this(
+                byteArray.getArray(),
+                byteArray.getOffset(),
+                byteArray.getLength(),
+                isImmutable,
+                hashCode(byteArray),
+                encoding);
     }
 
     @Override
@@ -104,17 +118,16 @@ private boolean isPerfectFit() {
     }
 
     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);
+        if (isImmutable && isPerfectFit()) {
+            return new TBytesKey(bytes, 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);
+        return new TBytesKey(bytes, offset, length, isImmutable, bytesHashCode, encoding);
     }
 
     public TruffleString toTruffleString() {

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

@@ -9,6 +9,7 @@
  */
 package org.truffleruby.core.string;
 
+import com.oracle.truffle.api.strings.InternalByteArray;
 import com.oracle.truffle.api.strings.TruffleString;
 import org.truffleruby.collections.WeakValueCache;
 import org.truffleruby.core.encoding.Encodings;
@@ -74,11 +75,20 @@ public TruffleString getTString(TruffleString string, RubyEncoding rubyEncoding)
         assert rubyEncoding != null;
 
         var byteArray = string.getInternalByteArrayUncached(rubyEncoding.tencoding);
-        final TBytesKey key = new TBytesKey(byteArray, rubyEncoding);
+        final TBytesKey key = new TBytesKey(byteArray, TStringUtils.hasImmutableInternalByteArray(string),
+                rubyEncoding);
 
         return getTString(key);
     }
 
+    @TruffleBoundary
+    public TruffleString getTString(InternalByteArray byteArray, boolean isImmutable, RubyEncoding rubyEncoding) {
+        assert rubyEncoding != null;
+
+        return getTString(new TBytesKey(byteArray, isImmutable, rubyEncoding));
+    }
+
+    @TruffleBoundary
     public TruffleString getTString(byte[] bytes, RubyEncoding rubyEncoding) {
         assert rubyEncoding != null;