Preserve comments inside Sorbet signatures

lib/rbi/parser.rb

@@ -93,7 +93,8 @@ def parse(source, file:)
         raise ParseError.new(message, location)
       end
 
-      visitor = TreeBuilder.new(source, comments: result.comments, file: file)
+      comments_by_line = result.comments.to_h { |c| [c.location.start_line, c] }
+      visitor = TreeBuilder.new(source, comments_by_line: comments_by_line, file: file)
       visitor.visit(result.value)
       visitor.tree
     rescue ParseError => e
@@ -112,12 +113,13 @@ def parse(source, file:)
     end
 
     class Visitor < Prism::Visitor
-      #: (String source, file: String) -> void
-      def initialize(source, file:)
+      #: (String source, file: String, ?comments_by_line: Hash[Integer, Prism::Comment]) -> void
+      def initialize(source, file:, comments_by_line: {})
         super()
 
         @source = source
         @file = file
+        @comments_by_line = comments_by_line
       end
 
       private
@@ -174,6 +176,146 @@ def self?(node)
       def t_sig_without_runtime?(node)
         !!(node.is_a?(Prism::ConstantPathNode) && node_string(node) =~ /(::)?T::Sig::WithoutRuntime/)
       end
+
+      #: (Prism::Node node, ?min_line: Integer?) -> Array[Comment]
+      def node_comments(node, min_line: nil)
+        comments = []
+
+        latest_line, earliest_line = comment_lookup_range(node, min_line: min_line)
+
+        rbs_continuation = [] #: Array[Prism::Comment]
+
+        latest_line.downto(earliest_line) do |line|
+          comment = @comments_by_line[line]
+          break unless comment
+
+          text = comment.location.slice
+
+          # If we find a RBS comment continuation `#|`, we store it until we find the start with `#:`
+          if text.start_with?("#|")
+            rbs_continuation << comment
+            @comments_by_line.delete(line)
+            next
+          end
+
+          loc = Loc.from_prism(@file, comment.location)
+
+          # If we find the start of a RBS comment, we create a new RBSComment
+          # Note that we ignore RDoc directives such as `:nodoc:`
+          # See https://ruby.github.io/rdoc/RDoc/MarkupReference.html#class-RDoc::MarkupReference-label-Directives
+          if text.start_with?("#:") && !(text =~ /^#:[a-z_]+:/)
+            text = text.sub(/^#: ?/, "").rstrip
+
+            # If we found continuation comments, we merge them in reverse order (since we go from bottom to top)
+            rbs_continuation.reverse_each do |rbs_comment|
+              continuation_text = rbs_comment.location.slice.sub(/^#\| ?/, "").strip
+              continuation_loc = Loc.from_prism(@file, rbs_comment.location)
+              loc = loc.join(continuation_loc)
+              text = "#{text}#{continuation_text}"
+            end
+
+            rbs_continuation.clear
+            comments.unshift(RBSComment.new(text, loc: loc))
+          else
+            # If we have unused continuation comments, we should inject them back to not lose them
+            rbs_continuation.each do |rbs_comment|
+              comments.unshift(parse_comment(rbs_comment))
+            end
+
+            rbs_continuation.clear
+            comments.unshift(parse_comment(comment))
+          end
+
+          @comments_by_line.delete(line)
+        end
+
+        # If we have unused continuation comments, we should inject them back to not lose them
+        rbs_continuation.each do |rbs_comment|
+          comments.unshift(parse_comment(rbs_comment))
+        end
+        rbs_continuation.clear
+
+        comments
+      end
+
+      #: (Prism::Node node, ?min_line: Integer?) -> [Integer, Integer]
+      def comment_lookup_range(node, min_line: nil)
+        node_start_line = node.location.start_line
+        latest_line = node_start_line
+        comment = @comments_by_line[latest_line]
+        # Start comment lookup from line above the node if there are no trailing comments.
+        latest_line -= 1 unless comment && inline_comment_for_node?(node, comment)
+
+        earliest_line = min_line || 1
+        # End comment lookup at the line after `min_line` if `min_line` solely
+        # belongs to the surrounding node.
+        # The parent comment should not attach to the `foo` node.
+        # Example:
+        # `parent( # parent comment`
+        # `  foo: Integer # foo comment`
+        earliest_line += 1 if min_line && min_line < node_start_line
+
+        [latest_line, earliest_line]
+      end
+
+      #: (Prism::Node node, Prism::Comment comment) -> bool
+      def inline_comment_for_node?(node, comment)
+        return false unless node.location.start_line == comment.location.start_line
+
+        between = @source[node.location.end_offset...comment.location.start_offset] || ""
+        between.each_char.all? do |char|
+          char == "," || char.match?(/\s/)
+        end
+      end
+
+      #: (Prism::Node node) -> Array[Comment]
+      def comments_inside(node)
+        comments = [] #: Array[Comment]
+
+        node.location.start_line.upto(node.location.end_line) do |line|
+          comment = @comments_by_line[line]
+          next unless comment
+
+          location = comment.location
+          next if location.start_offset < node.location.start_offset
+          next if location.end_offset > node.location.end_offset
+
+          comments << parse_comment(comment)
+          @comments_by_line.delete(line)
+        end
+
+        comments
+      end
+
+      #: (Loc? inner, Loc? outer) -> bool
+      def loc_inside?(inner, outer)
+        return false unless inner && outer
+
+        inner_begin_line = inner.begin_line
+        inner_begin_column = inner.begin_column
+        inner_end_line = inner.end_line
+        inner_end_column = inner.end_column
+        outer_begin_line = outer.begin_line
+        outer_begin_column = outer.begin_column
+        outer_end_line = outer.end_line
+        outer_end_column = outer.end_column
+        return false unless inner_begin_line && inner_begin_column && inner_end_line && inner_end_column
+        return false unless outer_begin_line && outer_begin_column && outer_end_line && outer_end_column
+
+        begins_inside = inner_begin_line > outer_begin_line ||
+          (inner_begin_line == outer_begin_line && inner_begin_column >= outer_begin_column)
+        ends_inside = inner_end_line < outer_end_line ||
+          (inner_end_line == outer_end_line && inner_end_column <= outer_end_column)
+
+        begins_inside && ends_inside
+      end
+
+      #: (Prism::Comment node) -> Comment
+      def parse_comment(node)
+        text = node.location.slice.sub(/^# ?/, "").rstrip
+        loc = Loc.from_prism(@file, node.location)
+        Comment.new(text, loc: loc)
+      end
     end
 
     class TreeBuilder < Visitor
@@ -183,11 +325,15 @@ class TreeBuilder < Visitor
       #: Prism::Node?
       attr_reader :last_node
 
-      #: (String source, comments: Array[Prism::Comment], file: String) -> void
-      def initialize(source, comments:, file:)
-        super(source, file: file)
+      #: (String source, file: String, ?comments: Array[Prism::Comment]?,
+      #| ?comments_by_line: Hash[Integer, Prism::Comment]) -> void
+      def initialize(source, file:, comments: nil, comments_by_line: {})
+        if comments
+          comments_by_line = comments.to_h { |comment| [comment.location.start_line, comment] }
+        end
+
+        super(source, file: file, comments_by_line: comments_by_line)
 
-        @comments_by_line = comments.to_h { |c| [c.location.start_line, c] } #: Hash[Integer, Prism::Comment]
         @tree = Tree.new #: Tree
 
         @scopes_stack = [@tree] #: Array[Tree]
@@ -608,82 +754,14 @@ def detach_comments_from_sigs(sigs)
         comments = [] #: Array[Comment]
 
         sigs.each do |sig|
-          comments += sig.comments.dup
-          sig.comments.clear
+          inside, outside = sig.comments.partition { |comment| loc_inside?(comment.loc, sig.loc) }
+          comments.concat(outside)
+          sig.comments.replace(inside)
         end
 
         comments
       end
 
-      #: (Prism::Node node) -> Array[Comment]
-      def node_comments(node)
-        comments = []
-
-        start_line = node.location.start_line
-        start_line -= 1 unless @comments_by_line.key?(start_line)
-
-        rbs_continuation = [] #: Array[Prism::Comment]
-
-        start_line.downto(1) do |line|
-          comment = @comments_by_line[line]
-          break unless comment
-
-          text = comment.location.slice
-
-          # If we find a RBS comment continuation `#|`, we store it until we find the start with `#:`
-          if text.start_with?("#|")
-            rbs_continuation << comment
-            @comments_by_line.delete(line)
-            next
-          end
-
-          loc = Loc.from_prism(@file, comment.location)
-
-          # If we find the start of a RBS comment, we create a new RBSComment
-          # Note that we ignore RDoc directives such as `:nodoc:`
-          # See https://ruby.github.io/rdoc/RDoc/MarkupReference.html#class-RDoc::MarkupReference-label-Directives
-          if text.start_with?("#:") && !(text =~ /^#:[a-z_]+:/)
-            text = text.sub(/^#: ?/, "").rstrip
-
-            # If we found continuation comments, we merge them in reverse order (since we go from bottom to top)
-            rbs_continuation.reverse_each do |rbs_comment|
-              continuation_text = rbs_comment.location.slice.sub(/^#\| ?/, "").strip
-              continuation_loc = Loc.from_prism(@file, rbs_comment.location)
-              loc = loc.join(continuation_loc)
-              text = "#{text}#{continuation_text}"
-            end
-
-            rbs_continuation.clear
-            comments.unshift(RBSComment.new(text, loc: loc))
-          else
-            # If we have unused continuation comments, we should inject them back to not lose them
-            rbs_continuation.each do |rbs_comment|
-              comments.unshift(parse_comment(rbs_comment))
-            end
-
-            rbs_continuation.clear
-            comments.unshift(parse_comment(comment))
-          end
-
-          @comments_by_line.delete(line)
-        end
-
-        # If we have unused continuation comments, we should inject them back to not lose them
-        rbs_continuation.each do |rbs_comment|
-          comments.unshift(parse_comment(rbs_comment))
-        end
-        rbs_continuation.clear
-
-        comments
-      end
-
-      #: (Prism::Comment node) -> Comment
-      def parse_comment(node)
-        text = node.location.slice.sub(/^# ?/, "").rstrip
-        loc = Loc.from_prism(@file, node.location)
-        Comment.new(text, loc: loc)
-      end
-
       #: (Prism::Node? node) -> Array[Arg]
       def parse_send_args(node)
         args = [] #: Array[Arg]
@@ -783,10 +861,10 @@ def parse_params(node)
 
       #: (Prism::CallNode node) -> Sig
       def parse_sig(node)
-        builder = SigBuilder.new(@source, file: @file)
+        builder = SigBuilder.new(@source, comments_by_line: @comments_by_line, file: @file)
         builder.current.loc = node_loc(node)
         builder.visit_call_node(node)
-        builder.current.comments = node_comments(node)
+        builder.current.comments = node_comments(node) + builder.current.comments
         builder.current
       end
 
@@ -936,11 +1014,16 @@ class SigBuilder < Visitor
       #: Sig
       attr_reader :current
 
-      #: (String content, file: String) -> void
-      def initialize(content, file:)
+      # Bounds sig param comment lookup to comments inside the current `params(...)` call.
+      #: Integer?
+      attr_reader :params_start_line
+
+      #: (String content, comments_by_line: Hash[Integer, Prism::Comment], file: String) -> void
+      def initialize(content, comments_by_line:, file:)
         super
 
         @current = Sig.new #: Sig
+        @params_start_line = nil #: Integer?
       end
 
       # @override
@@ -974,7 +1057,9 @@ def visit_call_node(node)
         when "overridable"
           @current.is_overridable = true
         when "params"
+          @params_start_line = node.location.start_line
           visit(node.arguments)
+          @params_start_line = nil
         when "returns"
           return_type = sig_type_argument(node)
           @current.return_type = return_type if return_type
@@ -991,6 +1076,7 @@ def visit_call_node(node)
 
         visit(node.receiver)
         visit(node.block)
+        @current.comments.concat(comments_inside(node)) if node.message == "sig"
       end
 
       # @override
@@ -999,6 +1085,8 @@ def visit_assoc_node(node)
         @current.params << SigParam.new(
           sig_param_name(node.key),
           node_string!(node.value),
+          loc: node_loc(node),
+          comments: node_comments(node, min_line: params_start_line),
         )
       end