Fixed Prism::Translation::RubyParser's comment processing

Tests were failing in Flay under Prism.

Author: zenspider

lib/prism/translation/ruby_parser.rb

@@ -35,6 +35,16 @@ def initialize(file, in_def: false, in_pattern: false)
           @in_pattern = in_pattern
         end
 
+        def attach_comments sexp, node # :nodoc:
+          return unless node.comments
+          return if node.comments.empty?
+
+          extra = node.location.start_line - node.comments.last.location.start_line
+          comments = node.comments.map(&:slice)
+          comments.concat [nil] * extra.clamp(0, nil)
+          sexp.comments = comments.join "\n"
+        end
+
         # ```
         # alias foo bar
         # ^^^^^^^^^^^^^
@@ -415,14 +425,19 @@ def visit_class_node(node)
               visit(node.constant_path)
             end
 
-          if node.body.nil?
-            s(node, :class, name, visit(node.superclass))
-          elsif node.body.is_a?(StatementsNode)
-            compiler = copy_compiler(in_def: false)
-            s(node, :class, name, visit(node.superclass)).concat(node.body.body.map { |child| child.accept(compiler) })
-          else
-            s(node, :class, name, visit(node.superclass), node.body.accept(copy_compiler(in_def: false)))
-          end
+          result =
+            if node.body.nil?
+              s(node, :class, name, visit(node.superclass))
+            elsif node.body.is_a?(StatementsNode)
+              compiler = copy_compiler(in_def: false)
+              s(node, :class, name, visit(node.superclass)).concat(node.body.body.map { |child| child.accept(compiler) })
+            else
+              s(node, :class, name, visit(node.superclass), node.body.accept(copy_compiler(in_def: false)))
+            end
+
+          attach_comments(result, node)
+
+          result
         end
 
         # ```
@@ -611,7 +626,10 @@ def visit_def_node(node)
               s(node, :defs, visit(node.receiver), name)
             end
 
+          attach_comments(result, node)
+
           result.line(node.name_loc.start_line)
+
           if node.parameters.nil?
             result << s(node, :args).line(node.name_loc.start_line)
           else
@@ -1270,14 +1288,19 @@ def visit_module_node(node)
               visit(node.constant_path)
             end
 
-          if node.body.nil?
-            s(node, :module, name)
-          elsif node.body.is_a?(StatementsNode)
-            compiler = copy_compiler(in_def: false)
-            s(node, :module, name).concat(node.body.body.map { |child| child.accept(compiler) })
-          else
-            s(node, :module, name, node.body.accept(copy_compiler(in_def: false)))
-          end
+          result =
+            if node.body.nil?
+              s(node, :module, name)
+            elsif node.body.is_a?(StatementsNode)
+              compiler = copy_compiler(in_def: false)
+              s(node, :module, name).concat(node.body.body.map { |child| child.accept(compiler) })
+            else
+              s(node, :module, name, node.body.accept(copy_compiler(in_def: false)))
+            end
+
+          attach_comments(result, node)
+
+          result
         end
 
         # ```
@@ -1898,6 +1921,16 @@ def parse_file(filepath)
         translate(Prism.parse_file(filepath, partial_script: true), filepath)
       end
 
+      # Parse the give file and translate it into the
+      # seattlerb/ruby_parser gem's Sexp format. This method is
+      # provided for API compatibility to RubyParser and takes an
+      # optional +timeout+ argument.
+      def process ruby, file="(string)", timeout=nil
+        Timeout.timeout timeout do
+          parse ruby, file
+        end
+      end
+
       class << self
         # Parse the given source and translate it into the seattlerb/ruby_parser
         # gem's Sexp format.
@@ -1922,6 +1955,8 @@ def translate(result, filepath)
           raise ::RubyParser::SyntaxError, "#{filepath}:#{error.location.start_line} :: #{error.message}"
         end
 
+        result.attach_comments!
+
         result.value.accept(Compiler.new(filepath))
       end
     end