From 641f0ebe157cc5006302f9b505891f2374dca903 Mon Sep 17 00:00:00 2001
From: Marc-Andre Lafortune <github@marc-andre.ca>
Date: Tue, 21 Jul 2020 12:30:16 -0400
Subject: [PATCH] Add `ProcessedSource#comment_at_line`. Optimize
 `line_with_comment?`

---
 CHANGELOG.md                              |  1 +
 lib/rubocop/ast/processed_source.rb       | 15 +++++++++++----
 spec/rubocop/ast/processed_source_spec.rb | 11 +++++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 919cf8f4c..40a35ada5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@
 * [#70](https://github.com/rubocop-hq/rubocop-ast/pull/70): Add `NextNode` ([@marcandre][])
 * [#85](https://github.com/rubocop-hq/rubocop-ast/pull/85): Add `IntNode#value` and `FloatNode#value`. ([@fatkodima][])
 * [#82](https://github.com/rubocop-hq/rubocop-ast/pull/82): `NodePattern`: Allow comments ([@marcandre][])
+* [#83](https://github.com/rubocop-hq/rubocop-ast/pull/83): Add `ProcessedSource#comment_at_line` ([@marcandre][])
 
 ### Bug fixes
 
diff --git a/lib/rubocop/ast/processed_source.rb b/lib/rubocop/ast/processed_source.rb
index 928a07da8..c7c368412 100644
--- a/lib/rubocop/ast/processed_source.rb
+++ b/lib/rubocop/ast/processed_source.rb
@@ -76,7 +76,7 @@ def each_comment(&block)
         comments.each(&block)
       end
 
-      # @deprecated Use `comments.find`
+      # @deprecated Use `comments.find` or `comment_at_line`
       def find_comment(&block)
         comments.find(&block)
       end
@@ -99,9 +99,14 @@ def blank?
         ast.nil?
       end
 
+      # @return [Comment, nil] the comment at that line, if any.
+      def comment_at_line(line)
+        comment_index[line]
+      end
+
       # @return [Boolean] if the given line number has a comment.
       def line_with_comment?(line)
-        comment_lines.include?(line)
+        comment_index.include?(line)
       end
 
       # @return [Boolean] if any of the lines in the given `source_range` has a comment.
@@ -144,8 +149,10 @@ def line_indentation(line_number)
 
       private
 
-      def comment_lines
-        @comment_lines ||= comments.map { |c| c.location.line }
+      def comment_index
+        @comment_index ||= {}.tap do |hash|
+          comments.each { |c| hash[c.location.line] = c }
+        end
       end
 
       def parse(source, ruby_version)
diff --git a/spec/rubocop/ast/processed_source_spec.rb b/spec/rubocop/ast/processed_source_spec.rb
index 4127b2564..cefb77961 100644
--- a/spec/rubocop/ast/processed_source_spec.rb
+++ b/spec/rubocop/ast/processed_source_spec.rb
@@ -257,6 +257,17 @@ def some_method
       end
     end
 
+    describe '#comment_at_line' do
+      it 'returns the comment at the given line number' do
+        expect(processed_source.comment_at_line(1).text).to eq '# comment one'
+        expect(processed_source.comment_at_line(4).text).to eq '# comment two'
+      end
+
+      it 'returns nil if line has no comment' do
+        expect(processed_source.comment_at_line(3)).to be nil
+      end
+    end
+
     describe '#line_with_comment?' do
       it 'returns true for lines with comments' do
         expect(processed_source.line_with_comment?(1)).to be true