Move parser comment handling to shared visitor

Author: KaanOzkan

lib/rbi/parser.rb

@@ -112,12 +112,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: nil)
         super()
 
         @source = source
         @file = file
+        @comments_by_line = comments_by_line || {} #: Hash[Integer, Prism::Comment]
       end
 
       private
@@ -159,6 +160,75 @@ def self?(node)
       def t_sig_without_runtime?(node)
         !!(node.is_a?(Prism::ConstantPathNode) && node_string(node) =~ /(::)?T::Sig::WithoutRuntime/)
       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
     end
 
     class TreeBuilder < Visitor
@@ -170,9 +240,12 @@ class TreeBuilder < Visitor
 
       #: (String source, comments: Array[Prism::Comment], file: String) -> void
       def initialize(source, comments:, file:)
-        super(source, file: file)
+        super(
+          source,
+          comments_by_line: comments.to_h { |c| [c.location.start_line, c] },
+          file: file,
+        )
 
-        @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]
@@ -600,75 +673,6 @@ def detach_comments_from_sigs(sigs)
         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]