3 files changed, 297 insertions, 43 deletions

diff --git a/lib/prism/parse_result/comments.rb b/lib/prism/parse_result/comments.rb
index 7a3a47de50..df80792d39 100644
--- a/lib/prism/parse_result/comments.rb
+++ b/lib/prism/parse_result/comments.rb
@@ -1,7 +1,10 @@
 # frozen_string_literal: true
+# :markup: markdown
+#--
+# rbs_inline: enabled
 
 module Prism
-  class ParseResult
+  class ParseResult < Result
     # When we've parsed the source, we have both the syntax tree and the list of
     # comments that we found in the source. This class is responsible for
     # walking the tree and finding the nearest location to attach each comment.
@@ -17,82 +20,125 @@ module Prism
     # the comment. Otherwise it will favor attaching to the nearest location
     # that is after the comment.
     class Comments
+      # @rbs!
+      #    # An internal interface for a target that comments can be attached
+      #    # to. This is either going to be a NodeTarget or a CommentTarget.
+      #    interface _CommentTarget
+      #      def start_offset: () -> Integer
+      #      def end_offset: () -> Integer
+      #      def encloses?: (Comment) -> bool
+      #      def leading_comment: (Comment) -> void
+      #      def trailing_comment: (Comment) -> void
+      #    end
+
       # A target for attaching comments that is based on a specific node's
       # location.
       class NodeTarget # :nodoc:
-        attr_reader :node
+        attr_reader :node #: node
 
+        #: (node node) -> void
         def initialize(node)
           @node = node
         end
 
+        #: () -> Integer
         def start_offset
-          node.location.start_offset
+          node.start_offset
         end
 
+        #: () -> Integer
         def end_offset
-          node.location.end_offset
+          node.end_offset
         end
 
+        #: (Comment comment) -> bool
         def encloses?(comment)
           start_offset <= comment.location.start_offset &&
             comment.location.end_offset <= end_offset
         end
 
-        def <<(comment)
-          node.location.comments << comment
+        #: (Comment comment) -> void
+        def leading_comment(comment)
+          node.location.leading_comment(comment)
+        end
+
+        #: (Comment comment) -> void
+        def trailing_comment(comment)
+          node.location.trailing_comment(comment)
         end
       end
 
       # A target for attaching comments that is based on a location field on a
       # node. For example, the `end` token of a ClassNode.
       class LocationTarget # :nodoc:
-        attr_reader :location
+        attr_reader :location #: Location
 
+        #: (Location location) -> void
         def initialize(location)
           @location = location
         end
 
+        #: () -> Integer
         def start_offset
           location.start_offset
         end
 
+        #: () -> Integer
         def end_offset
           location.end_offset
         end
 
+        #: (Comment comment) -> bool
         def encloses?(comment)
           false
         end
 
-        def <<(comment)
-          location.comments << comment
+        #: (Comment comment) -> void
+        def leading_comment(comment)
+          location.leading_comment(comment)
+        end
+
+        #: (Comment comment) -> void
+        def trailing_comment(comment)
+          location.trailing_comment(comment)
         end
       end
 
       # The parse result that we are attaching comments to.
-      attr_reader :parse_result
+      attr_reader :parse_result #: ParseResult
 
       # Create a new Comments object that will attach comments to the given
       # parse result.
+      #--
+      #: (ParseResult parse_result) -> void
       def initialize(parse_result)
         @parse_result = parse_result
       end
 
       # Attach the comments to their respective locations in the tree by
       # mutating the parse result.
+      #--
+      #: () -> void
       def attach!
         parse_result.comments.each do |comment|
           preceding, enclosing, following = nearest_targets(parse_result.value, comment)
-          target =
-            if comment.trailing?
-              preceding || following || enclosing || NodeTarget.new(parse_result.value)
+
+          if comment.trailing?
+            if preceding
+              preceding.trailing_comment(comment)
             else
-              # If a comment exists on its own line, prefer a leading comment.
-              following || preceding || enclosing || NodeTarget.new(parse_result.value)
+              (following || enclosing || NodeTarget.new(parse_result.value)).leading_comment(comment)
             end
-
-          target << comment
+          else
+            # If a comment exists on its own line, prefer a leading comment.
+            if following
+              following.leading_comment(comment)
+            elsif preceding
+              preceding.trailing_comment(comment)
+            else
+              (enclosing || NodeTarget.new(parse_result.value)).leading_comment(comment)
+            end
+          end
         end
       end
 
@@ -100,11 +146,13 @@ module Prism
 
       # Responsible for finding the nearest targets to the given comment within
       # the context of the given encapsulating node.
+      #--
+      #: (node node, Comment comment) -> [_CommentTarget?, _CommentTarget?, _CommentTarget?]
       def nearest_targets(node, comment)
         comment_start = comment.location.start_offset
         comment_end = comment.location.end_offset
 
-        targets = []
+        targets = [] #: Array[_CommentTarget]
         node.comment_targets.map do |value|
           case value
           when StatementsNode
@@ -117,8 +165,8 @@ module Prism
         end
 
         targets.sort_by!(&:start_offset)
-        preceding = nil
-        following = nil
+        preceding = nil #: _CommentTarget?
+        following = nil #: _CommentTarget?
 
         left = 0
         right = targets.length
@@ -134,6 +182,7 @@ module Prism
           target_end = target.end_offset
 
           if target.encloses?(comment)
+            # @type var target: NodeTarget
             # The comment is completely contained by this target. Abandon the
             # binary search at this level.
             return nearest_targets(target.node, comment)
@@ -166,12 +215,5 @@ module Prism
         [preceding, NodeTarget.new(node), following]
       end
     end
-
-    private_constant :Comments
-
-    # Attach the list of comments to their respective locations in the tree.
-    def attach_comments!
-      Comments.new(self).attach!
-    end
   end
 end
diff --git a/lib/prism/parse_result/errors.rb b/lib/prism/parse_result/errors.rb
new file mode 100644
index 0000000000..388309d23d
--- /dev/null
+++ b/lib/prism/parse_result/errors.rb
@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+# :markup: markdown
+#--
+# rbs_inline: enabled
+
+require "stringio"
+
+module Prism
+  class ParseResult < Result
+    # An object to represent the set of errors on a parse result. This object
+    # can be used to format the errors in a human-readable way.
+    class Errors
+      # The parse result that contains the errors.
+      attr_reader :parse_result #: ParseResult
+
+      # Initialize a new set of errors from the given parse result.
+      #--
+      #: (ParseResult parse_result) -> void
+      def initialize(parse_result)
+        @parse_result = parse_result
+      end
+
+      # Formats the errors in a human-readable way and return them as a string.
+      #--
+      #: () -> String
+      def format
+        error_lines = {} #: Hash[Integer, Array[ParseError]]
+        parse_result.errors.each do |error|
+          location = error.location
+          (location.start_line..location.end_line).each do |line|
+            error_lines[line] ||= []
+            error_lines[line] << error
+          end
+        end
+
+        source_lines = parse_result.source.source.lines
+        source_lines << "" if error_lines.key?(source_lines.size + 1)
+
+        io = StringIO.new
+        source_lines.each.with_index(1) do |line, line_number|
+          io.puts(line)
+
+          (error_lines.delete(line_number) || []).each do |error|
+            location = error.location
+
+            case line_number
+            when location.start_line
+              io.print(" " * location.start_column + "^")
+
+              if location.start_line == location.end_line
+                if location.start_column != location.end_column
+                  io.print("~" * (location.end_column - location.start_column - 1))
+                end
+
+                io.puts(" " + error.message)
+              else
+                io.puts("~" * (line.bytesize - location.start_column))
+              end
+            when location.end_line
+              io.puts("~" * location.end_column + " " + error.message)
+            else
+              io.puts("~" * line.bytesize)
+            end
+          end
+        end
+
+        io.puts
+        io.string
+      end
+    end
+  end
+end
diff --git a/lib/prism/parse_result/newlines.rb b/lib/prism/parse_result/newlines.rb
index ca05f5b702..450c790226 100644
--- a/lib/prism/parse_result/newlines.rb
+++ b/lib/prism/parse_result/newlines.rb
@@ -1,7 +1,10 @@
 # frozen_string_literal: true
+# :markup: markdown
+#--
+# rbs_inline: enabled
 
 module Prism
-  class ParseResult
+  class ParseResult < Result
     # The :line tracepoint event gets fired whenever the Ruby VM encounters an
     # expression on a new line. The types of expressions that can trigger this
     # event are:
@@ -17,48 +20,185 @@ module Prism
     # Note that the logic in this file should be kept in sync with the Java
     # MarkNewlinesVisitor, since that visitor is responsible for marking the
     # newlines for JRuby/TruffleRuby.
+    #
+    # This file is autoloaded only when `mark_newlines!` is called, so the
+    # re-opening of the various nodes in this file will only be performed in
+    # that case. We do that to avoid storing the extra `@newline` instance
+    # variable on every node if we don't need it.
     class Newlines < Visitor
+      # The map of lines indices to whether or not they have been marked as
+      # emitting a newline event.
+      # @rbs @lines: Array[bool]
+
       # Create a new Newlines visitor with the given newline offsets.
-      def initialize(newline_marked)
-        @newline_marked = newline_marked
+      #--
+      #: (Integer lines) -> void
+      def initialize(lines)
+        @lines = Array.new(1 + lines, false)
       end
 
-      # Permit block/lambda nodes to mark newlines within themselves.
+      # Permit block nodes to mark newlines within themselves.
+      #--
+      #: (BlockNode node) -> void
       def visit_block_node(node)
-        old_newline_marked = @newline_marked
-        @newline_marked = Array.new(old_newline_marked.size, false)
+        old_lines = @lines
+        @lines = Array.new(old_lines.size, false)
 
         begin
           super(node)
         ensure
-          @newline_marked = old_newline_marked
+          @lines = old_lines
         end
       end
 
-      alias_method :visit_lambda_node, :visit_block_node
+      # Permit lambda nodes to mark newlines within themselves.
+      #--
+      #: (LambdaNode node) -> void
+      def visit_lambda_node(node)
+        old_lines = @lines
+        @lines = Array.new(old_lines.size, false)
 
-      # Mark if/unless nodes as newlines.
+        begin
+          super(node)
+        ensure
+          @lines = old_lines
+        end
+      end
+
+      # Mark if nodes as newlines.
+      #--
+      #: (IfNode node) -> void
       def visit_if_node(node)
-        node.set_newline_flag(@newline_marked)
+        node.newline_flag!(@lines)
         super(node)
       end
 
-      alias_method :visit_unless_node, :visit_if_node
+      # Mark unless nodes as newlines.
+      #--
+      #: (UnlessNode node) -> void
+      def visit_unless_node(node)
+        node.newline_flag!(@lines)
+        super(node)
+      end
 
       # Permit statements lists to mark newlines within themselves.
+      #--
+      #: (StatementsNode node) -> void
       def visit_statements_node(node)
         node.body.each do |child|
-          child.set_newline_flag(@newline_marked)
+          child.newline_flag!(@lines)
         end
         super(node)
       end
     end
+  end
+
+  class Node
+    # Tracks whether or not this node should emit a newline event when the
+    # instructions that it represents are executed.
+    # @rbs @newline_flag: bool
+
+    #: () -> bool
+    def newline_flag? # :nodoc:
+      !!defined?(@newline_flag)
+    end
+
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      line = location.start_line
+      unless lines[line]
+        lines[line] = true
+        @newline_flag = true
+      end
+    end
+  end
+
+  class BeginNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      # Never mark BeginNode with a newline flag, mark children instead.
+    end
+  end
+
+  class ParenthesesNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      # Never mark ParenthesesNode with a newline flag, mark children instead.
+    end
+  end
+
+  class IfNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      predicate.newline_flag!(lines)
+    end
+  end
 
-    private_constant :Newlines
+  class UnlessNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      predicate.newline_flag!(lines)
+    end
+  end
+
+  class UntilNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      predicate.newline_flag!(lines)
+    end
+  end
+
+  class WhileNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      predicate.newline_flag!(lines)
+    end
+  end
+
+  class RescueModifierNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      expression.newline_flag!(lines)
+    end
+  end
+
+  class InterpolatedMatchLastLineNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      first = parts.first
+      first.newline_flag!(lines) if first
+    end
+  end
+
+  class InterpolatedRegularExpressionNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      first = parts.first
+      first.newline_flag!(lines) if first
+    end
+  end
+
+  class InterpolatedStringNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      first = parts.first
+      first.newline_flag!(lines) if first
+    end
+  end
+
+  class InterpolatedSymbolNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      first = parts.first
+      first.newline_flag!(lines) if first
+    end
+  end
 
-    # Walk the tree and mark nodes that are on a new line.
-    def mark_newlines!
-      value.accept(Newlines.new(Array.new(1 + source.offsets.size, false)))
+  class InterpolatedXStringNode < Node
+    #: (Array[bool] lines) -> void
+    def newline_flag!(lines) # :nodoc:
+      first = parts.first
+      first.newline_flag!(lines) if first
     end
   end
 end