Enable flattening/unflattening with custom unflatten procs

Now flattenables are serialized using a string name, so that
flattenables do not necessarily need to be registered before
serialization.  They just need to override getTypeName().

Allows custom unflatten procs to be set on the SkReadBuffer.
This is optional if the flattenable is registered, but otherwise
must be called.

This was split off from:
https://codereview.chromium.org/1837913003/

BUG=skia:
GOLD_TRYBOT_URL= https://gold.skia.org/search2?unt=true&query=source_type%3Dgm&master=false&issue=1858323002

Review URL: https://codereview.chromium.org/1858323002

11 files changed, 272 insertions, 65 deletions

diff --git a/include/core/SkFlattenable.h b/include/core/SkFlattenable.h
index c76f119c13..0ba83da087 100644
--- a/include/core/SkFlattenable.h
+++ b/include/core/SkFlattenable.h
@@ -92,9 +92,15 @@ public:
      */
     virtual Factory getFactory() const = 0;
 
-    /** Returns the name of the object's class
-      */
-    const char* getTypeName() const { return FactoryToName(getFactory()); }
+    /**
+     *  Returns the name of the object's class.
+     *
+     *  Subclasses should override this function if they intend to provide
+     *  support for flattening without using the global registry.
+     *
+     *  If the flattenable is registered, there is no need to override.
+     */
+    virtual const char* getTypeName() const { return FactoryToName(getFactory()); }
 
     static Factory NameToFactory(const char name[]);
     static const char* FactoryToName(Factory);
diff --git a/include/core/SkWriteBuffer.h b/include/core/SkWriteBuffer.h
index 306a3e37a9..c931ad3fd2 100644
--- a/include/core/SkWriteBuffer.h
+++ b/include/core/SkWriteBuffer.h
@@ -16,6 +16,7 @@
 #include "SkPixelSerializer.h"
 #include "SkRefCnt.h"
 #include "SkWriter32.h"
+#include "../private/SkTHash.h"
 
 class SkBitmap;
 class SkBitmapHeap;
@@ -27,7 +28,6 @@ class SkWriteBuffer {
 public:
     enum Flags {
         kCrossProcess_Flag  = 1 << 0,
-        kValidation_Flag    = 1 << 1,
     };
 
     SkWriteBuffer(uint32_t flags = 0);
@@ -35,7 +35,7 @@ public:
     ~SkWriteBuffer();
 
     bool isCrossProcess() const {
-        return this->isValidating() || SkToBool(fFlags & kCrossProcess_Flag);
+        return SkToBool(fFlags & kCrossProcess_Flag);
     }
 
     SkWriter32* getWriter32() { return &fWriter; }
@@ -108,8 +108,6 @@ public:
     SkPixelSerializer* getPixelSerializer() const { return fPixelSerializer; }
 
 private:
-    bool isValidating() const { return SkToBool(fFlags & kValidation_Flag); }
-
     const uint32_t fFlags;
     SkFactorySet* fFactorySet;
     SkWriter32 fWriter;
@@ -118,6 +116,9 @@ private:
     SkRefCntSet* fTFSet;
 
     SkAutoTUnref<SkPixelSerializer> fPixelSerializer;
+
+    // Only used if we do not have an fFactorySet
+    SkTHashMap<SkString, uint32_t> fFlattenableDict;
 };
 
 #endif // SkWriteBuffer_DEFINED
diff --git a/src/core/SkFlattenableSerialization.cpp b/src/core/SkFlattenableSerialization.cpp
index e9ce09ff1f..06e8c10706 100644
--- a/src/core/SkFlattenableSerialization.cpp
+++ b/src/core/SkFlattenableSerialization.cpp
@@ -12,7 +12,7 @@
 #include "SkWriteBuffer.h"
 
 SkData* SkValidatingSerializeFlattenable(SkFlattenable* flattenable) {
-    SkWriteBuffer writer(SkWriteBuffer::kValidation_Flag);
+    SkWriteBuffer writer;
     writer.writeFlattenable(flattenable);
     size_t size = writer.bytesWritten();
     auto data = SkData::MakeUninitialized(size);
diff --git a/src/core/SkReadBuffer.cpp b/src/core/SkReadBuffer.cpp
index bafcbc2b91..70b0d13a40 100644
--- a/src/core/SkReadBuffer.cpp
+++ b/src/core/SkReadBuffer.cpp
@@ -106,6 +106,11 @@ int32_t SkReadBuffer::read32() {
     return fReader.readInt();
 }
 
+uint8_t SkReadBuffer::peekByte() {
+    SkASSERT(fReader.available() > 0);
+    return *((uint8_t*) fReader.peek());
+}
+
 void SkReadBuffer::readString(SkString* string) {
     size_t len;
     const char* strContents = fReader.readString(&len);
@@ -348,9 +353,32 @@ SkFlattenable* SkReadBuffer::readFlattenable(SkFlattenable::Type ft) {
         }
         factory = fFactoryArray[index];
     } else {
-        factory = (SkFlattenable::Factory)readFunctionPtr();
-        if (nullptr == factory) {
-            return nullptr; // writer failed to give us the flattenable
+        SkString name;
+        if (this->peekByte()) {
+            // If the first byte is non-zero, the flattenable is specified by a string.
+            this->readString(&name);
+
+            // Add the string to the dictionary.
+            fFlattenableDict.set(fFlattenableDict.count() + 1, name);
+        } else {
+            // Read the index.  We are guaranteed that the first byte
+            // is zeroed, so we must shift down a byte.
+            uint32_t index = fReader.readU32() >> 8;
+            if (0 == index) {
+                return nullptr; // writer failed to give us the flattenable
+            }
+
+            SkString* namePtr = fFlattenableDict.find(index);
+            SkASSERT(namePtr);
+            name = *namePtr;
+        }
+
+        // Check if a custom Factory has been specified for this flattenable.
+        if (!(factory = this->getCustomFactory(name))) {
+            // If there is no custom Factory, check for a default.
+            if (!(factory = SkFlattenable::NameToFactory(name.c_str()))) {
+                return nullptr; // writer failed to give us the flattenable
+            }
         }
     }
 
diff --git a/src/core/SkReadBuffer.h b/src/core/SkReadBuffer.h
index e73f4cf9d9..8f15a8d366 100644
--- a/src/core/SkReadBuffer.h
+++ b/src/core/SkReadBuffer.h
@@ -22,6 +22,7 @@
 #include "SkReader32.h"
 #include "SkRefCnt.h"
 #include "SkShader.h"
+#include "SkTHash.h"
 #include "SkWriteBuffer.h"
 #include "SkXfermode.h"
 
@@ -116,6 +117,9 @@ public:
     virtual uint32_t readUInt();
     virtual int32_t read32();
 
+    // peek
+    virtual uint8_t peekByte();
+
     // strings -- the caller is responsible for freeing the string contents
     virtual void readString(SkString* string);
     virtual void* readEncodedString(size_t* length, SkPaint::TextEncoding encoding);
@@ -200,6 +204,20 @@ public:
     }
 
     /**
+     *  For an input flattenable (specified by name), set a custom factory proc
+     *  to use when unflattening.  Will make a copy of |name|.
+     *
+     *  If the global registry already has a default factory for the flattenable,
+     *  this will override that factory.  If a custom factory has already been
+     *  set for the flattenable, this will override that factory.
+     *
+     *  Custom factories can be removed by calling setCustomFactory("...", nullptr).
+     */
+    void setCustomFactory(const SkString& name, SkFlattenable::Factory factory) {
+        fCustomFactory.set(name, factory);
+    }
+
+    /**
      *  Provide a function to decode an SkBitmap from encoded data. Only used if the writer
      *  encoded the SkBitmap. If the proper decoder cannot be used, a red bitmap with the
      *  appropriate size will be used.
@@ -217,8 +235,27 @@ public:
     }
 
 protected:
+    /**
+     *  Allows subclass to check if we are using factories for expansion
+     *  of flattenables.
+     */
+    int factoryCount() { return fFactoryCount; }
+
+
+    /**
+     *  Checks if a custom factory has been set for a given flattenable.
+     *  Returns the custom factory if it exists, or nullptr otherwise.
+     */
+    SkFlattenable::Factory getCustomFactory(const SkString& name) {
+        SkFlattenable::Factory* factoryPtr = fCustomFactory.find(name);
+        return factoryPtr ? *factoryPtr : nullptr;
+    }
+
     SkReader32 fReader;
 
+    // Only used if we do not have an fFactoryArray.
+    SkTHashMap<uint32_t, SkString> fFlattenableDict;
+
 private:
     bool readArray(void* value, size_t size, size_t elementSize);
 
@@ -234,6 +271,9 @@ private:
     SkFlattenable::Factory* fFactoryArray;
     int                     fFactoryCount;
 
+    // Only used if we do not have an fFactoryArray.
+    SkTHashMap<SkString, SkFlattenable::Factory> fCustomFactory;
+
     SkPicture::InstallPixelRefProc fBitmapDecoder;
 
 #ifdef DEBUG_NON_DETERMINISTIC_ASSERT
diff --git a/src/core/SkValidatingReadBuffer.cpp b/src/core/SkValidatingReadBuffer.cpp
index b85e532d0c..bc4042a2d7 100644
--- a/src/core/SkValidatingReadBuffer.cpp
+++ b/src/core/SkValidatingReadBuffer.cpp
@@ -86,6 +86,14 @@ int32_t SkValidatingReadBuffer::read32() {
     return this->readInt();
 }
 
+uint8_t SkValidatingReadBuffer::peekByte() {
+    if (fReader.available() <= 0) {
+        fError = true;
+        return 0;
+    }
+    return *((uint8_t*) fReader.peek());
+}
+
 void SkValidatingReadBuffer::readString(SkString* string) {
     const size_t len = this->readUInt();
     const void* ptr = fReader.peek();
@@ -226,12 +234,37 @@ bool SkValidatingReadBuffer::validateAvailable(size_t size) {
 }
 
 SkFlattenable* SkValidatingReadBuffer::readFlattenable(SkFlattenable::Type type) {
-    SkString name;
-    this->readString(&name);
+    // The validating read buffer always uses strings and string-indices for unflattening.
+    SkASSERT(0 == this->factoryCount());
+
+    uint8_t firstByte = this->peekByte();
     if (fError) {
         return nullptr;
     }
 
+    SkString name;
+    if (firstByte) {
+        // If the first byte is non-zero, the flattenable is specified by a string.
+        this->readString(&name);
+        if (fError) {
+            return nullptr;
+        }
+
+        // Add the string to the dictionary.
+        fFlattenableDict.set(fFlattenableDict.count() + 1, name);
+    } else {
+        // Read the index.  We are guaranteed that the first byte
+        // is zeroed, so we must shift down a byte.
+        uint32_t index = fReader.readU32() >> 8;
+        if (0 == index) {
+            return nullptr; // writer failed to give us the flattenable
+        }
+
+        SkString* namePtr = fFlattenableDict.find(index);
+        SkASSERT(namePtr);
+        name = *namePtr;
+    }
+
     // Is this the type we wanted ?
     const char* cname = name.c_str();
     SkFlattenable::Type baseType;
@@ -239,28 +272,25 @@ SkFlattenable* SkValidatingReadBuffer::readFlattenable(SkFlattenable::Type type)
         return nullptr;
     }
 
-    SkFlattenable::Factory factory = SkFlattenable::NameToFactory(cname);
-    if (nullptr == factory) {
-        return nullptr; // writer failed to give us the flattenable
+    // Get the factory for this flattenable.
+    SkFlattenable::Factory factory = this->getCustomFactory(name);
+    if (!factory) {
+        factory = SkFlattenable::NameToFactory(cname);
+        if (!factory) {
+            return nullptr; // writer failed to give us the flattenable
+        }
     }
 
-    // if we get here, factory may still be null, but if that is the case, the
-    // failure was ours, not the writer.
+    // If we get here, the factory is non-null.
     sk_sp<SkFlattenable> obj;
     uint32_t sizeRecorded = this->readUInt();
-    if (factory) {
-        size_t offset = fReader.offset();
-        obj = (*factory)(*this);
-        // check that we read the amount we expected
-        size_t sizeRead = fReader.offset() - offset;
-        this->validate(sizeRecorded == sizeRead);
-        if (fError) {
-            obj = nullptr;
-        }
-    } else {
-        // we must skip the remaining data
-        this->skip(sizeRecorded);
-        SkASSERT(false);
+    size_t offset = fReader.offset();
+    obj = (*factory)(*this);
+    // check that we read the amount we expected
+    size_t sizeRead = fReader.offset() - offset;
+    this->validate(sizeRecorded == sizeRead);
+    if (fError) {
+        obj = nullptr;
     }
     return obj.release();
 }
diff --git a/src/core/SkValidatingReadBuffer.h b/src/core/SkValidatingReadBuffer.h
index 785b15e291..aaea1493ef 100644
--- a/src/core/SkValidatingReadBuffer.h
+++ b/src/core/SkValidatingReadBuffer.h
@@ -37,6 +37,9 @@ public:
     uint32_t readUInt() override;
     int32_t read32() override;
 
+    // peek
+    uint8_t peekByte() override;
+
     // strings -- the caller is responsible for freeing the string contents
     void readString(SkString* string) override;
     void* readEncodedString(size_t* length, SkPaint::TextEncoding encoding) override;
diff --git a/src/core/SkWriteBuffer.cpp b/src/core/SkWriteBuffer.cpp
index 1674b931ee..b90a81e631 100644
--- a/src/core/SkWriteBuffer.cpp
+++ b/src/core/SkWriteBuffer.cpp
@@ -258,47 +258,53 @@ void SkWriteBuffer::setPixelSerializer(SkPixelSerializer* serializer) {
 
 void SkWriteBuffer::writeFlattenable(const SkFlattenable* flattenable) {
     /*
-     *  If we have a factoryset, then the first 32bits tell us...
+     *  The first 32 bits tell us...
      *       0: failure to write the flattenable
-     *      >0: (1-based) index into the SkFactorySet or SkNamedFactorySet
-     *  If we don't have a factoryset, then the first "ptr" is either the
-     *  factory, or null for failure.
-     *
-     *  The distinction is important, since 0-index is 32bits (always), but a
-     *  0-functionptr might be 32 or 64 bits.
+     *      >0: index (1-based) into fFactorySet or fFlattenableDict or
+     *          the first character of a string
      */
     if (nullptr == flattenable) {
-        if (this->isValidating()) {
-            this->writeString("");
-        } else if (fFactorySet != nullptr) {
-            this->write32(0);
-        } else {
-            this->writeFunctionPtr(nullptr);
-        }
+        this->write32(0);
         return;
     }
 
-    SkFlattenable::Factory factory = flattenable->getFactory();
-    SkASSERT(factory != nullptr);
-
     /*
-     *  We can write 1 of 3 versions of the flattenable:
-     *  1.  function-ptr : this is the fastest for the reader, but assumes that
-     *      the writer and reader are in the same process.
-     *  2.  index into fFactorySet : This is assumes the writer will later
+     *  We can write 1 of 2 versions of the flattenable:
+     *  1.  index into fFactorySet : This assumes the writer will later
      *      resolve the function-ptrs into strings for its reader. SkPicture
      *      does exactly this, by writing a table of names (matching the indices)
      *      up front in its serialized form.
-     *  3.  index into fNamedFactorySet. fNamedFactorySet will also store the
-     *      name. SkGPipe uses this technique so it can write the name to its
-     *      stream before writing the flattenable.
+     *  2.  string name of the flattenable or index into fFlattenableDict:  We
+     *      store the string to allow the reader to specify its own factories
+     *      after write time.  In order to improve compression, if we have
+     *      already written the string, we write its index instead.
      */
-    if (this->isValidating()) {
-        this->writeString(flattenable->getTypeName());
-    } else if (fFactorySet) {
+    if (fFactorySet) {
+        SkFlattenable::Factory factory = flattenable->getFactory();
+        SkASSERT(factory);
         this->write32(fFactorySet->add(factory));
     } else {
-        this->writeFunctionPtr((void*)factory);
+        const char* name = flattenable->getTypeName();
+        SkASSERT(name);
+        SkString key(name);
+        if (uint32_t* indexPtr = fFlattenableDict.find(key)) {
+            // We will write the index as a 32-bit int.  We want the first byte
+            // that we send to be zero - this will act as a sentinel that we
+            // have an index (not a string).  This means that we will send the
+            // the index shifted left by 8.  The remaining 24-bits should be
+            // plenty to store the index.  Note that this strategy depends on
+            // being little endian.
+            SkASSERT(0 == *indexPtr >> 24);
+            this->write32(*indexPtr << 8);
+        } else {
+            // Otherwise write the string.  Clients should not use the empty
+            // string as a name, or we will have a problem.
+            SkASSERT(strcmp("", name));
+            this->writeString(name);
+
+            // Add key to dictionary.
+            fFlattenableDict.set(key, fFlattenableDict.count() + 1);
+        }
     }
 
     // make room for the size of the flattened object
diff --git a/tests/BitmapHeapTest.cpp b/tests/BitmapHeapTest.cpp
index 89e6faf764..f9d4ee19bb 100644
--- a/tests/BitmapHeapTest.cpp
+++ b/tests/BitmapHeapTest.cpp
@@ -88,9 +88,7 @@ DEF_TEST(BitmapHeap, reporter) {
     index = dictionary.find(*bitmapShader);
     heap.endAddingOwnersDeferral(false);
 
-    // The dictionary should report the same index since the new entry is identical.
     // The bitmap heap should contain the bitmap, but with no references.
-    REPORTER_ASSERT(reporter, 1 == index);
     REPORTER_ASSERT(reporter, heap.count() == 1);
     REPORTER_ASSERT(reporter, SkBitmapHeapTester::GetRefCount(heap.getEntry(0)) == 0);
 }
diff --git a/tests/FlattenableCustomFactory.cpp b/tests/FlattenableCustomFactory.cpp
new file mode 100644
index 0000000000..794f76872a
--- /dev/null
+++ b/tests/FlattenableCustomFactory.cpp
@@ -0,0 +1,95 @@
+/*
+ * Copyright 2016 Google Inc.
+ *
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include "SkFlattenable.h"
+#include "SkReadBuffer.h"
+#include "SkWriteBuffer.h"
+#include "Test.h"
+
+class IntFlattenable : public SkFlattenable {
+public:
+    IntFlattenable(uint32_t a, uint32_t b, uint32_t c, uint32_t d)
+        : fA(a)
+        , fB(b)
+        , fC(c)
+        , fD(d)
+    {}
+
+    void flatten(SkWriteBuffer& buffer) const override {
+        buffer.writeUInt(fA);
+        buffer.writeUInt(fB);
+        buffer.writeUInt(fC);
+        buffer.writeUInt(fD);
+    }
+
+    Factory getFactory() const override { return nullptr; }
+
+    uint32_t a() const { return fA; }
+    uint32_t b() const { return fB; }
+    uint32_t c() const { return fC; }
+    uint32_t d() const { return fD; }
+
+    const char* getTypeName() const override { return "IntFlattenable"; }
+
+private:
+    uint32_t fA;
+    uint32_t fB;
+    uint32_t fC;
+    uint32_t fD;
+};
+
+static sk_sp<SkFlattenable> custom_create_proc(SkReadBuffer& buffer) {
+    uint32_t a = buffer.readUInt();
+    uint32_t b = buffer.readUInt();
+    uint32_t c = buffer.readUInt();
+    uint32_t d = buffer.readUInt();
+    return sk_sp<SkFlattenable>(new IntFlattenable(a + 1, b + 1, c + 1, d + 1));
+}
+
+DEF_TEST(UnflattenWithCustomFactory, r) {
+    // Create and flatten the test flattenable
+    SkWriteBuffer writeBuffer;
+    SkAutoTUnref<SkFlattenable> flattenable1(new IntFlattenable(1, 2, 3, 4));
+    writeBuffer.writeFlattenable(flattenable1);
+    SkAutoTUnref<SkFlattenable> flattenable2(new IntFlattenable(2, 3, 4, 5));
+    writeBuffer.writeFlattenable(flattenable2);
+    SkAutoTUnref<SkFlattenable> flattenable3(new IntFlattenable(3, 4, 5, 6));
+    writeBuffer.writeFlattenable(flattenable3);
+
+    // Copy the contents of the write buffer into a read buffer
+    sk_sp<SkData> data = SkData::MakeUninitialized(writeBuffer.bytesWritten());
+    writeBuffer.writeToMemory(data->writable_data());
+    SkReadBuffer readBuffer(data->data(), data->size());
+
+    // Register a custom factory with the read buffer
+    readBuffer.setCustomFactory(SkString("IntFlattenable"), &custom_create_proc);
+
+    // Unflatten and verify the flattenables
+    SkAutoTUnref<IntFlattenable> out1((IntFlattenable*) readBuffer.readFlattenable(
+            SkFlattenable::kSkUnused_Type));
+    REPORTER_ASSERT(r, out1);
+    REPORTER_ASSERT(r, 2 == out1->a());
+    REPORTER_ASSERT(r, 3 == out1->b());
+    REPORTER_ASSERT(r, 4 == out1->c());
+    REPORTER_ASSERT(r, 5 == out1->d());
+
+    SkAutoTUnref<IntFlattenable> out2((IntFlattenable*) readBuffer.readFlattenable(
+            SkFlattenable::kSkUnused_Type));
+    REPORTER_ASSERT(r, out2);
+    REPORTER_ASSERT(r, 3 == out2->a());
+    REPORTER_ASSERT(r, 4 == out2->b());
+    REPORTER_ASSERT(r, 5 == out2->c());
+    REPORTER_ASSERT(r, 6 == out2->d());
+
+    SkAutoTUnref<IntFlattenable> out3((IntFlattenable*) readBuffer.readFlattenable(
+            SkFlattenable::kSkUnused_Type));
+    REPORTER_ASSERT(r, out3);
+    REPORTER_ASSERT(r, 4 == out3->a());
+    REPORTER_ASSERT(r, 5 == out3->b());
+    REPORTER_ASSERT(r, 6 == out3->c());
+    REPORTER_ASSERT(r, 7 == out3->d());
+}
diff --git a/tests/SerializationTest.cpp b/tests/SerializationTest.cpp
index 9bdfe1e0e6..a8f253f025 100644
--- a/tests/SerializationTest.cpp
+++ b/tests/SerializationTest.cpp
@@ -139,7 +139,7 @@ template<> struct SerializationTestUtils<SkString, true> {
 
 template<typename T, bool testInvalid>
 static void TestObjectSerializationNoAlign(T* testObj, skiatest::Reporter* reporter) {
-    SkWriteBuffer writer(SkWriteBuffer::kValidation_Flag);
+    SkWriteBuffer writer;
     SerializationUtils<T>::Write(writer, testObj);
     size_t bytesWritten = writer.bytesWritten();
     REPORTER_ASSERT(reporter, SkAlign4(bytesWritten) == bytesWritten);
@@ -177,7 +177,7 @@ static void TestObjectSerialization(T* testObj, skiatest::Reporter* reporter) {
 template<typename T>
 static T* TestFlattenableSerialization(T* testObj, bool shouldSucceed,
                                        skiatest::Reporter* reporter) {
-    SkWriteBuffer writer(SkWriteBuffer::kValidation_Flag);
+    SkWriteBuffer writer;
     SerializationUtils<T>::Write(writer, testObj);
     size_t bytesWritten = writer.bytesWritten();
     REPORTER_ASSERT(reporter, SkAlign4(bytesWritten) == bytesWritten);
@@ -215,7 +215,7 @@ static T* TestFlattenableSerialization(T* testObj, bool shouldSucceed,
 
 template<typename T>
 static void TestArraySerialization(T* data, skiatest::Reporter* reporter) {
-    SkWriteBuffer writer(SkWriteBuffer::kValidation_Flag);
+    SkWriteBuffer writer;
     SerializationUtils<T>::Write(writer, data, kArraySize);
     size_t bytesWritten = writer.bytesWritten();
     // This should write the length (in 4 bytes) and the array
@@ -533,7 +533,7 @@ DEF_TEST(Serialization, reporter) {
         sk_sp<SkPicture> pict(recorder.finishRecordingAsPicture());
 
         // Serialize picture
-        SkWriteBuffer writer(SkWriteBuffer::kValidation_Flag);
+        SkWriteBuffer writer;
         pict->flatten(writer);
         size_t size = writer.bytesWritten();
         SkAutoTMalloc<unsigned char> data(size);