From 708e36827655603da1ee01e55b4514a6016f6f5a Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Wed, 5 Aug 2020 12:58:54 +0300
Subject: [PATCH] Add methods to `ProcessedSource` to get info about tokens of
 concrete nodes

---
 CHANGELOG.md                              |  1 +
 lib/rubocop/ast/processed_source.rb       | 39 ++++++++++
 spec/rubocop/ast/processed_source_spec.rb | 87 ++++++++++++++++++++++-
 3 files changed, 126 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 21e02d24d..999c29a92 100644
--- a/CHANGELOG.md
+++ b/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)
diff --git a/lib/rubocop/ast/processed_source.rb b/lib/rubocop/ast/processed_source.rb
index 9540f6dd8..1fbf81c92 100644
--- a/lib/rubocop/ast/processed_source.rb
+++ b/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
diff --git a/spec/rubocop/ast/processed_source_spec.rb b/spec/rubocop/ast/processed_source_spec.rb
index 9d9bf57d9..d0e76ba8f 100644
--- a/spec/rubocop/ast/processed_source_spec.rb
+++ b/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