Use SkFlatDictionary in SkGPipe to take advantage of its new features.

Add a controller class to perform the allocation/unallocation for the dictionary and to provide an entry to be replaced, if replacements are allowed.

TODO:
Use LRU caching in my custom controller so replacements will be done less often.
More refactoring on SkFlatDictionary so picture recording's use of the dictionary does not require going through the path to replace.

Review URL: https://codereview.appspot.com/6345102

git-svn-id: http://skia.googlecode.com/svn/trunk@4639 2bbb7eff-a529-9590-31e7-b0007b416f81

6 files changed, 337 insertions, 207 deletions

diff --git a/src/core/SkPictureFlat.cpp b/src/core/SkPictureFlat.cpp
index 34d05fce66..763ed60c5c 100644
--- a/src/core/SkPictureFlat.cpp
+++ b/src/core/SkPictureFlat.cpp
@@ -60,11 +60,10 @@ SkRefCnt* SkRefCntPlayback::set(int index, SkRefCnt* obj) {
 
 ///////////////////////////////////////////////////////////////////////////////
 
-SkFlatData* SkFlatData::Create(SkChunkAlloc* heap, const void* obj,
+SkFlatData* SkFlatData::Create(SkFlatController* controller, const void* obj,
         int index, void (*flattenProc)(SkOrderedWriteBuffer&, const void*),
         SkRefCntSet* refCntRecorder, SkRefCntSet* faceRecorder,
-        uint32_t writeBufferflags) {
-
+        uint32_t writeBufferflags, SkFactorySet* fset) {
     // a buffer of 256 bytes should be sufficient for most paints, regions,
     // and matrices.
     intptr_t storage[256];
@@ -76,7 +75,8 @@ SkFlatData* SkFlatData::Create(SkChunkAlloc* heap, const void* obj,
         buffer.setTypefaceRecorder(faceRecorder);
     }
     buffer.setFlags(writeBufferflags);
-
+    buffer.setFactoryRecorder(fset);
+    
     flattenProc(buffer, obj);
     uint32_t size = buffer.size();
     SkASSERT(SkIsAlign4(size));
@@ -88,7 +88,7 @@ SkFlatData* SkFlatData::Create(SkChunkAlloc* heap, const void* obj,
      *  3. 4-byte sentinel
      */
     size_t allocSize = sizeof(SkFlatData) + size + sizeof(uint32_t);
-    SkFlatData* result = (SkFlatData*) heap->allocThrow(allocSize);
+    SkFlatData* result = (SkFlatData*) controller->allocThrow(allocSize);
 
     result->fIndex = index;
     result->fFlatSize = size;
diff --git a/src/core/SkPictureFlat.h b/src/core/SkPictureFlat.h
index b8f0b083f8..c746c01251 100644
--- a/src/core/SkPictureFlat.h
+++ b/src/core/SkPictureFlat.h
@@ -142,17 +142,40 @@ private:
 //                   which is agnostic to the type of data is is indexing. It is
 //                   also responsible for flattening/unflattening objects but
 //                   details of that operation are hidden in the provided procs
-// SkFlatDictionary: is a abstract templated dictionary that maintains a
+// SkFlatDictionary: is an abstract templated dictionary that maintains a
 //                   searchable set of SkFlataData objects of type T.
+// SkFlatController: is an interface provided to SkFlatDictionary which handles
+//                   allocation and unallocation in some cases
 //
 // NOTE: any class that wishes to be used in conjunction with SkFlatDictionary
 // must subclass the dictionary and provide the necessary flattening procs.
 // The end of this header contains dictionary subclasses for some common classes
-// like SkBitmap, SkMatrix, SkPaint, and SkRegion.
+// like SkBitmap, SkMatrix, SkPaint, and SkRegion. SkFlatController must also
+// be implemented, or SkChunkFlatController can be used to use an
+// SkChunkAllocator and never do replacements.
 //
 //
 ///////////////////////////////////////////////////////////////////////////////
 
+class SkFlatData;
+
+class SkFlatController : public SkRefCnt {
+public:
+    /**
+     * Provide a new block of memory for the SkFlatDictionary to use.
+     */
+    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.
+     */
+    virtual void unalloc(void* ptr) = 0;
+};
+
 class SkFlatData {
 public:
     /**
@@ -192,6 +215,8 @@ public:
     void* data() { return (char*)this + sizeof(*this); }
     // Our data is always 32bit aligned, so we can offer this accessor
     uint32_t* data32() { return (uint32_t*)this->data(); }
+    // Returns the size of the flattened data.
+    size_t flatSize() const { return fFlatSize; }
 
     void setSentinelInCache() {
         this->setSentinel(kInCache_Sentinel);
@@ -210,16 +235,22 @@ public:
     }
 #endif
 
-    static SkFlatData* Create(SkChunkAlloc* heap, const void* obj, int index,
+    static SkFlatData* Create(SkFlatController* controller, const void* obj, int index,
                               void (*flattenProc)(SkOrderedWriteBuffer&, const void*),
                               SkRefCntSet* refCntRecorder = NULL,
                               SkRefCntSet* faceRecorder = NULL,
-                              uint32_t writeBufferflags = 0);
+                              uint32_t writeBufferflags = 0,
+                              SkFactorySet* fset = NULL);
+
     void unflatten(void* result,
                    void (*unflattenProc)(SkOrderedReadBuffer&, void*),
                    SkRefCntPlayback* refCntPlayback = NULL,
                    SkTypefacePlayback* facePlayback = NULL) const;
 
+    // When we purge an entry, we want to reuse an old index for the new entry,
+    // so we expose this setter.
+    void setIndex(int index) { fIndex = index; }
+
     // for unittesting
     friend bool operator==(const SkFlatData& a, const SkFlatData& b) {
         size_t N = (const char*)a.dataStop() - (const char*)a.dataToCompare();
@@ -257,15 +288,30 @@ private:
 template <class T>
 class SkFlatDictionary {
 public:
-    SkFlatDictionary(SkChunkAlloc* heap) {
+    SkFlatDictionary(SkFlatController* controller, SkRefCntSet* refSet = NULL,
+                     SkRefCntSet* typeFaceSet = NULL,
+                     SkFactorySet* factorySet = NULL)
+    : fController(controller), fRefSet(refSet), fTypefaceSet(typeFaceSet)
+    , fFactorySet(factorySet) {
         fFlattenProc = NULL;
         fUnflattenProc = NULL;
-        fHeap = heap;
+        SkASSERT(controller);
+        fController->ref();
+        SkSafeRef(refSet);
+        SkSafeRef(typeFaceSet);
+        SkSafeRef(factorySet);
         // set to 1 since returning a zero from find() indicates failure
         fNextIndex = 1;
         sk_bzero(fHash, sizeof(fHash));
     }
 
+    virtual ~SkFlatDictionary() {
+        fController->unref();
+        SkSafeUnref(fRefSet);
+        SkSafeUnref(fTypefaceSet);
+        SkSafeUnref(fFactorySet);
+    }
+
     int count() const { return fData.count(); }
 
     const SkFlatData*  operator[](int index) const {
@@ -282,49 +328,66 @@ public:
         fNextIndex = 1;
         sk_bzero(fHash, sizeof(fHash));
     }
-    
+
     /**
-     * Given an element of type T it returns its index in the dictionary. If
-     * the element wasn't previously in the dictionary it is automatically added
-     *
-     *  To make the Compare function fast, we write a sentinel value at the end
-     *  of each block. The blocks in our fData[] all have a 0 sentinel. The
-     *  newly created block we're comparing against has a -1 in the sentinel.
-     *
-     *  This trick allows Compare to always loop until failure. If it fails on
-     *  the sentinal value, we know the blocks are equal.
+     * Similar to find. Allows the caller to specify an SkFlatData to replace in
+     * the case of an add. Also tells the caller whether a new SkFlatData was
+     * added and whether the old one was replaced. The parameters added and
+     * replaced are required to be non-NULL. Rather than returning the index of
+     * the entry in the dictionary, it returns the actual SkFlatData.
      */
-    int find(const T& element, SkRefCntSet* refCntRecorder = NULL,
-             SkRefCntSet* faceRecorder = NULL, uint32_t writeBufferflags = 0) {
-
-        SkFlatData* flat = SkFlatData::Create(fHeap, &element, fNextIndex,
-                fFlattenProc, refCntRecorder, faceRecorder, writeBufferflags);
-
-        int hashIndex = ChecksumToHashIndex(flat->checksum());
-        const SkFlatData* candidate = fHash[hashIndex];
-        if (candidate && !SkFlatData::Compare(flat, candidate)) {
-            (void)fHeap->unalloc(flat);
-            return candidate->index();
-        }
-        
-        int index = SkTSearch<SkFlatData>((const SkFlatData**) fData.begin(),
-                fData.count(), flat, sizeof(flat), &SkFlatData::Compare);
-        if (index >= 0) {
-            (void)fHeap->unalloc(flat);
-            fHash[hashIndex] = fData[index];
-            return fData[index]->index();
+    const SkFlatData* findAndReplace(const T& element,
+                                     uint32_t writeBufferFlags,
+                                     const SkFlatData* toReplace, bool* added,
+                                     bool* replaced) {
+        SkASSERT(added != NULL && replaced != NULL);
+        int oldCount = fData.count();
+        const SkFlatData* flat = this->findAndReturnFlat(element,
+                                                         writeBufferFlags);
+        *added = fData.count() == oldCount + 1;
+        *replaced = false;
+        if (*added && toReplace != NULL) {
+            // First, find the index of the one to replace
+            int indexToReplace = fData.find(toReplace);
+            if (indexToReplace >= 0) {
+                // findAndReturnFlat set the index to fNextIndex and increased
+                // fNextIndex by one. Reuse the index from the one being
+                // replaced and reset fNextIndex to the proper value.
+                const_cast<SkFlatData*>(flat)->setIndex(toReplace->index());
+                fNextIndex--;
+                // Remove from the array.
+                fData.remove(indexToReplace);
+                // Remove from the hash table.
+                int oldHash = ChecksumToHashIndex(toReplace->checksum());
+                if (fHash[oldHash] == toReplace) {
+                    fHash[oldHash] = NULL;
+                }
+                // Delete the actual object.
+                fController->unalloc((void*)toReplace);
+                *replaced = true;
+            }
         }
+        return flat;
+    }
 
-        index = ~index;
-        *fData.insert(index) = flat;
-        flat->setSentinelInCache();
-        fHash[hashIndex] = flat;
-        SkASSERT(fData.count() == fNextIndex);
-        return fNextIndex++;
+    /**
+     * Given an element of type T return its 1-based index in the dictionary. If
+     * the element wasn't previously in the dictionary it is automatically
+     * added.
+     *
+     * To make the Compare function fast, we write a sentinel value at the end
+     * of each block. The blocks in our fData[] all have a 0 sentinel. The
+     * newly created block we're comparing against has a -1 in the sentinel.
+     *
+     * This trick allows Compare to always loop until failure. If it fails on
+     * the sentinal value, we know the blocks are equal.
+     */
+    int find(const T& element, uint32_t writeBufferflags = 0) {
+        return this->findAndReturnFlat(element, writeBufferflags)->index();
     }
 
     /**
-     * Given a pointer to a array of type T we allocate the array and fill it
+     * Given a pointer to an array of type T we allocate the array and fill it
      * with the unflattened dictionary contents. The return value is the size of
      * the allocated array.
      */
@@ -349,9 +412,45 @@ protected:
     void (*fUnflattenProc)(SkOrderedReadBuffer&, void*);
 
 private:
-    SkChunkAlloc* fHeap;
-    int fNextIndex;
+    SkFlatController * const     fController;
+    int                          fNextIndex;
     SkTDArray<const SkFlatData*> fData;
+    SkRefCntSet*                 fRefSet;
+    SkRefCntSet*                 fTypefaceSet;
+    SkFactorySet*                fFactorySet;
+
+    const SkFlatData* findAndReturnFlat(const T& element,
+                                        uint32_t writeBufferflags) {
+        SkFlatData* flat = SkFlatData::Create(fController, &element, fNextIndex,
+                                              fFlattenProc, fRefSet,
+                                              fTypefaceSet, writeBufferflags,
+                                              fFactorySet);
+        
+        int hashIndex = ChecksumToHashIndex(flat->checksum());
+        const SkFlatData* candidate = fHash[hashIndex];
+        if (candidate && !SkFlatData::Compare(flat, candidate)) {
+            fController->unalloc(flat);
+            return candidate;
+        }
+        
+        int index = SkTSearch<SkFlatData>((const SkFlatData**) fData.begin(),
+                                          fData.count(), flat, sizeof(flat),
+                                          &SkFlatData::Compare);
+        if (index >= 0) {
+            fController->unalloc(flat);
+            fHash[hashIndex] = fData[index];
+            return fData[index];
+        }
+        
+        index = ~index;
+        *fData.insert(index) = flat;
+        SkASSERT(fData.count() == fNextIndex);
+        fNextIndex++;
+        flat->setSentinelInCache();
+        fHash[hashIndex] = flat;
+        return flat;
+    }
+    
 
     enum {
         // Determined by trying diff values on picture-recording benchmarks
@@ -393,9 +492,29 @@ static void SkUnflattenObjectProc(SkOrderedReadBuffer& buffer, void* obj) {
     ((T*)obj)->unflatten(buffer);
 }
 
+class SkChunkFlatController : public SkFlatController {
+public:
+    SkChunkFlatController(size_t minSize)
+    : fHeap(minSize) {}
+
+    virtual void* allocThrow(size_t bytes) {
+        return fHeap.allocThrow(bytes);
+    }
+
+    virtual void unalloc(void* ptr) {
+        (void) fHeap.unalloc(ptr);
+    }
+    void reset() { fHeap.reset(); }
+private:
+    SkChunkAlloc fHeap;
+};
+
 class SkBitmapDictionary : public SkFlatDictionary<SkBitmap> {
 public:
-    SkBitmapDictionary(SkChunkAlloc* heap) : SkFlatDictionary<SkBitmap>(heap) {
+    SkBitmapDictionary(SkFlatController* controller, SkRefCntSet* refSet = NULL,
+                       SkRefCntSet* typefaceSet = NULL,
+                       SkFactorySet* factorySet = NULL)
+    : SkFlatDictionary<SkBitmap>(controller, refSet, typefaceSet, factorySet) {
         fFlattenProc = &SkFlattenObjectProc<SkBitmap>;
         fUnflattenProc = &SkUnflattenObjectProc<SkBitmap>;
     }
@@ -403,7 +522,8 @@ public:
 
 class SkMatrixDictionary : public SkFlatDictionary<SkMatrix> {
  public:
-    SkMatrixDictionary(SkChunkAlloc* heap) : SkFlatDictionary<SkMatrix>(heap) {
+    SkMatrixDictionary(SkFlatController* controller)
+    : SkFlatDictionary<SkMatrix>(controller) {
         fFlattenProc = &flattenMatrix;
         fUnflattenProc = &unflattenMatrix;
     }
@@ -419,7 +539,9 @@ class SkMatrixDictionary : public SkFlatDictionary<SkMatrix> {
 
 class SkPaintDictionary : public SkFlatDictionary<SkPaint> {
  public:
-    SkPaintDictionary(SkChunkAlloc* heap) : SkFlatDictionary<SkPaint>(heap) {
+    SkPaintDictionary(SkFlatController* controller, SkRefCntSet* refSet,
+                      SkRefCntSet* typefaceSet)
+    : SkFlatDictionary<SkPaint>(controller, refSet, typefaceSet) {
         fFlattenProc = &SkFlattenObjectProc<SkPaint>;
         fUnflattenProc = &SkUnflattenObjectProc<SkPaint>;
     }
@@ -427,7 +549,8 @@ class SkPaintDictionary : public SkFlatDictionary<SkPaint> {
 
 class SkRegionDictionary : public SkFlatDictionary<SkRegion> {
  public:
-    SkRegionDictionary(SkChunkAlloc* heap) : SkFlatDictionary<SkRegion>(heap) {
+    SkRegionDictionary(SkFlatController* controller)
+    : SkFlatDictionary<SkRegion>(controller) {
         fFlattenProc = &flattenRegion;
         fUnflattenProc = &unflattenRegion;
     }
diff --git a/src/core/SkPictureRecord.cpp b/src/core/SkPictureRecord.cpp
index ed528f293c..9f5a1c22dc 100644
--- a/src/core/SkPictureRecord.cpp
+++ b/src/core/SkPictureRecord.cpp
@@ -18,9 +18,9 @@ enum {
 
 SkPictureRecord::SkPictureRecord(uint32_t flags) :
         fHeap(HEAP_BLOCK_SIZE),
-        fBitmaps(&fHeap),
+        fBitmaps(&fHeap, &fRCSet),
         fMatrices(&fHeap),
-        fPaints(&fHeap),
+        fPaints(&fHeap, &fRCSet, &fTFSet),
         fRegions(&fHeap),
         fWriter(MIN_WRITER_SIZE),
         fRecordFlags(flags) {
@@ -555,7 +555,7 @@ void SkPictureRecord::addPaint(const SkPaint& paint) {
 }
 
 void SkPictureRecord::addPaintPtr(const SkPaint* paint) {
-    this->addInt(paint ? fPaints.find(*paint, &fRCSet, &fTFSet) : 0);
+    this->addInt(paint ? fPaints.find(*paint) : 0);
 }
 
 void SkPictureRecord::addPath(const SkPath& path) {
@@ -666,7 +666,7 @@ int SkPictureRecord::find(const SkBitmap& bitmap) {
     
     uint32_t writeFlags = flattenPixels ?
         SkFlattenableWriteBuffer::kForceFlattenBitmapPixels_Flag : 0;
-    int index = fBitmaps.find(bitmap, &fRCSet, NULL, writeFlags);
+    int index = fBitmaps.find(bitmap, writeFlags);
 
     if (flattenPixels) {
         entry.fIndex = index;
diff --git a/src/core/SkPictureRecord.h b/src/core/SkPictureRecord.h
index 13b9db55c5..23d0828b97 100644
--- a/src/core/SkPictureRecord.h
+++ b/src/core/SkPictureRecord.h
@@ -187,8 +187,11 @@ public:
 #endif
 
 private:
-    SkChunkAlloc fHeap;
+    SkChunkFlatController fHeap;
 
+    SkRefCntSet fRCSet;
+    SkRefCntSet fTFSet;
+    
     SkTDArray<PixelRefDictionaryEntry> fPixelRefDictionary;
     SkBitmapDictionary fBitmaps;
     SkMatrixDictionary fMatrices;
@@ -201,9 +204,6 @@ private:
     // we ref each item in these arrays
     SkTDArray<SkPicture*> fPictureRefs;
 
-    SkRefCntSet fRCSet;
-    SkRefCntSet fTFSet;
-
     uint32_t fRecordFlags;
     int fInitialSaveCount;
 
diff --git a/src/pipe/SkGPipeRead.cpp b/src/pipe/SkGPipeRead.cpp
index 6f79e51f9a..e2cd151d28 100644
--- a/src/pipe/SkGPipeRead.cpp
+++ b/src/pipe/SkGPipeRead.cpp
@@ -105,14 +105,19 @@ public:
     }
 
     void addBitmap(int index) {
-        SkASSERT(fBitmaps.count() == index);
-        SkBitmap* bm = new SkBitmap();
+        index--;
+        SkBitmap* bm;
+        if(fBitmaps.count() == index) {
+            bm = SkNEW(SkBitmap);
+            *fBitmaps.append() = bm;
+        } else {
+            bm = fBitmaps[index];
+        }
         bm->unflatten(*fReader);
-        *fBitmaps.append() = bm;
     }
 
     SkBitmap* getBitmap(unsigned index) {
-        return fBitmaps[index];
+        return fBitmaps[index - 1];
     }
 
     void addTypeface() {
diff --git a/src/pipe/SkGPipeWrite.cpp b/src/pipe/SkGPipeWrite.cpp
index 13b1f843e6..9007163365 100644
--- a/src/pipe/SkGPipeWrite.cpp
+++ b/src/pipe/SkGPipeWrite.cpp
@@ -26,6 +26,7 @@
 #include "SkRasterizer.h"
 #include "SkShader.h"
 #include "SkOrderedWriteBuffer.h"
+#include "SkPictureFlat.h"
 
 static SkFlattenable* get_paintflat(const SkPaint& paint, unsigned paintFlat) {
     SkASSERT(paintFlat < kCount_PaintFlats);
@@ -58,6 +59,94 @@ static size_t writeTypeface(SkWriter32* writer, SkTypeface* typeface) {
 
 ///////////////////////////////////////////////////////////////////////////////
 
+class FlattenableHeap : public SkFlatController {
+public:
+    FlattenableHeap(int numFlatsToKeep) : fNumFlatsToKeep(numFlatsToKeep) {}
+
+    ~FlattenableHeap() {
+        fPointers.freeAll();
+    }
+
+    virtual void* allocThrow(size_t bytes) SK_OVERRIDE;
+
+    virtual void unalloc(void* ptr) SK_OVERRIDE;
+
+    const SkFlatData* flatToReplace() const;
+
+    // Mark an SkFlatData as one that should not be returned by flatToReplace.
+    // Takes the result of SkFlatData::index() as its parameter.
+    void markFlatForKeeping(int index) {
+        *fFlatsThatMustBeKept.append() = index;
+    }
+
+    void markAllFlatsSafeToDelete() {
+        fFlatsThatMustBeKept.reset();
+    }
+
+private:
+    // Keep track of the indices (i.e. the result of SkFlatData::index()) of
+    // flats that must be kept, since they are on the current paint.
+    SkTDArray<int>   fFlatsThatMustBeKept;
+    SkTDArray<void*> fPointers;
+    const int        fNumFlatsToKeep;
+};
+
+void FlattenableHeap::unalloc(void* ptr) {
+    int indexToRemove = fPointers.rfind(ptr);
+    if (indexToRemove >= 0) {
+        sk_free(ptr);
+        fPointers.remove(indexToRemove);
+    }
+}
+
+void* FlattenableHeap::allocThrow(size_t bytes) {
+    void* ptr = sk_malloc_throw(bytes);
+    *fPointers.append() = ptr;
+    return ptr;
+}
+
+const SkFlatData* FlattenableHeap::flatToReplace() const {
+    // First, determine whether we should replace one.
+    if (fPointers.count() > fNumFlatsToKeep) {
+        // Look through the flattenable heap.
+        // TODO: Return the LRU flat.
+        for (int i = 0; i < fPointers.count(); i++) {
+            SkFlatData* potential = (SkFlatData*)fPointers[i];
+            // Make sure that it is not one that must be kept.
+            bool mustKeep = false;
+            for (int j = 0; j < fFlatsThatMustBeKept.count(); j++) {
+                if (potential->index() == fFlatsThatMustBeKept[j]) {
+                    mustKeep = true;
+                    break;
+                }
+            }
+            if (!mustKeep) {
+                return potential;
+            }
+        }
+    }
+    return NULL;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+
+class FlatDictionary : public SkFlatDictionary<SkFlattenable> {
+public:
+    FlatDictionary(FlattenableHeap* heap, SkFactorySet* factorySet)
+            : SkFlatDictionary<SkFlattenable>(heap, NULL, NULL, factorySet) {
+        fFlattenProc = &flattenFlattenableProc;
+        // No need to define fUnflattenProc since the writer will never
+        // unflatten the data.
+    }
+    static void flattenFlattenableProc(SkOrderedWriteBuffer& buffer,
+                                       const void* obj) {
+        buffer.writeFlattenable((SkFlattenable*)obj);
+    }
+
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
 /*
  * Shared heap for storing large things that can be shared, for a stream
  * used by multiple readers.
@@ -65,15 +154,15 @@ static size_t writeTypeface(SkWriter32* writer, SkTypeface* typeface) {
  * TODO: Store paths (others?)
  * TODO: Generalize the LRU caching mechanism
  */
-class Heap {
+class SharedHeap {
 public:
-    Heap(bool shallow, int numOfReaders)
+    SharedHeap(bool shallow, int numOfReaders)
         : fBitmapCount(0)
         , fMostRecentlyUsed(NULL)
         , fLeastRecentlyUsed(NULL)
         , fCanDoShallowCopies(shallow)
         , fNumberOfReaders(numOfReaders) {}
-    ~Heap() {
+    ~SharedHeap() {
         BitmapInfo* iter = fMostRecentlyUsed;
         while (iter != NULL) {
             BitmapInfo* next = iter->fLessRecentlyUsed;
@@ -179,7 +268,7 @@ private:
 };
 
 // We just "used" info. Update our LRU accordingly
-void Heap::setMostRecentlyUsed(BitmapInfo* info) {
+void SharedHeap::setMostRecentlyUsed(BitmapInfo* info) {
     SkASSERT(info != NULL);
     if (info == fMostRecentlyUsed) {
         return;
@@ -214,7 +303,7 @@ void Heap::setMostRecentlyUsed(BitmapInfo* info) {
  * Given a new bitmap to be added to the cache, return an existing one that
  * should be removed to make room, or NULL if there is already room.
  */
-BitmapInfo* Heap::bitmapToReplace(const SkBitmap& bm) const {
+BitmapInfo* SharedHeap::bitmapToReplace(const SkBitmap& bm) const {
     // Arbitrarily set a limit of 5. We should test to find the best tradeoff
     // between time and space. A lower limit means that we use less space, but
     // it also means that we may have to insert the same bitmap into the heap
@@ -312,8 +401,7 @@ public:
     virtual void drawData(const void*, size_t) SK_OVERRIDE;
 
 private:
-    Heap fHeap;
-    SkFactorySet* fFactorySet;  // optional, only used if cross-process
+    SharedHeap fSharedHeap;
     SkGPipeController* fController;
     SkWriter32& fWriter;
     size_t      fBlockSize; // amount allocated for writer
@@ -345,27 +433,16 @@ private:
         }
     }
 
-    struct FlatData {
-        uint32_t    fIndex; // always > 0
-        PaintFlats  fPaintFlat; // deliberately before fSize so that Compare
-                                // will ignore it.
-        uint32_t    fSize;
-
-        void*       data() { return (char*)this + sizeof(*this); }
-
-        static int Compare(const FlatData* a, const FlatData* b) {
-            return memcmp(&a->fSize, &b->fSize, a->fSize + sizeof(a->fSize));
-        }
-    };
-
-    SkTDArray<FlatData*> fBitmapArray;
+    // These are only used when in cross process, but with no shared address
+    // space, so bitmaps are flattened.
+    FlattenableHeap    fBitmapHeap;
+    SkBitmapDictionary fBitmapDictionary;
     int flattenToIndex(const SkBitmap&);
 
-    SkTDArray<FlatData*> fFlatArray;
-    size_t fBytesOfFlatData;
+    FlattenableHeap fFlattenableHeap;
+    FlatDictionary  fFlatDictionary;
     int fCurrFlatIndex[kCount_PaintFlats];
     int flattenToIndex(SkFlattenable* obj, PaintFlats);
-    int flattenableToReplace(const FlatData& newFlat);
 
     SkPaint fPaint;
     void writePaint(const SkPaint&);
@@ -384,58 +461,18 @@ private:
 
 int SkGPipeCanvas::flattenToIndex(const SkBitmap & bitmap) {
     SkASSERT(shouldFlattenBitmaps(fFlags));
-    SkOrderedWriteBuffer tmpWriter(1024);
-    tmpWriter.setFlags((SkFlattenableWriteBuffer::Flags)
-                       (SkFlattenableWriteBuffer::kInlineFactoryNames_Flag
-                        | SkFlattenableWriteBuffer::kCrossProcess_Flag));
-    tmpWriter.setFactoryRecorder(fFactorySet);
-    bitmap.flatten(tmpWriter);
-
-    size_t len = tmpWriter.size();
-    size_t allocSize = len + sizeof(FlatData);
-
-    SkAutoSMalloc<1024> storage(allocSize);
-    FlatData* flat = (FlatData*)storage.get();
-    flat->fSize = len;
-    tmpWriter.flatten(flat->data());
-    
-    int index = SkTSearch<FlatData>((const FlatData**)fBitmapArray.begin(),
-                                     fBitmapArray.count(), flat, sizeof(flat),
-                                     &FlatData::Compare);
-    if (index < 0) {
-        index = ~index;
-        FlatData* copy = (FlatData*)sk_malloc_throw(allocSize);
-        memcpy(copy, flat, allocSize);
-        // For bitmaps, we can use zero based indices, since we will never ask
-        // for a NULL bitmap (unlike with paint flattenables).
-        copy->fIndex = fBitmapArray.count();
-        *fBitmapArray.insert(index) = copy;
-        if (this->needOpBytes(len)) {
-            this->writeOp(kDef_Bitmap_DrawOp, 0, copy->fIndex);
-            fWriter.write(copy->data(), len);
-        }
-    }
-    return fBitmapArray[index]->fIndex;
-}
-
-// Return -1 if there is no need to replace a flattenable, or there was not an
-// appropriate one to replace. Otherwise return the index of a flattenable to
-// replace.
-int SkGPipeCanvas::flattenableToReplace(const FlatData& newFlat) {
-    // For now, set an arbitrary limit on the size of FlatData we have stored.
-    // Note that this is currently a soft limit. If we have reached the limit,
-    // we replace one, but do not ensure that we return to below the limit.
-    if (fBytesOfFlatData + fFlatArray.bytes() > 1024) {
-        for (int i = 0; i < fFlatArray.count(); i++) {
-            // Only replace the same paint flat. Since a paint can only have
-            // one of each type, replacing one of the same type means that
-            // we will not be purging a flat on the same paint.
-            if (newFlat.fPaintFlat == fFlatArray[i]->fPaintFlat) {
-                return i;
-            }
-        }
-    }
-    return -1;
+    uint32_t flags = SkFlattenableWriteBuffer::kInlineFactoryNames_Flag
+        | SkFlattenableWriteBuffer::kCrossProcess_Flag;
+    bool added, replaced;
+    const SkFlatData* flat = fBitmapDictionary.findAndReplace(
+        bitmap, flags, fBitmapHeap.flatToReplace(), &added, &replaced);
+
+    int index = flat->index();
+    if (added && this->needOpBytes(flat->flatSize())) {
+        this->writeOp(kDef_Bitmap_DrawOp, 0, index);
+        fWriter.write(flat->data(), flat->flatSize());
+    }
+    return index;
 }
 
 // return 0 for NULL (or unflattenable obj), or index-base-1
@@ -445,79 +482,41 @@ int SkGPipeCanvas::flattenToIndex(SkFlattenable* obj, PaintFlats paintflat) {
         return 0;
     }
 
-    SkOrderedWriteBuffer tmpWriter(1024);
-    
+    uint32_t writeBufferFlags;
     if (fFlags & SkGPipeWriter::kCrossProcess_Flag) {
-        tmpWriter.setFlags((SkFlattenableWriteBuffer::Flags)
-                           (SkFlattenableWriteBuffer::kInlineFactoryNames_Flag
-                           | SkFlattenableWriteBuffer::kCrossProcess_Flag));
-        tmpWriter.setFactoryRecorder(fFactorySet);
+        writeBufferFlags = (SkFlattenableWriteBuffer::kInlineFactoryNames_Flag
+                            | SkFlattenableWriteBuffer::kCrossProcess_Flag);
     } else {
         // Needed for bitmap shaders.
-        tmpWriter.setFlags(SkFlattenableWriteBuffer::kForceFlattenBitmapPixels_Flag);
+        writeBufferFlags = SkFlattenableWriteBuffer::kForceFlattenBitmapPixels_Flag;
     }
 
-    tmpWriter.writeFlattenable(obj);
-    size_t len = tmpWriter.size();
-    size_t allocSize = len + sizeof(FlatData);
-
-    SkAutoSMalloc<1024> storage(allocSize);
-    FlatData* flat = (FlatData*)storage.get();
-    flat->fSize = len;
-    tmpWriter.flatten(flat->data());
-
-    int index = SkTSearch<FlatData>((const FlatData**)fFlatArray.begin(),
-                                    fFlatArray.count(), flat, sizeof(FlatData*),
-                                    &FlatData::Compare);
-    bool replacedAFlat = false;
-    if (index < 0) {
-        index = ~index;
-        FlatData* copy = (FlatData*)sk_malloc_throw(allocSize);
-        memcpy(copy, flat, allocSize);
-        copy->fPaintFlat = paintflat;
-        int indexToReplace = this->flattenableToReplace(*copy);
-        if (indexToReplace >= 0) {
-            replacedAFlat = true;
-            FlatData* oldData = fFlatArray[indexToReplace];
-            copy->fIndex = oldData->fIndex;
-            fBytesOfFlatData -= (sizeof(FlatData) + oldData->fSize);
-            sk_free(oldData);
-            fFlatArray.remove(indexToReplace);
-            if (indexToReplace < index) {
-                index--;
-            }
-        }
-        *fFlatArray.insert(index) = copy;
-        fBytesOfFlatData += allocSize;
-        if (!replacedAFlat) {
-            // Call this after the insert, so that count() will have been grown
-            // (unless we replaced one, in which case fIndex has already been
-            // set properly).
-            copy->fIndex = fFlatArray.count();
-//          SkDebugf("--- add flattenable[%d] size=%d index=%d\n", paintflat, len, copy->fIndex);
-        }
-
-        if (this->needOpBytes(len)) {
-            this->writeOp(kDef_Flattenable_DrawOp, paintflat, copy->fIndex);
-            fWriter.write(copy->data(), len);
-        }
+    bool added, replaced;
+    const SkFlatData* flat = fFlatDictionary.findAndReplace(
+            *obj, writeBufferFlags, fFlattenableHeap.flatToReplace(), &added, &replaced);
+    int index = flat->index();
+    if (added && this->needOpBytes(flat->flatSize())) {
+        this->writeOp(kDef_Flattenable_DrawOp, paintflat, index);
+        fWriter.write(flat->data(), flat->flatSize());
     }
-    int retVal = fFlatArray[index]->fIndex;
-    if (replacedAFlat) {
-        retVal = ~retVal;
+    if (replaced) {
+        index = ~index;
     }
-    return retVal;
+    return index;
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 
 #define MIN_BLOCK_SIZE  (16 * 1024)
+#define BITMAPS_TO_KEEP 5
+#define FLATTENABLES_TO_KEEP 10
 
 SkGPipeCanvas::SkGPipeCanvas(SkGPipeController* controller,
                              SkWriter32* writer, SkFactorySet* fset, uint32_t flags)
-: fHeap(!(flags & SkGPipeWriter::kCrossProcess_Flag), controller->numberOfReaders())
-, fWriter(*writer), fFlags(flags), fBytesOfFlatData(0) {
-    fFactorySet = fset;
+: fSharedHeap(!(flags & SkGPipeWriter::kCrossProcess_Flag), controller->numberOfReaders())
+, fWriter(*writer), fFlags(flags)
+, fBitmapHeap(BITMAPS_TO_KEEP), fBitmapDictionary(&fBitmapHeap, NULL, NULL, fset)
+, fFlattenableHeap(FLATTENABLES_TO_KEEP), fFlatDictionary(&fFlattenableHeap, fset) {
     fController = controller;
     fDone = false;
     fBlockSize = 0; // need first block from controller
@@ -539,9 +538,6 @@ SkGPipeCanvas::SkGPipeCanvas(SkGPipeController* controller,
 
 SkGPipeCanvas::~SkGPipeCanvas() {
     this->finish();
-
-    fFlatArray.freeAll();
-    fBitmapArray.freeAll();
 }
 
 bool SkGPipeCanvas::needOpBytes(size_t needed) {
@@ -791,7 +787,7 @@ void SkGPipeCanvas::drawBitmap(const SkBitmap& bm, SkScalar left, SkScalar top,
     if (flatten) {
         bitmapIndex = this->flattenToIndex(bm);
     } else {
-        ptr = fHeap.addBitmap(bm);
+        ptr = fSharedHeap.addBitmap(bm);
         if (NULL == ptr) {
             return;
         }
@@ -825,7 +821,7 @@ void SkGPipeCanvas::drawBitmapRect(const SkBitmap& bm, const SkIRect* src,
     if (flatten) {
         bitmapIndex = this->flattenToIndex(bm);
     } else {
-        ptr = fHeap.addBitmap(bm);
+        ptr = fSharedHeap.addBitmap(bm);
         if (NULL == ptr) {
             return;
         }
@@ -874,7 +870,7 @@ void SkGPipeCanvas::drawBitmapNine(const SkBitmap& bm, const SkIRect& center,
     if (flatten) {
         bitmapIndex = this->flattenToIndex(bm);
     } else {
-        ptr = fHeap.addBitmap(bm);
+        ptr = fSharedHeap.addBitmap(bm);
         if (NULL == ptr) {
             return;
         }
@@ -911,7 +907,7 @@ void SkGPipeCanvas::drawSprite(const SkBitmap& bm, int left, int top,
     if (flatten) {
         bitmapIndex = this->flattenToIndex(bm);
     } else {
-        ptr = fHeap.addBitmap(bm);
+        ptr = fSharedHeap.addBitmap(bm);
         if (NULL == ptr) {
             return;
         }
@@ -1162,13 +1158,19 @@ void SkGPipeCanvas::writePaint(const SkPaint& paint) {
         base.setTypeface(paint.getTypeface());
     }
 
+    // This is a new paint, so all old flats can be safely purged, if necessary.
+    fFlattenableHeap.markAllFlatsSafeToDelete();
     for (int i = 0; i < kCount_PaintFlats; i++) {
         int index = this->flattenToIndex(get_paintflat(paint, i), (PaintFlats)i);
         bool replaced = index < 0;
         if (replaced) {
             index = ~index;
         }
-        SkASSERT(index >= 0 && index <= fFlatArray.count());
+        // Store the index of any flat that needs to be kept. 0 means no flat.
+        if (index > 0) {
+            fFlattenableHeap.markFlatForKeeping(index);
+        }
+        SkASSERT(index >= 0 && index <= fFlatDictionary.count());
         if (index != fCurrFlatIndex[i] || replaced) {
             *ptr++ = PaintOp_packOpFlagData(kFlatIndex_PaintOp, i, index);
             fCurrFlatIndex[i] = index;