Add methods to ProcessedSource to get info about tokens of concrete…

… nodes

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### New features
 
 * [#88](https://github.com/rubocop-hq/rubocop-ast/pull/88): Add `RescueNode`. Add `ResbodyNode#exceptions` and `ResbodyNode#branch_index`. ([@fatkodima][])
+* [#92](https://github.com/rubocop-hq/rubocop-ast/pull/92): Add methods to `ProcessedSource` to get info about tokens of concrete nodes. ([@fatkodima][])
 * [#89](https://github.com/rubocop-hq/rubocop-ast/pull/89): Support right hand assignment for Ruby 2.8 (3.0) parser. ([@koic][])
 
 ## 0.3.0 (2020-08-01)

lib/rubocop/ast/processed_source.rb

@@ -159,6 +159,20 @@ def line_indentation(line_number)
           .length
       end
 
+      def tokens_within(source_range)
+        begin_index = first_token_index(source_range)
+        end_index = last_token_index(source_range)
+        sorted_tokens[begin_index..end_index]
+      end
+
+      def first_token_of(source_range)
+        sorted_tokens[first_token_index(source_range)]
+      end
+
+      def last_token_of(source_range)
+        sorted_tokens[last_token_index(source_range)]
+      end
+
       private
 
       def comment_index
@@ -240,6 +254,21 @@ def create_parser(ruby_version)
           end
         end
       end
+
+      def first_token_index(source_range)
+        sorted_tokens.bsearch_index { |token| token.begin_pos >= source_range.begin_pos }
+      end
+
+      def last_token_index(source_range)
+        sorted_tokens.bsearch_index { |token| token.end_pos >= source_range.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
     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,74 @@ def some_method
       expect(processed_source.following_line(brace_token)).to eq '# line 3'
     end
   end
+
+  describe '#tokens_within' do
+    context 'when no heredoc present' 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
+    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
+  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
+  end
 end