Add ByteBasedCharSequence to deal with YARP giving us byte offsets but Truffle wanting CharSequence

Author: eregon

src/main/java/org/truffleruby/debug/TruffleDebugNodes.java

@@ -33,7 +33,6 @@
 import com.oracle.truffle.api.object.DynamicObjectLibrary;
 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.TruffleString;
 import org.graalvm.collections.Pair;
 import org.truffleruby.Layouts;
@@ -302,20 +301,15 @@ public abstract static class YARPExecuteNode extends CoreMethodArrayArgumentsNod
         @Specialization(guards = "strings.isRubyString(code)", limit = "1")
         protected Object yarpExecute(VirtualFrame frame, Object code,
                 @Cached RubyStringLibrary strings) {
-            var tstring = strings.getTString(code);
-            var encoding = strings.getEncoding(code);
-
-            return doExecute(tstring, encoding, RubyArguments.getMethod(frame));
+            return doExecute(code, RubyArguments.getMethod(frame));
         }
 
         @TruffleBoundary
-        private Object doExecute(AbstractTruffleString tstring, RubyEncoding encoding, InternalMethod method) {
-            String sourceString = TStringUtils.toJavaStringOrThrow(tstring, encoding);
-            Source source = Source.newBuilder("ruby", sourceString, "<parse_ast>").build();
+        private Object doExecute(Object code, InternalMethod method) {
             TranslatorEnvironment.resetTemporaryVariablesIndex();
 
             final RootCallTarget callTarget = RubyLanguage.getCurrentContext().getCodeLoader().parseWithYARP(
-                    new RubySource(source, source.getName()),
+                    code,
                     ParserContext.TOP_LEVEL,
                     null,
                     RubyLanguage.getCurrentContext().getRootLexicalScope(),
@@ -1465,23 +1459,22 @@ private RubyRootNode parse(String sourceCode) {
     @CoreMethod(names = "parse_with_yarp_and_dump_truffle_ast", onSingleton = true, required = 3, lowerFixnum = 3)
     public abstract static class ParseWithYARPAndDumpTruffleASTNode extends CoreMethodArrayArgumentsNode {
 
-        @Specialization
         @TruffleBoundary
-        protected Object parseAndDump(Object sourceCode, Object focusedNodeClassName, int index,
+        @Specialization(guards = "strings.isRubyString(code)", limit = "1")
+        protected Object parseAndDump(Object code, Object focusedNodeClassName, int index,
+                @Cached RubyStringLibrary strings,
                 @Cached TruffleString.FromJavaStringNode fromJavaStringNode) {
-            String sourceCodeString = RubyGuards.getJavaString(sourceCode);
             String nodeClassNameString = RubyGuards.getJavaString(focusedNodeClassName);
-            RubyRootNode rootNode = parse(sourceCodeString);
+            RubyRootNode rootNode = parse(code);
             String output = TruffleASTPrinter.dump(rootNode, nodeClassNameString, index);
             return createString(fromJavaStringNode, output, Encodings.UTF_8);
         }
 
-        private RubyRootNode parse(String sourceCode) {
-            Source source = Source.newBuilder("ruby", sourceCode, "<parse_ast>").build();
+        private RubyRootNode parse(Object code) {
             TranslatorEnvironment.resetTemporaryVariablesIndex();
 
             final RootCallTarget callTarget = RubyLanguage.getCurrentContext().getCodeLoader().parseWithYARP(
-                    new RubySource(source, source.getName()),
+                    code,
                     ParserContext.TOP_LEVEL,
                     null,
                     RubyLanguage.getCurrentContext().getRootLexicalScope(),

src/main/java/org/truffleruby/language/loader/ByteBasedCharSequence.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright (c) 2023 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.language.loader;
+
+import com.oracle.truffle.api.strings.TruffleString;
+import org.truffleruby.core.encoding.Encodings;
+import org.truffleruby.core.encoding.RubyEncoding;
+import org.truffleruby.core.string.TStringWithEncoding;
+
+/** YARP gives use byte offsets but Truffle wants a CharSequence, this tries to integrate both as much as possible.
+ * Using a java.lang.String instead would mean computing char offsets, which is prohibitively expensive. */
+public final class ByteBasedCharSequence implements CharSequence {
+
+    private final byte[] bytes;
+    private final int offset;
+    private final int length;
+    private final RubyEncoding encoding;
+
+    public ByteBasedCharSequence(TStringWithEncoding tstringWithEnc) {
+        this(tstringWithEnc.getBytesOrCopy(), 0, tstringWithEnc.byteLength(), tstringWithEnc.encoding);
+
+        // Ensure it can be converted to a Java String early
+        if (tstringWithEnc.encoding == Encodings.BINARY) {
+            tstringWithEnc.toJavaStringOrThrow();
+        }
+    }
+
+    private ByteBasedCharSequence(byte[] bytes, int offset, int length, RubyEncoding encoding) {
+        this.bytes = bytes;
+        this.offset = offset;
+        this.length = length;
+        this.encoding = encoding;
+    }
+
+    @Override
+    public int length() {
+        return bytes.length;
+    }
+
+    @Override
+    public char charAt(int index) {
+        assert index >= offset && index < offset + length;
+        return (char) bytes[offset + index];
+    }
+
+    @Override
+    public CharSequence subSequence(int start, int end) {
+        return new ByteBasedCharSequence(bytes, start, end - start, encoding);
+    }
+
+    @Override
+    public String toString() {
+        return TruffleString.fromByteArrayUncached(bytes, offset, length, encoding.tencoding, false)
+                .toJavaStringUncached();
+    }
+}

src/main/java/org/truffleruby/language/loader/CodeLoader.java

@@ -17,11 +17,13 @@
 import org.truffleruby.core.string.TStringWithEncoding;
 import org.truffleruby.language.LexicalScope;
 import org.truffleruby.language.Nil;
+import org.truffleruby.language.RubyGuards;
 import org.truffleruby.language.RubyNode;
 import org.truffleruby.language.RubyRootNode;
 import org.truffleruby.annotations.Visibility;
 import org.truffleruby.language.arguments.EmptyArgumentsDescriptor;
 import org.truffleruby.language.arguments.RubyArguments;
+import org.truffleruby.language.library.RubyStringLibrary;
 import org.truffleruby.language.methods.DeclarationContext;
 import org.truffleruby.language.methods.InternalMethod;
 import org.truffleruby.language.methods.SharedMethodInfo;
@@ -88,14 +90,18 @@ public RootCallTarget parse(RubySource source,
     }
 
     @TruffleBoundary
-    public RootCallTarget parseWithYARP(RubySource source,
+    public RootCallTarget parseWithYARP(Object code,
             ParserContext parserContext,
             MaterializedFrame parentFrame,
             LexicalScope lexicalScope,
             Node currentNode) {
-        final YARPTranslatorDriver translator = new YARPTranslatorDriver(context, source);
-        return translator
-                .parse(source, parserContext, null, parentFrame, lexicalScope, currentNode);
+        var tstringWithEnc = new TStringWithEncoding(RubyGuards.asTruffleStringUncached(code),
+                RubyStringLibrary.getUncached().getEncoding(code));
+        var charSequence = new ByteBasedCharSequence(tstringWithEnc);
+        Source source = Source.newBuilder("ruby", charSequence, "<parse_ast>").build();
+        var rubySource = new RubySource(source, source.getName(), tstringWithEnc);
+        final YARPTranslatorDriver translator = new YARPTranslatorDriver(context, rubySource);
+        return translator.parse(rubySource, parserContext, null, parentFrame, lexicalScope, currentNode);
     }
 
     @TruffleBoundary

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

@@ -19,6 +19,7 @@
 
 import com.oracle.truffle.api.source.Source;
 import org.truffleruby.core.encoding.RubyEncoding;
+import org.truffleruby.core.encoding.TStringUtils;
 import org.truffleruby.core.string.TStringWithEncoding;
 
 public final class RubySource {
@@ -29,6 +30,7 @@ public final class RubySource {
      * Sources in the future (but then we'll need to still use this path in Ruby backtraces). */
     private final String sourcePath;
     private final TruffleString code;
+    private byte[] bytes;
     private final RubyEncoding encoding;
     private final boolean isEval;
     private final int lineOffset;
@@ -69,6 +71,7 @@ public boolean hasTruffleString() {
     }
 
     public TruffleString getTruffleString() {
+        assert hasTruffleString();
         return code;
     }
 
@@ -77,6 +80,15 @@ public TStringWithEncoding getTStringWithEncoding() {
         return new TStringWithEncoding(code, encoding);
     }
 
+    public byte[] getBytes() {
+        assert hasTruffleString();
+        if (bytes != null) {
+            return bytes;
+        } else {
+            return bytes = TStringUtils.getBytesOrCopy(code, encoding);
+        }
+    }
+
     public RubyEncoding getEncoding() {
         assert hasTruffleString();
         return encoding;

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

@@ -184,7 +184,7 @@ public RootCallTarget parse(RubySource rubySource, ParserContext parserContext,
         // Needs the magic comment to be parsed
         parseEnvironment.allowTruffleRubyPrimitives = parserConfiguration.allowTruffleRubyPrimitives;
 
-        final SourceSection sourceSection = source.createSection(0, source.getCharacters().length());
+        final SourceSection sourceSection = source.createSection(0, rubySource.getBytes().length);
         final SourceIndexLength sourceIndexLength = SourceIndexLength.fromSourceSection(sourceSection);
 
         final String modulePath = staticLexicalScope == null || staticLexicalScope == context.getRootLexicalScope()
@@ -226,7 +226,7 @@ public RootCallTarget parse(RubySource rubySource, ParserContext parserContext,
 
         // Translate to Ruby Truffle nodes
 
-        byte[] sourceBytes = rubySource.getSource().getCharacters().toString().getBytes();
+        byte[] sourceBytes = rubySource.getBytes();
         final YARPTranslator translator = new YARPTranslator(
                 language,
                 null,
@@ -395,7 +395,7 @@ public static org.yarp.Nodes.Node parseToYARPAST(RubyContext context, RubyLangua
         TruffleSafepoint.poll(DummyNode.INSTANCE);
 
         // YARP begin
-        byte[] sourceBytes = rubySource.getSource().getCharacters().toString().getBytes();
+        byte[] sourceBytes = rubySource.getBytes();
         org.yarp.Parser.loadLibrary(language.getRubyHome() + "/lib/libyarp" + Platform.LIB_SUFFIX);
         byte[] serializedBytes = Parser.parseAndSerialize(sourceBytes);
 

src/main/java/org/truffleruby/parser/lexer/LexerSource.java

@@ -74,7 +74,6 @@ public LexerSource(RubySource rubySource) {
             rubyEncoding = Encodings.UTF_8;
             // TODO CS 5-Sep-17 can we get the bytes directly rather than using getCharacters -> toString -> getBytes?
             var sourceString = source.getCharacters().toString();
-            // this.sourceBytes = TStringUtils.fromByteArray(sourceString.getBytes(StandardCharsets.UTF_8), TruffleString.Encoding.UTF_8);
             this.sourceBytes = TStringUtils.utf8TString(sourceString);
         }
         this.sourceByteLength = sourceBytes.byteLength(rubyEncoding.tencoding);