Start from scratch on a faster SkFlatDictionary.

src/core/SkPictureFlat.h

Index: src/core/SkPictureFlat.h
diff --git a/src/core/SkPictureFlat.h b/src/core/SkPictureFlat.h
index 0a147defe8a104da2fcddc205e79dbe86d1f0388..11afe7447e9a388b8883a94d9ce33b251103c6af 100644
--- a/src/core/SkPictureFlat.h
+++ b/src/core/SkPictureFlat.h
@@ -151,9 +151,9 @@ private:
 //                   also responsible for flattening/unflattening objects but
 //                   details of that operation are hidden in the provided procs
 // SkFlatDictionary: is an abstract templated dictionary that maintains a
-//                   searchable set of SkFlataData objects of type T.
+//                   searchable set of SkFlatData objects of type T.
 // SkFlatController: is an interface provided to SkFlatDictionary which handles
-//                   allocation and unallocation in some cases. It also holds
+//                   allocation (and unallocation in some cases). It also holds
 //                   ref count recorders and the like.
 //
 // NOTE: any class that wishes to be used in conjunction with SkFlatDictionary
@@ -176,15 +176,14 @@ public:
     virtual ~SkFlatController();
     /**
      * Provide a new block of memory for the SkFlatDictionary to use.
+     * This memory is owned by the controller and has the same lifetime unless you
+     * call unalloc(), in which case it may be freed early.

[tomhudson, 2013/07/22 16:36:02]

Nit: While we're fixing up this comment, the function isn't really 'Provid[ing]'
memory. Do we have a better verb?

[mtklein, 2013/07/22 17:43:56]

How about a nice, neutral "Return"?

      */
     virtual void* allocThrow(size_t bytes) = 0;
 
     /**
-     * Unallocate a previously allocated block, returned by allocThrow.
-     * Implementation should at least perform an unallocation if passed the last
-     * pointer returned by allocThrow. If findAndReplace() is intended to be
-     * used, unalloc should also be able to unallocate the SkFlatData that is
-     * provided.
+     * Hint that this block, which was allocated with allocThrow, is no longer needed.
+     * The implementation may choose to free this memory any time beteween now and destruction.
      */
     virtual void unalloc(void* ptr) = 0;
 
@@ -400,28 +399,30 @@ private:
         SkASSERT(SkIsAlign4(fFlatSize));
         this->data32()[fFlatSize >> 2] = value;
     }
+
+    void init(int index, int32_t size);
+    template <class T> friend class SkFlatDictionary;  // For init().
 };
 
 template <class T>
 class SkFlatDictionary {
+    static const size_t kWriteBufferGrowthBytes = 1024;
+
 public:
     SkFlatDictionary(SkFlatController* controller)
-    : fController(controller) {
-        fFlattenProc = NULL;
-        fUnflattenProc = NULL;
-        SkASSERT(controller);
-        fController->ref();
-        // set to 1 since returning a zero from find() indicates failure
-        fNextIndex = 1;
+    : fFlattenProc(NULL)
+    , fUnflattenProc(NULL)
+    , fController(SkRef(controller))
+    , fNextIndex(1)  // set to 1 since returning a zero from find() indicates failure
+    , fDataSize(0)
+    , fFlatData(this->allocFlatData(fDataSize))
+    , fWriteBuffer(kWriteBufferGrowthBytes)
+    , fWriteBufferReady(false) {
         sk_bzero(fHash, sizeof(fHash));
         // index 0 is always empty since it is used as a signal that find failed
         fIndexedData.push(NULL);
     }
 
-    virtual ~SkFlatDictionary() {
-        fController->unref();
-    }
-
     int count() const {
         SkASSERT(fIndexedData.count() == fSortedData.count()+1);
         return fSortedData.count();
@@ -532,33 +533,36 @@ public:
     }
 
     const SkFlatData* findAndReturnFlat(const T& element) {
-        SkFlatData* flat = SkFlatData::Create(fController, &element, fNextIndex, fFlattenProc);
+        const SkFlatData& temp = this->getTemporary(element, fNextIndex);
 
-        int hashIndex = ChecksumToHashIndex(flat->checksum());
+        // See if we have it in the hash?
+        const int hashIndex = ChecksumToHashIndex(temp.checksum());
         const SkFlatData* candidate = fHash[hashIndex];
-        if (candidate && !SkFlatData::Compare(*flat, *candidate)) {
-            fController->unalloc(flat);
-            return candidate;
-        }
+        if (candidate != NULL && SkFlatData::Compare(temp, *candidate) == 0) return candidate;
 
-        int index = SkTSearch<const SkFlatData,
-                              SkFlatData::Less>((const SkFlatData**) fSortedData.begin(),
-                                                fSortedData.count(), flat, sizeof(flat));
+        // See if we have it at all?
+        const int index = SkTSearch<const SkFlatData, SkFlatData::Less>(fSortedData.begin(),
+                                                                        fSortedData.count(),
+                                                                        &temp,
+                                                                        sizeof(&temp));
         if (index >= 0) {
-            fController->unalloc(flat);
+            // Found.  Update hash before we return.
             fHash[hashIndex] = fSortedData[index];
             return fSortedData[index];
         }
 
-        index = ~index;
-        *fSortedData.insert(index) = flat;
-        *fIndexedData.insert(flat->index()) = flat;
+        // We don't have it.  Add it.
+        SkFlatData* durable = this->getDurable();
+        *fSortedData.insert(~index) = durable;  // SkTSearch returned bit-not of where to insert.
+        *fIndexedData.insert(durable->index()) = durable;
+        fHash[hashIndex] = durable;
+
+        SkASSERT(durable->index() == fNextIndex);
         SkASSERT(fSortedData.count() == fNextIndex);
+        SkASSERT(fIndexedData.count() == fNextIndex+1);
         fNextIndex++;
-        flat->setSentinelInCache();
-        fHash[hashIndex] = flat;
-        SkASSERT(fIndexedData.count() == fSortedData.count()+1);
-        return flat;
+
+        return durable;
     }
 
 protected:
@@ -566,6 +570,63 @@ protected:
     void (*fUnflattenProc)(SkOrderedReadBuffer&, void*);
 
 private:
+    // Layout: [ SkFlatData header, 20 bytes ] [ data ... ] [ sentinel, 4 bytes]

[tomhudson, 2013/07/22 16:36:02]

Nit: Do we need to worry (or assert?) about the alignment of the sentinel, given
that we're treating the data as uint8?

[mtklein, 2013/07/22 17:43:56]

Oh, good call.  Using SkWriter32 guarantees that dataSize will be 4-byte
aligned, which is indeed something we're relying on to align the sentinel. 
setSentinel does this same assert itself, so we're covered already, but I think
it probably helps comprehension to document the alignment and double the assert
here.

+    uint8_t* allocFlatData(size_t dataSize) {
+        return (uint8_t*) fController->allocThrow(sizeof(SkFlatData) + dataSize + sizeof(uint32_t));
+    }
+
+    // Returns a pointer to where we should start writing flattened data (just past the header).
+    static uint8_t* dataStart(uint8_t* flatData) {
+        return flatData + sizeof(SkFlatData);
+    }
+
+    // We have to delay fWriteBuffer's initialization until its first use; fController might not
+    // be fully set up by the time we get it in the constructor.
+    void lazyWriteBufferInit() {
+        if (fWriteBufferReady) return;
+        // Without a bitmap heap, we'll flatten bitmaps into paints.  That's never what you want.
+        SkASSERT(fController->getBitmapHeap() != NULL);
+        fWriteBuffer.setBitmapHeap(fController->getBitmapHeap());
+        fWriteBuffer.setTypefaceRecorder(fController->getTypefaceSet());
+        fWriteBuffer.setNamedFactoryRecorder(fController->getNamedFactorySet());
+        fWriteBuffer.setFlags(fController->getWriteBufferFlags());
+        fWriteBufferReady = true;
+    }
+
+    // This reference is valid only until the next call to getTemporary() or getDurable().
+    const SkFlatData& getTemporary(const T& element, int index) {
+        lazyWriteBufferInit();

[tomhudson, 2013/07/22 16:36:02]

this->lazyWriteBufferInit();

[mtklein, 2013/07/22 17:43:56]

Done.

+
+        // If all the flattened bytes fit into fFlatData, we can skip the call to writeToMemory.
+        fWriteBuffer.reset(this->dataStart(fFlatData), fDataSize);
+        fFlattenProc(fWriteBuffer, &element);
+        const size_t size = fWriteBuffer.size();
+
+        if (!fWriteBuffer.wroteOnlyToStorage()) {
+            // It didn't all fit.  Copy into a bigger replacement fFlatData.
+            SkASSERT(size > fDataSize);
+            uint8_t* bigger = this->allocFlatData(size);
+            fWriteBuffer.writeToMemory(this->dataStart(bigger));
+            fController->unalloc(fFlatData);
+            fDataSize = size;
+            fFlatData = bigger;
+        }
+
+        // The data is in fFlatData now, but we need to stamp its header and trailing sentinel.
+        SkFlatData* temp = (SkFlatData*)fFlatData;
+        temp->init(index, size);
+        return *temp;
+    }
+
+    // This result is owned by fController and lives as long as it does (unless unalloc'd).
+    SkFlatData* getDurable() {
+        SkFlatData* durable = (SkFlatData*)fFlatData;
+        durable->setSentinelInCache();
+        // The next flat data we see has a good chance of being at least this big.
+        fFlatData = this->allocFlatData(fDataSize);
+        return durable;
+    }
+
     void unflatten(T* dst, const SkFlatData* element) const {
         element->unflatten(dst, fUnflattenProc,
                            fController->getBitmapHeap(),
@@ -584,14 +645,18 @@ private:
         }
     }
 
-    SkFlatController * const     fController;
-    int                          fNextIndex;
+    SkAutoTUnref<SkFlatController> fController;
+    size_t   fDataSize;
+    uint8_t* fFlatData;  // length == sizeof(SkFlatData) + fDataSize + sizeof(uint32_t)
+    SkOrderedWriteBuffer fWriteBuffer;
+    bool                 fWriteBufferReady;
 
     // SkFlatDictionary has two copies of the data one indexed by the
     // SkFlatData's index and the other sorted. The sorted data is used
     // for finding and uniquification while the indexed copy is used
     // for standard array-style lookups based on the SkFlatData's index
     // (as in 'unflatten').
+    int fNextIndex;
     SkTDArray<const SkFlatData*> fIndexedData;
     // fSortedData is sorted by checksum/size/data.
     SkTDArray<const SkFlatData*> fSortedData;
@@ -645,20 +710,19 @@ public:
         this->setTypefacePlayback(&fTypefacePlayback);
     }
 
-    ~SkChunkFlatController() {
-        fTypefaceSet->unref();
-    }
-
     virtual void* allocThrow(size_t bytes) SK_OVERRIDE {
-        return fHeap.allocThrow(bytes);
+        fLastAllocated = fHeap.allocThrow(bytes);
+        return fLastAllocated;
     }
 
     virtual void unalloc(void* ptr) SK_OVERRIDE {
-        (void) fHeap.unalloc(ptr);
+        // fHeap can only free a pointer if it was the last one allocated.  Otherwise, we'll just
+        // have to wait until fHeap is destroyed.
+        if (ptr == fLastAllocated) (void)fHeap.unalloc(ptr);
     }
 
     void setupPlaybacks() const {
-        fTypefacePlayback.reset(fTypefaceSet);
+        fTypefacePlayback.reset(fTypefaceSet.get());
     }
 
     void setBitmapStorage(SkBitmapHeap* heap) {
@@ -667,19 +731,11 @@ public:
 
 private:
     SkChunkAlloc               fHeap;
-    SkRefCntSet*               fTypefaceSet;
+    SkAutoTUnref<SkRefCntSet>  fTypefaceSet;
+    void*                      fLastAllocated;
     mutable SkTypefacePlayback fTypefacePlayback;
 };
 
-class SkBitmapDictionary : public SkFlatDictionary<SkBitmap> {
-public:
-    SkBitmapDictionary(SkFlatController* controller)
-    : SkFlatDictionary<SkBitmap>(controller) {
-        fFlattenProc = &SkFlattenObjectProc<SkBitmap>;
-        fUnflattenProc = &SkUnflattenObjectProc<SkBitmap>;
-    }
-};
-
 class SkMatrixDictionary : public SkFlatDictionary<SkMatrix> {
  public:
     SkMatrixDictionary(SkFlatController* controller)