Fix StrNode#last_line and StrNode#line_count for heredocs

Author: fatkodima

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## master (unreleased)
 
+### Bug fixes
+
+* [#40](https://github.com/rubocop-hq/rubocop-ast/pull/40): Fix `StrNode#last_line` and `StrNode#line_count` for heredocs. ([@fatkodima][])
+
 ## 1.3.0 (2020-11-30)
 
 ### Changes

lib/rubocop/ast/node/str_node.rb

@@ -11,6 +11,18 @@ class StrNode < Node
       def heredoc?
         loc.is_a?(Parser::Source::Map::Heredoc)
       end
+
+      def last_line
+        if heredoc?
+          loc.heredoc_end.line
+        else
+          super
+        end
+      end
+
+      def line_count
+        last_line - first_line + 1
+      end
     end
   end
 end

spec/rubocop/ast/str_node_spec.rb

@@ -56,4 +56,73 @@
       it { is_expected.to be_heredoc }
     end
   end
+
+  describe '#first_line' do
+    context 'with a normal string' do
+      let(:source) { "'foo'" }
+
+      it { expect(str_node.first_line).to eq(1) }
+    end
+
+    context 'with a heredoc' do
+      let(:source) do
+        <<~RUBY
+          <<-CODE
+            foo
+            bar
+          CODE
+        RUBY
+      end
+
+      it { expect(str_node.first_line).to eq(1) }
+    end
+  end
+
+  describe '#last_line' do
+    context 'with a normal string' do
+      let(:source) do
+        ['"foo"\\',
+         '"bar"'].join("\n")
+      end
+
+      it { expect(str_node.last_line).to eq(2) }
+    end
+
+    context 'with a heredoc' do
+      let(:source) do
+        <<~RUBY
+          <<-CODE
+            foo
+            bar
+          CODE
+        RUBY
+      end
+
+      it { expect(str_node.last_line).to eq(4) }
+    end
+  end
+
+  describe '#line_count' do
+    context 'with a normal string' do
+      let(:source) do
+        ['"foo"\\',
+         '"bar"'].join("\n")
+      end
+
+      it { expect(str_node.line_count).to eq(2) }
+    end
+
+    context 'with a heredoc' do
+      let(:source) do
+        <<~RUBY
+          <<-CODE
+            foo
+            bar
+          CODE
+        RUBY
+      end
+
+      it { expect(str_node.line_count).to eq(4) }
+    end
+  end
 end