From 1fe5fb717d88aebda9a834d02c6f8b45f3e0c246 Mon Sep 17 00:00:00 2001
From: Michael Ilseman <milseman@apple.com>
Date: Thu, 17 May 2018 14:27:02 -0700
Subject: [PATCH] [string] Skip allocation in reserveCapacity if smol

If the requested capacity is small enough to fit in our small string
representation, don't allocate a UTF-16 buffer, instead just return
early.
---
 stdlib/public/core/StringGuts.swift  | 10 +++++++---
 test/stdlib/NewStringAppending.swift | 14 ++++++--------
 validation-test/stdlib/String.swift  | 10 +++++-----
 3 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/stdlib/public/core/StringGuts.swift b/stdlib/public/core/StringGuts.swift
index 7c7a02834db..203feceeee3 100644
--- a/stdlib/public/core/StringGuts.swift
+++ b/stdlib/public/core/StringGuts.swift
@@ -966,18 +966,22 @@ extension _StringGuts {
       }
     }
 
-    // TODO (TODO: JIRA): check if we're small and still within capacity
+    // Small strings can accomodate small capacities
+    if capacity <= _SmallUTF8String.capacity {
+      return
+    }
 
+    let selfCount = self.count
     if isASCII {
       let storage = _copyToNativeStorage(
         of: UInt8.self,
-        from: 0..<self.count,
+        from: 0..<selfCount,
         unusedCapacity: Swift.max(capacity - count, 0))
       self = _StringGuts(_large: storage)
     } else {
       let storage = _copyToNativeStorage(
         of: UTF16.CodeUnit.self,
-        from: 0..<self.count,
+        from: 0..<selfCount,
         unusedCapacity: Swift.max(capacity - count, 0))
       self = _StringGuts(_large: storage)
     }
diff --git a/test/stdlib/NewStringAppending.swift b/test/stdlib/NewStringAppending.swift
index 3a2beb17b9c..85929de0147 100644
--- a/test/stdlib/NewStringAppending.swift
+++ b/test/stdlib/NewStringAppending.swift
@@ -69,27 +69,25 @@ var s = "⓪" // start non-empty
 
 // To make this test independent of the memory allocator implementation,
 // explicitly request initial capacity.
-s.reserveCapacity(8)
+s.reserveCapacity(16)
 
-// CHECK-NEXT: String(Native(owner: @[[storage0:[x0-9a-f]+]], count: 1, capacity: 8)) = "⓪"
+// CHECK-NEXT: String(Native(owner: @[[storage0:[x0-9a-f]+]], count: 1, capacity: 16)) = "⓪"
 print("\(repr(s))")
 
-// CHECK-NEXT: String(Native(owner: @[[storage0]], count: 2, capacity: 8)) = "⓪1"
+// CHECK-NEXT: String(Native(owner: @[[storage0]], count: 2, capacity: 16)) = "⓪1"
 s += "1"
 print("\(repr(s))")
 
-// CHECK-NEXT: String(Native(owner: @[[storage0]], count: 8, capacity: 8)) = "⓪1234567"
+// CHECK-NEXT: String(Native(owner: @[[storage0]], count: 8, capacity: 16)) = "⓪1234567"
 s += "234567"
 print("\(repr(s))")
 
-// -- expect a reallocation here
-
-// CHECK-NEXT: String(Native(owner: @[[storage1:[x0-9a-f]+]], count: 9, capacity: 16)) = "⓪12345678"
+// CHECK-NEXT: String(Native(owner: @[[storage0:[x0-9a-f]+]], count: 9, capacity: 16)) = "⓪12345678"
 // CHECK-NOT: @[[storage0]],
 s += "8"
 print("\(repr(s))")
 
-// CHECK-NEXT: String(Native(owner: @[[storage1]], count: 16, capacity: 16)) = "⓪123456789012345"
+// CHECK-NEXT: String(Native(owner: @[[storage0]], count: 16, capacity: 16)) = "⓪123456789012345"
 s += "9012345"
 print("\(repr(s))")
 
diff --git a/validation-test/stdlib/String.swift b/validation-test/stdlib/String.swift
index 0b62a959dbc..fd0a8405c09 100644
--- a/validation-test/stdlib/String.swift
+++ b/validation-test/stdlib/String.swift
@@ -948,7 +948,7 @@ StringTests.test("stringGutsReserve")
       base._guts._objectIdentifier != nil &&
       isUnique
     
-    base.reserveCapacity(0)
+    base.reserveCapacity(16)
     // Now it's unique
     
     // If it was already native and unique, no reallocation
@@ -1065,11 +1065,11 @@ StringTests.test("reserveCapacity") {
   expectNotEqual(id0, s.bufferID)
   s = ""
   print("empty capacity \(s.capacity)")
-  s.reserveCapacity(oldCap + 2)
-  print("reserving \(oldCap + 2) -> \(s.capacity), width = \(s._guts.byteWidth)")
+  s.reserveCapacity(oldCap + 18)
+  print("reserving \(oldCap + 18) -> \(s.capacity), width = \(s._guts.byteWidth)")
   let id1 = s.bufferID
-  s.insert(contentsOf: repeatElement(x, count: oldCap + 2), at: s.endIndex)
-  print("extending by \(oldCap + 2) -> \(s.capacity), width = \(s._guts.byteWidth)")
+  s.insert(contentsOf: repeatElement(x, count: oldCap + 18), at: s.endIndex)
+  print("extending by \(oldCap + 18) -> \(s.capacity), width = \(s._guts.byteWidth)")
   expectEqual(id1, s.bufferID)
   s.insert(contentsOf: repeatElement(x, count: s.capacity + 100), at: s.endIndex)
   expectNotEqual(id1, s.bufferID)