Make SkTArray maintain reserve count

Bug: skia:6690
Change-Id: I01f5bb56c654f513365d6ce9f19712d9be07a08d
Reviewed-on: https://skia-review.googlesource.com/20055
Commit-Queue: Brian Salomon <bsalomon@google.com>
Reviewed-by: Yuqian Li <liyuqian@google.com>

2 files changed, 87 insertions, 19 deletions

diff --git a/include/private/SkTArray.h b/include/private/SkTArray.h
index 4b700d2d29..a6e4dedf9f 100644
--- a/include/private/SkTArray.h
+++ b/include/private/SkTArray.h
@@ -96,12 +96,15 @@ public:
     }
 
     /**
-     * Resets to count() == 0
+     * Resets to count() == 0 and resets any reserve count.
      */
-    void reset() { this->pop_back_n(fCount); }
+    void reset() {
+        this->pop_back_n(fCount);
+        fReserved = false;
+    }
 
     /**
-     * Resets to count() = n newly constructed T objects.
+     * Resets to count() = n newly constructed T objects and resets any reserve count.
      */
     void reset(int n) {
         SkASSERT(n >= 0);
@@ -115,19 +118,11 @@ public:
         for (int i = 0; i < fCount; ++i) {
             new (fItemArray + i) T;
         }
+        fReserved = false;
     }
 
     /**
-     * Ensures there is enough reserved space for n elements.
-     */
-    void reserve(int n) {
-        if (fCount < n) {
-            this->checkRealloc(n - fCount);
-        }
-    }
-
-    /**
-     * Resets to a copy of a C array.
+     * Resets to a copy of a C array and resets any reserve count.
      */
     void reset(const T* array, int count) {
         for (int i = 0; i < fCount; ++i) {
@@ -137,6 +132,22 @@ public:
         this->checkRealloc(count);
         fCount = count;
         this->copy(array);
+        fReserved = false;
+    }
+
+    /**
+     * Ensures there is enough reserved space for n additional elements. The is guaranteed at least
+     * until the array size grows above n and subsequently shrinks below n, any version of reset()
+     * is called, or reserve() is called again.
+     */
+    void reserve(int n) {
+        SkASSERT(n >= 0);
+        if (n > 0) {
+            this->checkRealloc(n);
+            fReserved = fOwnMemory;
+        } else {
+            fReserved = false;
+        }
     }
 
     void removeShuffle(int n) {
@@ -431,10 +442,12 @@ private:
             fAllocCount = 0;
             fMemArray = nullptr;
             fOwnMemory = false;
+            fReserved = false;
         } else {
             fAllocCount = SkTMax(count, SkTMax(kMinHeapAllocCount, reserveCount));
             fMemArray = sk_malloc_throw(fAllocCount * sizeof(T));
             fOwnMemory = true;
+            fReserved = reserveCount > 0;
         }
     }
 
@@ -444,6 +457,7 @@ private:
         SkASSERT(preallocStorage);
         fCount = count;
         fMemArray = nullptr;
+        fReserved = false;
         if (count > preallocCount) {
             fAllocCount = SkTMax(count, kMinHeapAllocCount);
             fMemArray = sk_malloc_throw(fAllocCount * sizeof(T));
@@ -505,10 +519,10 @@ private:
         int newCount = fCount + delta;
 
         // We allow fAllocCount to be in the range [newCount, 3*newCount]. We also never shrink
-        // when we're currently using preallocated memory or would allocate less than
-        // kMinHeapAllocCount.
+        // when we're currently using preallocated memory, would allocate less than
+        // kMinHeapAllocCount, or a reserve count was specified that has yet to be exceeded.
         bool mustGrow = newCount > fAllocCount;
-        bool shouldShrink = fAllocCount > 3 * newCount && fOwnMemory;
+        bool shouldShrink = fAllocCount > 3 * newCount && fOwnMemory && !fReserved;
         if (!mustGrow && !shouldShrink) {
             return;
         }
@@ -531,15 +545,17 @@ private:
         }
         fMemArray = newMemArray;
         fOwnMemory = true;
+        fReserved = false;
     }
 
-    int fCount;
-    int fAllocCount;
-    bool fOwnMemory;
     union {
         T*       fItemArray;
         void*    fMemArray;
     };
+    int fCount;
+    int fAllocCount;
+    bool fOwnMemory : 1;
+    bool fReserved : 1;
 };
 
 template<typename T, bool MEM_MOVE> constexpr int SkTArray<T, MEM_MOVE>::kMinHeapAllocCount;
diff --git a/tests/TArrayTest.cpp b/tests/TArrayTest.cpp
index 0e18e9f96d..1b376e4d4b 100644
--- a/tests/TArrayTest.cpp
+++ b/tests/TArrayTest.cpp
@@ -5,6 +5,7 @@
  * found in the LICENSE file.
  */
 
+#include "SkRandom.h"
 #include "SkRefCnt.h"
 #include "SkTArray.h"
 #include "Test.h"
@@ -294,6 +295,52 @@ static void test_self_assignment(skiatest::Reporter* reporter) {
     REPORTER_ASSERT(reporter, a[0] == 1);
 }
 
+template <typename Array> static void test_array_reserve(skiatest::Reporter* reporter,
+                                                         Array* array, int reserveCount) {
+    SkRandom random;
+    REPORTER_ASSERT(reporter, array->allocCntForTest() >= reserveCount);
+    array->push_back();
+    REPORTER_ASSERT(reporter, array->allocCntForTest() >= reserveCount);
+    array->pop_back();
+    REPORTER_ASSERT(reporter, array->allocCntForTest() >= reserveCount);
+    while (array->count() < reserveCount) {
+        // Two steps forward, one step back
+        if (random.nextULessThan(3) < 2) {
+            array->push_back();
+        } else if (array->count() > 0) {
+            array->pop_back();
+        }
+        REPORTER_ASSERT(reporter, array->allocCntForTest() >= reserveCount);
+    }
+}
+
+template<typename Array> static void test_reserve(skiatest::Reporter* reporter) {
+    // Test that our allocated space stays >= to the reserve count until the array is filled to
+    // the reserve count
+    for (int reserveCount : {1, 2, 10, 100}) {
+        // Test setting reserve in constructor.
+        Array array1(reserveCount);
+        test_array_reserve(reporter, &array1, reserveCount);
+
+        // Test setting reserve after constructor.
+        Array array2;
+        array2.reserve(reserveCount);
+        test_array_reserve(reporter, &array2, reserveCount);
+
+        // Test increasing reserve after constructor.
+        Array array3(reserveCount/2);
+        array3.reserve(reserveCount);
+        test_array_reserve(reporter, &array3, reserveCount);
+
+        // Test setting reserve on non-empty array.
+        Array array4;
+        array4.push_back_n(reserveCount);
+        array4.reserve(reserveCount);
+        array4.pop_back_n(reserveCount);
+        test_array_reserve(reporter, &array4, 2 * reserveCount);
+    }
+}
+
 DEF_TEST(TArray, reporter) {
     TestTSet_basic<true>(reporter);
     TestTSet_basic<false>(reporter);
@@ -311,4 +358,9 @@ DEF_TEST(TArray, reporter) {
     test_unnecessary_alloc(reporter);
 
     test_self_assignment(reporter);
+
+    test_reserve<SkTArray<int>>(reporter);
+    test_reserve<SkSTArray<1, int>>(reporter);
+    test_reserve<SkSTArray<2, int>>(reporter);
+    test_reserve<SkSTArray<16, int>>(reporter);
 }