From 63f86320d74ff8391913f8994f4203cec3f9ef42 Mon Sep 17 00:00:00 2001
From: Afront <3943720+Afront@users.noreply.github.com>
Date: Sun, 2 Aug 2020 11:47:45 +0900
Subject: [PATCH] Optimize `push_comment` method

---
 lib/rbs/ast/comment.rb   |  6 ++++++
 lib/rbs/location.rb      | 15 +++++++++++++++
 lib/rbs/parser.y         | 14 ++++++--------
 test/rbs/comment_test.rb | 23 +++++++++++++++++++++++
 4 files changed, 50 insertions(+), 8 deletions(-)
 create mode 100644 test/rbs/comment_test.rb

diff --git a/lib/rbs/ast/comment.rb b/lib/rbs/ast/comment.rb
index fdd37224d..3feacb44b 100644
--- a/lib/rbs/ast/comment.rb
+++ b/lib/rbs/ast/comment.rb
@@ -22,6 +22,12 @@ def hash
       def to_json(*a)
         { string: string, location: location }.to_json(*a)
       end
+
+      def concat(string:, location:)  
+        @string.concat string
+        @location.concat location
+        self
+      end
     end
   end
 end
diff --git a/lib/rbs/location.rb b/lib/rbs/location.rb
index 9f918a9c2..2c8bc1b63 100644
--- a/lib/rbs/location.rb
+++ b/lib/rbs/location.rb
@@ -77,6 +77,21 @@ def self.concat(*locations)
       locations.inject {|l1, l2| l1 + l2 }
     end
 
+    def concat(*others)
+      others.each { |other| self << other }
+      self
+    end
+
+    def <<(other)
+      if other
+        raise "Invalid concat: buffer=#{buffer.name}, other.buffer=#{other.buffer.name}" unless other.buffer == buffer
+        @end_pos = other.end_pos
+        @source = nil
+        @end_loc = nil
+      end
+      self
+    end
+
     def pred?(loc)
       loc.is_a?(Location) &&
         loc.name == name &&
diff --git a/lib/rbs/parser.y b/lib/rbs/parser.y
index dbb8fe079..065ecb32b 100644
--- a/lib/rbs/parser.y
+++ b/lib/rbs/parser.y
@@ -1150,15 +1150,13 @@ def leading_comment(location)
 end
 
 def push_comment(string, location)
-  new_comment = AST::Comment.new(string: string+"\n", location: location)
-
-  if (prev_comment = leading_comment(location)) && prev_comment.location.start_column == location.start_column
-    @comments.delete prev_comment.location.end_line
-    new_comment = AST::Comment.new(string: prev_comment.string + new_comment.string,
-                                   location: prev_comment.location + new_comment.location)
+  if (comment = leading_comment(location)) && comment.location.start_column == location.start_column
+    comment.concat(string: "#{string}\n", location: location)
+    @comments[comment.location.end_line] = comment
+  else
+    new_comment = AST::Comment.new(string: "#{string}\n", location: location)
+    @comments[new_comment.location.end_line] = new_comment
   end
-
-  @comments[new_comment.location.end_line] = new_comment
 end
 
 def new_token(type, value = input.matched)
diff --git a/test/rbs/comment_test.rb b/test/rbs/comment_test.rb
new file mode 100644
index 000000000..65c3b0fc3
--- /dev/null
+++ b/test/rbs/comment_test.rb
@@ -0,0 +1,23 @@
+require "test_helper"
+
+class RBS::CommentTest < Minitest::Test
+  include TestHelper
+
+  def test_concat
+    buffer = RBS::Buffer.new(name: Pathname("foo.rbs"), content: "")
+
+    comment = RBS::AST::Comment.new(
+      string: 'foo',
+      location: RBS::Location.new(buffer: buffer, start_pos: 0, end_pos: 3)
+    )
+    
+    comment.concat(
+      string: 'bar',
+      location: RBS::Location.new(buffer: buffer, start_pos: 4, end_pos: 7)
+    )
+    
+    assert_equal "foobar", comment.string
+    assert_equal 0, comment.location.start_pos
+    assert_equal 7, comment.location.end_pos
+  end
+end