Add ProcessedSource#comment_at_line. Optimize line_with_comment?

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### New features
 
 * [#70](https://github.com/rubocop-hq/rubocop-ast/pull/70): Add `NextNode` ([@marcandre][])
+* [#83](https://github.com/rubocop-hq/rubocop-ast/pull/83): Add `ProcessedSource#comment_at_line` ([@marcandre][])
 
 ### Bug fixes
 

lib/rubocop/ast/processed_source.rb

@@ -76,7 +76,7 @@ def each_comment
         comments.each { |comment| yield comment }
       end
 
-      # @deprecated Use `comments.find`
+      # @deprecated Use `comments.find` or `comment_at_line`
       def find_comment
         comments.find { |comment| yield comment }
       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)

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

tasks/compile.rake

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+desc 'Sompile [pattern] to Ruby code'
+task :compile, [:pattern]  do |_task, args|
+  require_relative '../lib/rubocop/ast'
+  pattern = args.fetch(:pattern) { $stdin.gets }
+  require 'pry'
+  puts RuboCop::AST::NodePattern.const_get(:Compiler).new(pattern).emit_method_code
+end