+ Add associate_by_identity as an alternate to associate

lib/parser/source/comment.rb

@@ -48,6 +48,19 @@ def self.associate_locations(ast, comments)
         associator.associate_locations
       end
 
+      ##
+      # Associate `comments` with `ast` nodes using identity.
+      #
+      # @param [Parser::AST::Node] ast
+      # @param [Array<Comment>]    comments
+      # @return [Hash<Parser::Source::Node, Array<Comment>>]
+      # @see Parser::Source::Comment::Associator#associate_by_identity
+      #
+      def self.associate_by_identity(ast, comments)
+        associator = Associator.new(ast, comments)
+        associator.associate_by_identity
+      end
+
       ##
       # @param [Parser::Source::Range] range
       #

lib/parser/source/comment/associator.rb

@@ -84,12 +84,24 @@ def initialize(ast, comments)
       #
       # Note that {associate} produces unexpected result for nodes which are
       # equal but have distinct locations; comments for these nodes are merged.
+      # You may prefer using {associate_by_identity} or {associate_locations}.
       #
       # @return [Hash<Parser::AST::Node, Array<Parser::Source::Comment>>]
       # @deprecated Use {associate_locations}.
       #
       def associate
-        @map_using_locations = false
+        @map_using = :eql
+        do_associate
+      end
+
+      ##
+      # Same as {associate}, but compares by identity, thus producing an unambiguous
+      # result even in presence of equal nodes.
+      #
+      # @return [Hash<Parser::Source::Node, Array<Parser::Source::Comment>>]
+      #
+      def associate_locations
+        @map_using = :location
         do_associate
       end
 
@@ -100,8 +112,8 @@ def associate
       #
       # @return [Hash<Parser::Source::Map, Array<Parser::Source::Comment>>]
       #
-      def associate_locations
-        @map_using_locations = true
+      def associate_by_identity
+        @map_using = :identity
         do_associate
       end
 
@@ -122,6 +134,7 @@ def children_in_source_order(node)
 
       def do_associate
         @mapping     = Hash.new { |h, k| h[k] = [] }
+        @mapping.compare_by_identity if @map_using == :identity
         @comment_num = -1
         advance_comment
 
@@ -191,7 +204,7 @@ def current_comment_decorates?(node)
       end
 
       def associate_and_advance_comment(node)
-        key = @map_using_locations ? node.location : node
+        key = @map_using == :location ? node.location : node
         @mapping[key] << @current_comment
         advance_comment
       end

test/test_source_comment_associator.rb

@@ -27,6 +27,13 @@ def associate_locations(code)
     [ ast, associations ]
   end
 
+  def associate_by_identity(code)
+    ast, comments = parse_with_comments(code)
+    associations  = Parser::Source::Comment.associate_by_identity(ast, comments)
+
+    [ ast, associations ]
+  end
+
   def test_associate
     ast, associations = associate(<<-END)
 #!/usr/bin/env ruby
@@ -83,9 +90,8 @@ def bar
     ], associations[one_node].map(&:text)
   end
 
-  # The bug below is fixed by using associate_locations
-  def test_associate_dupe_statement
-    ast, associations = associate(<<-END)
+  def setup_dupe_statement(method = :associate)
+    @ast, @associations = send(method, <<-END)
 class Foo
   def bar
     f1 # comment on 1st call to f1
@@ -95,17 +101,30 @@ def bar
 end
     END
 
-    _klass_node        = ast
-    method_node        = ast.children[2]
-    body               = method_node.children[2]
-    f1_1_node          = body.children[0]
-    f1_2_node          = body.children[2]
+    _klass_node        = @ast
+    @method_node        = @ast.children[2]
+    @body               = @method_node.children[2]
+    @f1_1_node          = @body.children[0]
+    @f1_2_node          = @body.children[2]
+  end
 
-    assert_equal 1, associations.size
+  # The bug below is fixed by using associate_locations
+  def test_associate_dupe_statement
+    setup_dupe_statement
+    assert_equal 1, @associations.size
     assert_equal ['# comment on 1st call to f1', '# comment on 2nd call to f1'],
-                 associations[f1_1_node].map(&:text)
+                 @associations[@f1_1_node].map(&:text)
     assert_equal ['# comment on 1st call to f1', '# comment on 2nd call to f1'],
-                 associations[f1_2_node].map(&:text)
+                 @associations[@f1_2_node].map(&:text)
+  end
+
+  def test_associate_by_identity_dupe_statement
+    setup_dupe_statement(:associate_by_identity)
+    assert_equal 2, @associations.size
+    assert_equal ['# comment on 1st call to f1'],
+                 @associations[@f1_1_node].map(&:text)
+    assert_equal ['# comment on 2nd call to f1'],
+                 @associations[@f1_2_node].map(&:text)
   end
 
   def test_associate_locations