Simplify reference management in SkPDF

src/pdf/SkPDFDocument.cpp

Index: src/pdf/SkPDFDocument.cpp
diff --git a/src/pdf/SkPDFDocument.cpp b/src/pdf/SkPDFDocument.cpp
index c522b39516ee3cdbfa06458aea3131bb076ab37f..0b62c64e80a922453b9c118ae1c9514a0745537d 100644
--- a/src/pdf/SkPDFDocument.cpp
+++ b/src/pdf/SkPDFDocument.cpp
@@ -82,33 +82,6 @@ SkPDFDocument::~SkPDFDocument() {
     SkDELETE(fOtherPageResources);
 }
 
-namespace {
-class Streamer {
-public:
-    Streamer(SkPDFCatalog* cat, SkWStream* out)
-        : fCat(cat), fOut(out), fBaseOffset(SkToOffT(out->bytesWritten())) {
-    }
-
-    void stream(SkPDFObject* object) {
-        fCat->setFileOffset(object, this->offset());
-        SkPDFObject* realObject = fCat->getSubstituteObject(object);
-        fCat->emitObjectNumber(fOut, realObject);
-        fOut->writeText(" obj\n");
-        realObject->emitObject(fOut, fCat);
-        fOut->writeText("\nendobj\n");
-    }
-
-    off_t offset() {
-        return SkToOffT(fOut->bytesWritten()) - fBaseOffset;
-    }
-
-private:
-    SkPDFCatalog* const fCat;
-    SkWStream* const fOut;
-    const off_t fBaseOffset;
-};
-}  // namespace
-
 bool SkPDFDocument::emitPDF(SkWStream* stream) {
     if (fPages.isEmpty()) {
         return false;
@@ -187,47 +160,23 @@ bool SkPDFDocument::emitPDF(SkWStream* stream) {
         perform_font_subsetting(fCatalog.get(), fPages, &fSubstitutes);
     }
 
-    Streamer out(fCatalog, stream);
-    emitHeader(stream);
-
-    out.stream(fDocCatalog);
-    out.stream(fPages[0]);
-    out.stream(fPages[0]->getContentStream());
-
-    for (int i = 0; i < fFirstPageResources->count(); i++) {
-        out.stream((*fFirstPageResources)[i]);
-    }
-
-    SkTSet<SkPDFObject*>* firstPageSubstituteResources =
-            fCatalog->getSubstituteList(true);
-    for (int i = 0; i < firstPageSubstituteResources->count(); ++i) {
-        out.stream((*firstPageSubstituteResources)[i]);
-    }
-    // TODO(vandebo): Support linearized format
-    // if (fPages.size() > 1) {
-    //     // TODO(vandebo): Save the file offset for the first page xref table.
-    //     fCatalog->emitXrefTable(stream, true);
-    // }
-
-    for (int i = 0; i < fPageTree.count(); i++) {
-        out.stream(fPageTree[i]);
-    }
-
-    for (int i = 1; i < fPages.count(); i++) {
-        out.stream(fPages[i]->getContentStream());
+    SkPDFObjectSet resourceSet;
+    if (resourceSet.add(fDocCatalog)) {
+        fDocCatalog->addResources(&resourceSet, fCatalog);
     }
-
-    for (int i = 0; i < fOtherPageResources->count(); i++) {
-        out.stream((*fOtherPageResources)[i]);
-    }
-
-    SkTSet<SkPDFObject*>* otherSubstituteResources =
-            fCatalog->getSubstituteList(false);
-    for (int i = 0; i < otherSubstituteResources->count(); ++i) {
-        out.stream((*otherSubstituteResources)[i]);
+    off_t baseOffset = SkToOffT(stream->bytesWritten());
+    emitHeader(stream);
+    for (int i = 0; i < resourceSet.count(); ++i) {
+        SkPDFObject* object = resourceSet[i];
+        fCatalog->setFileOffset(object,
+                                SkToOffT(stream->bytesWritten()) - baseOffset);
+        SkASSERT(object == fCatalog->getSubstituteObject(object));
+        fCatalog->emitObjectNumber(stream, object);
+        stream->writeText(" obj\n");
+        object->emitObject(stream, fCatalog);
+        stream->writeText("\nendobj\n");
     }
-
-    fXRefFileOffset = out.offset();
+    fXRefFileOffset = SkToOffT(stream->bytesWritten()) - baseOffset;
     int64_t objCount = fCatalog->emitXrefTable(stream, fPages.count() > 1);
     emitFooter(stream, objCount);
     return true;