[ruby/yarp] Consider source encoding for slice

ruby/prism@8f59fc27cd Co-authored-by: Kevin Newton <kddnewton@users.noreply.github.com>
rm155 · Sep 6, 2023 · acd626a · acd626a
1 parent f1422e4
commit acd626a
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 6 deletions.
diff --git a/test/yarp/encoding_test.rb b/test/yarp/encoding_test.rb
@@ -96,5 +96,11 @@ def test_first_lexed_token
       encoding = YARP.lex("# encoding: ascii-8bit").value[0][0].value.encoding
       assert_equal Encoding.find("ascii-8bit"), encoding
     end
+
+    def test_slice_encoding
+      slice = YARP.parse("# encoding: Shift_JIS\nア").value.slice
+      assert_equal (+"ア").force_encoding(Encoding::SHIFT_JIS), slice
+      assert_equal Encoding::SHIFT_JIS, slice.encoding
+    end
   end
 end
diff --git a/yarp/extension.c b/yarp/extension.c
@@ -347,7 +347,7 @@ parse_input(yp_string_t *input, const char *filepath) {
     yp_node_t *node = yp_parse(&parser);
     rb_encoding *encoding = rb_enc_find(parser.encoding.name);
 
-    VALUE source = yp_source_new(&parser);
+    VALUE source = yp_source_new(&parser, encoding);
     VALUE result_argv[] = {
         yp_ast_new(&parser, node, encoding),
         parser_comments(&parser, source),

diff --git a/yarp/extension.h b/yarp/extension.h
@@ -7,7 +7,7 @@
 #include <ruby/encoding.h>
 #include "yarp.h"
 
-VALUE yp_source_new(yp_parser_t *parser);
+VALUE yp_source_new(yp_parser_t *parser, rb_encoding *encoding);
 VALUE yp_token_new(yp_parser_t *parser, yp_token_t *token, rb_encoding *encoding, VALUE source);
 VALUE yp_ast_new(yp_parser_t *parser, yp_node_t *node, rb_encoding *encoding);
 

diff --git a/yarp/templates/ext/yarp/api_node.c.erb b/yarp/templates/ext/yarp/api_node.c.erb
@@ -38,8 +38,8 @@ yp_string_new(yp_string_t *string, rb_encoding *encoding) {
 
 // Create a YARP::Source object from the given parser.
 VALUE
-yp_source_new(yp_parser_t *parser) {
-    VALUE source = rb_str_new((const char *) parser->start, parser->end - parser->start);
+yp_source_new(yp_parser_t *parser, rb_encoding *encoding) {
+    VALUE source = rb_enc_str_new((const char *) parser->start, parser->end - parser->start, encoding);
     VALUE offsets = rb_ary_new_capa(parser->newline_list.size);
 
     for (size_t index = 0; index < parser->newline_list.size; index++) {
@@ -78,7 +78,7 @@ yp_node_stack_pop(yp_node_stack_node_t **stack) {
 
 VALUE
 yp_ast_new(yp_parser_t *parser, yp_node_t *node, rb_encoding *encoding) {
-    VALUE source = yp_source_new(parser);
+    VALUE source = yp_source_new(parser, encoding);
     ID *constants = calloc(parser->constant_pool.size, sizeof(ID));
 
     for (size_t index = 0; index < parser->constant_pool.capacity; index++) {

diff --git a/yarp/templates/lib/yarp/serialize.rb.erb b/yarp/templates/lib/yarp/serialize.rb.erb
@@ -18,7 +18,13 @@ module YARP
     PATCH_VERSION = 0
 
     def self.load(input, serialized)
-      Loader.new(Source.new(input), serialized).load_result
+      input = input.dup
+      source = Source.new(input)
+      loader = Loader.new(source, serialized)
+      result = loader.load_result
+
+      input.force_encoding(loader.encoding)
+      result
     end
 
     def self.load_tokens(source, serialized)