Optimize String::Builder#write_byte

asterite · asterite · commit 165593521494 · 2017-08-19T13:40:28.000-03:00
diff --git a/spec/std/string_spec.cr b/spec/std/string_spec.cr
@@ -2163,6 +2163,17 @@ describe "String" do
     "fo\u{0000}".compare("FO", case_insensitive: true).should eq(1)
   end
 
+  it "builds with write_byte" do
+    string = String.build do |io|
+      255_u8.times do |byte|
+        io.write_byte(byte)
+      end
+    end
+    255.times do |i|
+      string.byte_at(i).should eq(i)
+    end
+  end
+
   it "raises if String.build negative capacity" do
     expect_raises(ArgumentError, "Negative capacity") do
       String.build(-1) { }
diff --git a/src/string/builder.cr b/src/string/builder.cr
@@ -53,6 +53,19 @@ class String::Builder
     nil
   end
 
+  def write_byte(byte : UInt8)
+    new_bytesize = real_bytesize + 1
+    if new_bytesize > @capacity
+      resize_to_capacity(Math.pw2ceil(new_bytesize))
+    end
+
+    @buffer[real_bytesize] = byte
+
+    @bytesize += 1
+
+    nil
+  end
+
   def buffer
     @buffer + String::HEADER_SIZE
   end
