Repository: ruby/yarp
Pull request #1428 (2023-09-08 18:35:44)

Reorganize parse result mutations
* Moves `MarkNewlinesVisitor` into `lib/yarp/parse_result/newlines.rb`
* Moves `CommentsAttacher` into `lib/yarp/parse_result/comments.rb`
* Template out a `comment_targets` method instead of relying on `deconstruct_keys` to find the comment targets
* Rename `mark_newlines` to `mark_newlines!` to indicate that it mutates `self`


Patch:
diff --git a/bin/parse b/bin/parse
index 05403e12..54824563 100755
--- a/bin/parse
+++ b/bin/parse
@@ -15,7 +15,7 @@ else
   result = YARP.parse_file(ARGV[0] || "test.rb")
 end
 
-result.mark_newlines if ENV["MARK_NEWLINES"]
+result.mark_newlines! if ENV["MARK_NEWLINES"]
 
 value = result.value
 value = value.accept(YARP::DesugarVisitor.new) if ENV["DESUGAR"]
diff --git a/lib/yarp.rb b/lib/yarp.rb
index 32d4a1b5..f878e719 100644
--- a/lib/yarp.rb
+++ b/lib/yarp.rb
@@ -248,193 +248,6 @@ module YARP
     def failure?
       !success?
     end
-
-    # CommentAttacher is a utility class to attach comments to locations in the AST
-    class CommentAttacher
-      attr_reader :parse_result
-
-      def initialize(parse_result)
-        @parse_result = parse_result
-      end
-
-      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)
-            else
-              # If a comment exists on its own line, prefer a leading comment.
-              following || preceding || enclosing || NodeTarget.new(parse_result.value)
-            end
-
-          target << comment
-        end
-      end
-
-      # A target for attaching comments that is based on a specific node
-      class NodeTarget
-        attr_reader :node
-
-        def initialize(node)
-          @node = node
-        end
-
-        def start_offset
-          node.location.start_offset
-        end
-
-        def end_offset
-          node.location.end_offset
-        end
-
-        def encloses?(comment)
-          start_offset <= comment.location.start_offset && comment.location.end_offset <= end_offset
-        end
-
-        def <<(comment)
-          node.location.comments << comment
-        end
-      end
-
-      # A target for attaching comments that is based on a location, which could be a part of a node. For example, the
-      # `end` token of a ClassNode
-      class LocationTarget
-        attr_reader :location
-
-        def initialize(location)
-          @location = location
-        end
-
-        def start_offset
-          location.start_offset
-        end
-
-        def end_offset
-          location.end_offset
-        end
-
-        def encloses?(comment)
-          false
-        end
-
-        def <<(comment)
-          location.comments << comment
-        end
-      end
-
-      private
-
-      # Responsible for finding the nearest targets to the given comment within the context of the given encapsulating
-      # node.
-      def nearest_targets(node, comment)
-        comment_start = comment.location.start_offset
-        comment_end = comment.location.end_offset
-
-        targets = []
-        node.deconstruct_keys(nil).each do |key, value|
-          next if key == :location
-
-          case value
-          when StatementsNode
-            targets.concat(value.body.map { |node| NodeTarget.new(node) })
-          when Node
-            targets << NodeTarget.new(value)
-          when Location
-            targets << LocationTarget.new(value)
-          when Array
-            targets.concat(value.map { |node| NodeTarget.new(node) }) if value.first.is_a?(Node)
-          end
-        end
-
-        targets.sort_by!(&:start_offset)
-        preceding = nil
-        following = nil
-
-        left = 0
-        right = targets.length
-
-        # This is a custom binary search that finds the nearest nodes to the given comment. When it finds a node that
-        # completely encapsulates the comment, it recursed downward into the tree.
-        while left < right
-          middle = (left + right) / 2
-          target = targets[middle]
-
-          target_start = target.start_offset
-          target_end = target.end_offset
-
-          if target.encloses?(comment)
-            # The comment is completely contained by this target. Abandon the binary search at this level.
-            return nearest_targets(target.node, comment)
-          end
-
-          if target_end <= comment_start
-            # This target falls completely before the comment. Because we will never consider this target or any targets
-            # before it again, this target must be the closest preceding target we have encountered so far.
-            preceding = target
-            left = middle + 1
-            next
-          end
-
-          if comment_end <= target_start
-            # This target falls completely after the comment. Because we will never consider this target or any targets
-            # after it again, this target must be the closest following target we have encountered so far.
-            following = target
-            right = middle
-            next
-          end
-
-          # This should only happen if there is a bug in this parser.
-          raise "Comment location overlaps with target location"
-        end
-
-        [preceding, NodeTarget.new(node), following]
-      end
-    end
-
-    # Attach the list of comments to their respective locations in the AST
-    def attach_comments!
-      CommentAttacher.new(self).attach!
-    end
-
-    # Keep in sync with Java MarkNewlinesVisitor
-    class MarkNewlinesVisitor < YARP::Visitor
-      def initialize(newline_marked)
-        @newline_marked = newline_marked
-      end
-
-      def visit_block_node(node)
-        old_newline_marked = @newline_marked
-        @newline_marked = Array.new(old_newline_marked.size, false)
-        begin
-          super(node)
-        ensure
-          @newline_marked = old_newline_marked
-        end
-      end
-      alias_method :visit_lambda_node, :visit_block_node
-
-      def visit_if_node(node)
-        node.set_newline_flag(@newline_marked)
-        super(node)
-      end
-      alias_method :visit_unless_node, :visit_if_node
-
-      def visit_statements_node(node)
-        node.body.each do |child|
-          child.set_newline_flag(@newline_marked)
-        end
-        super(node)
-      end
-    end
-    private_constant :MarkNewlinesVisitor
-
-    def mark_newlines
-      newline_marked = Array.new(1 + @source.offsets.size, false)
-      visitor = MarkNewlinesVisitor.new(newline_marked)
-      value.accept(visitor)
-      value
-    end
   end
 
   # This represents a token from the Ruby source.
@@ -778,6 +591,9 @@ require_relative "yarp/serialize"
 require_relative "yarp/pack"
 require_relative "yarp/pattern"
 
+require_relative "yarp/parse_result/comments"
+require_relative "yarp/parse_result/newlines"
+
 if RUBY_ENGINE == "ruby" and !ENV["YARP_FFI_BACKEND"]
   require "yarp/yarp"
 else
diff --git a/lib/yarp/parse_result/comments.rb b/lib/yarp/parse_result/comments.rb
new file mode 100644
index 00000000..88240609
--- /dev/null
+++ b/lib/yarp/parse_result/comments.rb
@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+module YARP
+  class ParseResult
+    # 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.
+    #
+    # It does this by first finding the nearest locations to each comment.
+    # Locations can either come from nodes directly or from location fields on
+    # nodes. For example, a `ClassNode` has an overall location encompassing the
+    # entire class, but it also has a location for the `class` keyword.
+    #
+    # Once the nearest locations are found, it determines which one to attach
+    # to. If it's a trailing comment (a comment on the same line as other source
+    # code), it will favor attaching to the nearest location that occurs before
+    # the comment. Otherwise it will favor attaching to the nearest location
+    # that is after the comment.
+    class Comments
+      # A target for attaching comments that is based on a specific node's
+      # location.
+      class NodeTarget
+        attr_reader :node
+
+        def initialize(node)
+          @node = node
+        end
+
+        def start_offset
+          node.location.start_offset
+        end
+
+        def end_offset
+          node.location.end_offset
+        end
+
+        def encloses?(comment)
+          start_offset <= comment.location.start_offset &&
+            comment.location.end_offset <= end_offset
+        end
+
+        def <<(comment)
+          node.location.comments << 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
+        attr_reader :location
+
+        def initialize(location)
+          @location = location
+        end
+
+        def start_offset
+          location.start_offset
+        end
+
+        def end_offset
+          location.end_offset
+        end
+
+        def encloses?(comment)
+          false
+        end
+
+        def <<(comment)
+          location.comments << comment
+        end
+      end
+
+      attr_reader :parse_result
+
+      def initialize(parse_result)
+        @parse_result = parse_result
+      end
+
+      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)
+            else
+              # If a comment exists on its own line, prefer a leading comment.
+              following || preceding || enclosing || NodeTarget.new(parse_result.value)
+            end
+
+          target << comment
+        end
+      end
+
+      private
+
+      # Responsible for finding the nearest targets to the given comment within
+      # the context of the given encapsulating node.
+      def nearest_targets(node, comment)
+        comment_start = comment.location.start_offset
+        comment_end = comment.location.end_offset
+
+        targets = []
+        node.comment_targets.map do |value|
+          case value
+          when StatementsNode
+            targets.concat(value.body.map { |node| NodeTarget.new(node) })
+          when Node
+            targets << NodeTarget.new(value)
+          when Location
+            targets << LocationTarget.new(value)
+          end
+        end
+
+        targets.sort_by!(&:start_offset)
+        preceding = nil
+        following = nil
+
+        left = 0
+        right = targets.length
+
+        # This is a custom binary search that finds the nearest nodes to the
+        # given comment. When it finds a node that completely encapsulates the
+        # comment, it recurses downward into the tree.
+        while left < right
+          middle = (left + right) / 2
+          target = targets[middle]
+
+          target_start = target.start_offset
+          target_end = target.end_offset
+
+          if target.encloses?(comment)
+            # The comment is completely contained by this target. Abandon the
+            # binary search at this level.
+            return nearest_targets(target.node, comment)
+          end
+
+          if target_end <= comment_start
+            # This target falls completely before the comment. Because we will
+            # never consider this target or any targets before it again, this
+            # target must be the closest preceding target we have encountered so
+            # far.
+            preceding = target
+            left = middle + 1
+            next
+          end
+
+          if comment_end <= target_start
+            # This target falls completely after the comment. Because we will
+            # never consider this target or any targets after it again, this
+            # target must be the closest following target we have encountered so
+            # far.
+            following = target
+            right = middle
+            next
+          end
+
+          # This should only happen if there is a bug in this parser.
+          raise "Comment location overlaps with a target location"
+        end
+
+        [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/yarp/parse_result/newlines.rb b/lib/yarp/parse_result/newlines.rb
new file mode 100644
index 00000000..d16600af
--- /dev/null
+++ b/lib/yarp/parse_result/newlines.rb
@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module YARP
+  class ParseResult
+    # 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:
+    #
+    # * if statements
+    # * unless statements
+    # * nodes that are children of statements lists
+    #
+    # In order to keep track of the newlines, we have a list of offsets that
+    # come back from the parser. We assign these offsets to the first nodes that
+    # we find in the tree that are on those lines.
+    #
+    # 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.
+    class Newlines < Visitor
+      def initialize(newline_marked)
+        @newline_marked = newline_marked
+      end
+
+      def visit_block_node(node)
+        old_newline_marked = @newline_marked
+        @newline_marked = Array.new(old_newline_marked.size, false)
+
+        begin
+          super(node)
+        ensure
+          @newline_marked = old_newline_marked
+        end
+      end
+
+      alias_method :visit_lambda_node, :visit_block_node
+
+      def visit_if_node(node)
+        node.set_newline_flag(@newline_marked)
+        super(node)
+      end
+
+      alias_method :visit_unless_node, :visit_if_node
+
+      def visit_statements_node(node)
+        node.body.each do |child|
+          child.set_newline_flag(@newline_marked)
+        end
+        super(node)
+      end
+    end
+
+    private_constant :Newlines
+
+    # 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)))
+    end
+  end
+end
diff --git a/templates/lib/yarp/node.rb.erb b/templates/lib/yarp/node.rb.erb
index 0d404063..9cc44959 100644
--- a/templates/lib/yarp/node.rb.erb
+++ b/templates/lib/yarp/node.rb.erb
@@ -48,6 +48,16 @@ module YARP
       }.compact.join(", ") %>]
     end
 
+    # def comment_targets: () -> Array[Node | Location]
+    def comment_targets
+      [<%= node.fields.map { |field|
+        case field
+        when YARP::NodeField, YARP::LocationField then field.name
+        when YARP::OptionalNodeField, YARP::NodeListField, YARP::OptionalLocationField then "*#{field.name}"
+        end
+      }.compact.join(", ") %>]
+    end
+
     # def copy: (**params) -> <%= node.name %>
     def copy(**params)
       <%= node.name %>.new(
diff --git a/yarp.gemspec b/yarp.gemspec
index 33b47d67..0cd1cf2f 100644
--- a/yarp.gemspec
+++ b/yarp.gemspec
@@ -68,6 +68,8 @@ Gem::Specification.new do |spec|
     "lib/yarp/pattern.rb",
     "lib/yarp/ripper_compat.rb",
     "lib/yarp/serialize.rb",
+    "lib/yarp/parse_result/comments.rb",
+    "lib/yarp/parse_result/newlines.rb",
     "src/diagnostic.c",
     "src/enc/yp_big5.c",
     "src/enc/yp_euc_jp.c",


Test patch:
diff --git a/test/yarp/newline_test.rb b/test/yarp/newline_test.rb
index 5a85f856..ba82b5ef 100644
--- a/test/yarp/newline_test.rb
+++ b/test/yarp/newline_test.rb
@@ -37,7 +37,7 @@ module YARP
       result = YARP.parse_file(filepath)
       assert_empty result.errors
 
-      result.mark_newlines
+      result.mark_newlines!
       visitor = NewlineVisitor.new(result.source)
 
       result.value.accept(visitor)