Add ProcessedSource#tokens_within, ProcessedSource#first_token_of…

… and `ProcessedSource#last_token_of`

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### New features
 
+* [#92](https://github.com/rubocop-hq/rubocop-ast/pull/92): Add `ProcessedSource#tokens_within`, `ProcessedSource#first_token_of` and `ProcessedSource#last_token_of`. ([@fatkodima][])
 * [#88](https://github.com/rubocop-hq/rubocop-ast/pull/88): Add `RescueNode`. Add `ResbodyNode#exceptions` and `ResbodyNode#branch_index`. ([@fatkodima][])
 * [#89](https://github.com/rubocop-hq/rubocop-ast/pull/89): Support right hand assignment for Ruby 2.8 (3.0) parser. ([@koic][])
 

lib/rubocop/ast/processed_source.rb

@@ -159,6 +159,20 @@ def line_indentation(line_number)
           .length
       end
 
+      def tokens_within(range_or_node)
+        begin_index = first_token_index(range_or_node)
+        end_index = last_token_index(range_or_node)
+        sorted_tokens[begin_index..end_index]
+      end
+
+      def first_token_of(range_or_node)
+        sorted_tokens[first_token_index(range_or_node)]
+      end
+
+      def last_token_of(range_or_node)
+        sorted_tokens[last_token_index(range_or_node)]
+      end
+
       private
 
       def comment_index
@@ -240,6 +254,31 @@ def create_parser(ruby_version)
           end
         end
       end
+
+      def first_token_index(range_or_node)
+        begin_pos = source_range(range_or_node).begin_pos
+        sorted_tokens.bsearch_index { |token| token.begin_pos >= begin_pos }
+      end
+
+      def last_token_index(range_or_node)
+        end_pos = source_range(range_or_node).end_pos
+        sorted_tokens.bsearch_index { |token| token.end_pos >= end_pos }
+      end
+
+      # The tokens list is always sorted by token position, except for cases when heredoc
+      # is passed as a method argument. In this case tokens are interleaved by
+      # heredoc contents' tokens.
+      def sorted_tokens
+        @sorted_tokens ||= tokens.sort_by(&:begin_pos)
+      end
+
+      def source_range(range_or_node)
+        if range_or_node.respond_to?(:source_range)
+          range_or_node.source_range
+        else
+          range_or_node
+        end
+      end
     end
   end
 end

spec/rubocop/ast/processed_source_spec.rb

@@ -10,6 +10,7 @@ def some_method
     end
     some_method
   RUBY
+  let(:ast) { processed_source.ast }
   let(:path) { 'ast/and_node_spec.rb' }
 
   shared_context 'invalid encoding source' do
@@ -292,7 +293,6 @@ def some_method
     describe '#contains_comment?' do
       subject(:commented) { processed_source.contains_comment?(range) }
 
-      let(:ast) { processed_source.ast }
       let(:array) { ast }
       let(:hash) { array.children[1] }
 
@@ -464,4 +464,89 @@ def some_method
       expect(processed_source.following_line(brace_token)).to eq '# line 3'
     end
   end
+
+  describe '#tokens_within' do
+    let(:source) { <<~RUBY }
+      foo(1, 2)
+      bar(3)
+    RUBY
+
+    it 'returns tokens for node' do
+      node = ast.children[1]
+      tokens = processed_source.tokens_within(node.source_range)
+
+      expect(tokens.map(&:text)).to eq(['bar', '(', '3', ')'])
+    end
+
+    it 'accepts Node as an argument' do
+      node = ast.children[1]
+      tokens = processed_source.tokens_within(node)
+
+      expect(tokens.map(&:text)).to eq(['bar', '(', '3', ')'])
+    end
+
+    context 'when heredoc as argument is present' do
+      let(:source) { <<~RUBY }
+        foo(1, [before], <<~DOC, [after])
+          inside heredoc.
+        DOC
+        bar(2)
+      RUBY
+
+      it 'returns tokens for node before heredoc' do
+        node = ast.children[0].arguments[1]
+        tokens = processed_source.tokens_within(node.source_range)
+
+        expect(tokens.map(&:text)).to eq(['[', 'before', ']'])
+      end
+
+      it 'returns tokens for heredoc node' do
+        node = ast.children[0].arguments[2]
+        tokens = processed_source.tokens_within(node.source_range)
+
+        expect(tokens.map(&:text)).to eq(['<<"'])
+      end
+
+      it 'returns tokens for node after heredoc' do
+        node = ast.children[0].arguments[3]
+        tokens = processed_source.tokens_within(node.source_range)
+
+        expect(tokens.map(&:text)).to eq(['[', 'after', ']'])
+      end
+    end
+  end
+
+  describe '#first_token_of' do
+    let(:source) { <<~RUBY }
+      foo(1, 2)
+      bar(3)
+    RUBY
+
+    it 'returns first token for node' do
+      node = ast.children[1]
+      expect(processed_source.first_token_of(node.source_range).text).to eq('bar')
+    end
+
+    it 'accepts Node as an argument' do
+      node = ast.children[1]
+      expect(processed_source.first_token_of(node).text).to eq('bar')
+    end
+  end
+
+  describe '#last_token_of' do
+    let(:source) { <<~RUBY }
+      foo(1, 2)
+      bar = baz
+    RUBY
+
+    it 'returns last token for node' do
+      node = ast.children[1]
+      expect(processed_source.last_token_of(node.source_range).text).to eq('baz')
+    end
+
+    it 'accepts Node as an argument' do
+      node = ast.children[1]
+      expect(processed_source.last_token_of(node).text).to eq('baz')
+    end
+  end
 end