Simplify and optimize comment lookup

Highlight for style/safe_navigation_spec isn't perfect though

lib/rubocop/cop/layout/extra_spacing.rb

@@ -40,6 +40,7 @@ class ExtraSpacing < Base
         def on_new_investigation
           return if processed_source.blank?
 
+          @aligned_comments = aligned_locations(processed_source.comments.map(&:loc))
           @corrected = Set.new if force_equal_sign_alignment?
 
           processed_source.tokens.each_cons(2) do |token1, token2|
@@ -49,6 +50,18 @@ def on_new_investigation
 
         private
 
+        def aligned_locations(locs)
+          return [] if locs.empty?
+
+          aligned = Set[locs.first.line, locs.last.line]
+          locs.each_cons(3) do |before, loc, after|
+            col = loc.column
+            aligned << loc.line if col == before.column || # rubocop:disable Style/MultipleComparison
+                                   col == after.column
+          end
+          aligned
+        end
+
         def check_tokens(ast, token1, token2)
           return if token2.type == :tNL
 
@@ -95,7 +108,7 @@ def extra_space_range(token1, token2)
 
         def aligned_tok?(token)
           if token.comment?
-            aligned_comments?(token)
+            @aligned_comments.include?(token.line)
           else
             aligned_with_something?(token.pos)
           end
@@ -119,26 +132,6 @@ def ignored_ranges(ast)
           end.compact
         end
 
-        def aligned_comments?(comment_token)
-          ix = processed_source.comments.index do |comment|
-            comment.loc.expression.begin_pos == comment_token.begin_pos
-          end
-          aligned_with_previous_comment?(ix) || aligned_with_next_comment?(ix)
-        end
-
-        def aligned_with_previous_comment?(index)
-          index.positive? && comment_column(index - 1) == comment_column(index)
-        end
-
-        def aligned_with_next_comment?(index)
-          index < processed_source.comments.length - 1 &&
-            comment_column(index + 1) == comment_column(index)
-        end
-
-        def comment_column(index)
-          processed_source.comments[index].loc.column
-        end
-
         def force_equal_sign_alignment?
           cop_config['ForceEqualSignAlignment']
         end

lib/rubocop/cop/lint/redundant_cop_disable_directive.rb

@@ -41,13 +41,12 @@ def initialize(config = nil, options = nil, offenses = nil)
         def on_new_investigation
           return unless offenses_to_check
 
-          comments = processed_source.comments
           cop_disabled_line_ranges = processed_source.disabled_line_ranges
 
           redundant_cops = Hash.new { |h, k| h[k] = Set.new }
 
           each_redundant_disable(cop_disabled_line_ranges,
-                                 offenses_to_check, comments) do |comment, redundant_cop|
+                                 offenses_to_check) do |comment, redundant_cop|
             redundant_cops[comment].add(redundant_cop)
           end
 
@@ -88,25 +87,25 @@ def directive_range_in_list(range, ranges)
                                        newlines: false)
         end
 
-        def each_redundant_disable(cop_disabled_line_ranges, offenses, comments,
+        def each_redundant_disable(cop_disabled_line_ranges, offenses,
                                    &block)
           disabled_ranges = cop_disabled_line_ranges[COP_NAME] || [0..0]
 
           cop_disabled_line_ranges.each do |cop, line_ranges|
             each_already_disabled(line_ranges,
-                                  disabled_ranges, comments) do |comment|
+                                  disabled_ranges) do |comment|
               yield comment, cop
             end
 
-            each_line_range(line_ranges, disabled_ranges, offenses, comments,
+            each_line_range(line_ranges, disabled_ranges, offenses,
                             cop, &block)
           end
         end
 
-        def each_line_range(line_ranges, disabled_ranges, offenses, comments,
+        def each_line_range(line_ranges, disabled_ranges, offenses,
                             cop)
           line_ranges.each_with_index do |line_range, ix|
-            comment = comments.find { |c| c.loc.line == line_range.begin }
+            comment = processed_source.comment_at_line(line_range.begin)
             next if ignore_offense?(disabled_ranges, line_range)
 
             redundant_cop = find_redundant(comment, offenses, cop, line_range,
@@ -115,7 +114,7 @@ def each_line_range(line_ranges, disabled_ranges, offenses, comments,
           end
         end
 
-        def each_already_disabled(line_ranges, disabled_ranges, comments)
+        def each_already_disabled(line_ranges, disabled_ranges)
           line_ranges.each_cons(2) do |previous_range, range|
             next if ignore_offense?(disabled_ranges, range)
             next if previous_range.end != range.begin
@@ -124,14 +123,12 @@ def each_already_disabled(line_ranges, disabled_ranges, comments)
             # the end of the previous range, it means that the cop was
             # already disabled by an earlier comment. So it's redundant
             # whether there are offenses or not.
-            redundant_comment = comments.find do |c|
-              c.loc.line == range.begin &&
-                # Comments disabling all cops don't count since it's reasonable
-                # to disable a few select cops first and then all cops further
-                # down in the code.
-                !all_disabled?(c)
-            end
-            yield redundant_comment if redundant_comment
+            comment = processed_source.comment_at_line(range.begin)
+
+            # Comments disabling all cops don't count since it's reasonable
+            # to disable a few select cops first and then all cops further
+            # down in the code.
+            yield comment if comment && !all_disabled?(comment)
           end
         end
 

lib/rubocop/cop/mixin/line_length_help.rb

@@ -12,9 +12,7 @@ def ignore_cop_directives?
 
       def directive_on_source_line?(line_index)
         source_line_number = line_index + processed_source.buffer.first_line
-        comment =
-          processed_source.comments
-                          .detect { |e| e.location.line == source_line_number }
+        comment = processed_source.comment_at_line(source_line_number)
 
         return false unless comment
 

lib/rubocop/cop/mixin/multiline_literal_brace_layout.rb

@@ -27,8 +27,7 @@ def new_line_needed_before_closing_brace?(node)
         last_element_line =
           last_element_range_with_trailing_comma(node).last_line
 
-        last_element_commented =
-          processed_source.comments.any? { |c| c.loc.line == last_element_line }
+        last_element_commented = processed_source.comment_at_line(last_element_line)
 
         last_element_commented && (node.chained? || node.argument?)
       end

lib/rubocop/cop/mixin/percent_array.rb

@@ -28,12 +28,8 @@ def message(_node)
       end
 
       def comments_in_array?(node)
-        comments = processed_source.comments
-        array_range = node.source_range.to_a
-
-        comments.any? do |comment|
-          !(comment.loc.expression.to_a & array_range).empty?
-        end
+        line_span = node.source_range.first_line...node.source_range.last_line
+        processed_source.each_comment_in_lines(line_span).any?
       end
 
       def check_percent_array(node)

lib/rubocop/cop/mixin/trailing_comma.rb

@@ -74,10 +74,8 @@ def should_have_comma?(style, node)
       end
 
       def inside_comment?(range, comma_offset)
-        processed_source.comments.any? do |comment|
-          comment_offset = comment.loc.expression.begin_pos - range.begin_pos
-          comment_offset >= 0 && comment_offset < comma_offset
-        end
+        comment = processed_source.comment_at_line(range.line)
+        comment && comment.loc.expression.begin_pos < range.begin_pos + comma_offset
       end
 
       # Returns true if the node has round/square/curly brackets.

lib/rubocop/cop/style/safe_navigation.rb

@@ -136,10 +136,10 @@ def handle_comments(corrector, node, method_call)
         end
 
         def comments(node)
-          processed_source.comments.select do |comment|
-            comment.loc.first_line > node.loc.first_line &&
-              comment.loc.last_line < node.loc.last_line
-          end
+          processed_source.each_comment_in_lines(
+            node.loc.first_line...
+            node.loc.last_line
+          ).to_a
         end
 
         def allowed_if_condition?(node)

spec/rubocop/cop/style/safe_navigation_spec.rb

@@ -674,18 +674,19 @@
 
       it 'does not lose comments within if expression' do
         expect_offense(<<~RUBY, variable: variable)
-          if %{variable}
-          ^^^^{variable} Use safe navigation (`&.`) instead [...]
+          if %{variable} # hello
+          ^^^^{variable}^^^^^^^^ Use safe navigation (`&.`) instead [...]
             # this is a comment
             # another comment
             %{variable}.bar
-          end
+          end # bye!
         RUBY
 
         expect_correction(<<~RUBY)
+          # hello
           # this is a comment
           # another comment
-          #{variable}&.bar
+          #{variable}&.bar # bye!
         RUBY
       end