[graphite] Move id off of UniformData and into UniformCache

Also renames UniformCache::findOrCreate() to insert(), and has it return
the assigned id. ExtractCombo() now does not take a UniformCache, since
it just has to return the CPU data extracted from the PaintParams. The
caching/id assignment is the responsibility of the DrawPass.

Updates unit tests and DrawPass to use this instead.

Bug: skia:12466
Change-Id: Ib363fcad799f62d0ad31bb9d3e4c995e191cae80
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/475798
Commit-Queue: Michael Ludwig <michaelludwig@google.com>
Reviewed-by: Jim Van Verth <jvanverth@google.com>
diff --git a/experimental/graphite/src/ContextUtils.cpp b/experimental/graphite/src/ContextUtils.cpp
index 54cf6de..a42c6b4 100644
--- a/experimental/graphite/src/ContextUtils.cpp
+++ b/experimental/graphite/src/ContextUtils.cpp
@@ -11,7 +11,6 @@
 #include "experimental/graphite/src/DrawList.h" // TODO: split PaintParams out into their own header
 #include "experimental/graphite/src/DrawTypes.h"
 #include "experimental/graphite/src/Uniform.h"
-#include "experimental/graphite/src/UniformCache.h"
 #include "experimental/graphite/src/UniformManager.h"
 #include "include/core/SkPaint.h"
 
@@ -179,8 +178,7 @@
     return sk_sp<UniformData>(new UniformData(count, uniforms, offsets, data, dataSize));
 }
 
-std::tuple<Combination, sk_sp<UniformData>> ExtractCombo(UniformCache* cache,
-                                                         const PaintParams& p) {
+std::tuple<Combination, sk_sp<UniformData>> ExtractCombo(const PaintParams& p) {
     Combination result;
     sk_sp<UniformData> uniforms;
 
@@ -267,9 +265,7 @@
     }
 
     result.fBlendMode = p.blendMode();
-
-    sk_sp<UniformData> trueUD = cache->findOrCreate(std::move(uniforms));
-    return { result, std::move(trueUD) };
+    return { result, std::move(uniforms) };
 }
 
 namespace {
diff --git a/experimental/graphite/src/ContextUtils.h b/experimental/graphite/src/ContextUtils.h
index a9f7e84..2bf9f3c 100644
--- a/experimental/graphite/src/ContextUtils.h
+++ b/experimental/graphite/src/ContextUtils.h
@@ -17,7 +17,6 @@
 
 class PaintParams;
 class Uniform;
-class UniformCache;
 
 // A single, fully specified combination resulting from a PaintCombo (i.e., it corresponds to a
 // specific skgpu::PaintParams object (a subset of SkPaint))
@@ -43,7 +42,6 @@
 
 class UniformData : public SkRefCnt {
 public:
-    static constexpr uint32_t kInvalidUniformID = 0;
 
     // TODO: should we require a name (e.g., "gradient_uniforms") for each uniform block so
     // we can better name the Metal FS uniform struct?
@@ -57,11 +55,6 @@
         delete [] fData;
     }
 
-    void setID(uint32_t id) {   // TODO: maybe make privileged for only UniformCache
-        SkASSERT(fID == kInvalidUniformID);
-        fID = id;
-    }
-    uint32_t id() const { return fID; }
     int count() const { return fCount; }
     const Uniform* uniforms() const { return fUniforms; }
     uint32_t* offsets() { return fOffsets; }
@@ -85,7 +78,6 @@
             , fDataSize(dataSize) {
     }
 
-    uint32_t fID = kInvalidUniformID;
     const int fCount;
     const Uniform* fUniforms;
     uint32_t* fOffsets; // offset of each uniform in 'fData'
@@ -93,7 +85,7 @@
     const size_t fDataSize;
 };
 
-std::tuple<Combination, sk_sp<UniformData>> ExtractCombo(UniformCache*, const PaintParams&);
+std::tuple<Combination, sk_sp<UniformData>> ExtractCombo(const PaintParams&);
 std::string GetMSLUniformStruct(ShaderCombo::ShaderType);
 
 } // namespace skgpu
diff --git a/experimental/graphite/src/DrawPass.cpp b/experimental/graphite/src/DrawPass.cpp
index 95016d8..efbd516 100644
--- a/experimental/graphite/src/DrawPass.cpp
+++ b/experimental/graphite/src/DrawPass.cpp
@@ -34,10 +34,10 @@
 std::tuple<uint32_t, uint32_t> get_ids_from_paint(skgpu::Recorder* recorder,
                                                   skgpu::PaintParams params) {
     // TODO: perhaps just return the ids here rather than the sk_sps?
-    auto [ combo, uniformData] = ExtractCombo(recorder->uniformCache(), params);
+    auto [ combo, uniformData] = ExtractCombo(params);
     auto programInfo = recorder->programCache()->findOrCreateProgram(combo);
-
-    return { programInfo->id(), uniformData->id() };
+    auto uniformID = recorder->uniformCache()->insert(std::move(uniformData));
+    return { programInfo->id(), uniformID };
 }
 
 } // anonymous namespace
@@ -222,7 +222,7 @@
         // bound independently of those used by the rest of the RenderStep, then we can upload now
         // and remember the location for re-use on any RenderStep that does shading.
         uint32_t programID = ProgramCache::kInvalidProgramID;
-        uint32_t shadingUniformID = UniformData::kInvalidUniformID;
+        uint32_t shadingUniformID = UniformCache::kInvalidUniformID;
         if (draw.fPaintParams.has_value()) {
             std::tie(programID, shadingUniformID) = get_ids_from_paint(recorder,
                                                                        draw.fPaintParams.value());
@@ -238,7 +238,7 @@
             // providing shape, transform, scissor, and paint depth to RenderStep
             uint32_t geometryIndex = 0;
 
-            uint32_t shadingIndex = UniformData::kInvalidUniformID;
+            uint32_t shadingIndex = UniformCache::kInvalidUniformID;
 
             const bool performsShading = draw.fPaintParams.has_value() && step->performsShading();
             if (performsShading) {
@@ -275,8 +275,8 @@
 
     // Used to track when a new pipeline or dynamic state needs recording between draw steps
     uint32_t lastPipeline = 0;
-    uint32_t lastShadingUniforms = UniformData::kInvalidUniformID;
-    uint32_t lastGeometryUniforms = 0;
+    uint32_t lastShadingUniforms = UniformCache::kInvalidUniformID;
+    uint32_t lastGeometryUniforms = UniformCache::kInvalidUniformID;
     SkIRect lastScissor = SkIRect::MakeSize(target->dimensions());
 
     for (const SortKey& key : keys) {
@@ -302,8 +302,8 @@
         if (pipelineChange) {
             // TODO: Look up pipeline description from key's index and record binding it
             lastPipeline = key.pipeline();
-            lastShadingUniforms = UniformData::kInvalidUniformID;
-            lastGeometryUniforms = 0;
+            lastShadingUniforms = UniformCache::kInvalidUniformID;
+            lastGeometryUniforms = UniformCache::kInvalidUniformID;
         }
         if (stateChange) {
             if (key.geometryUniforms() != lastGeometryUniforms) {
diff --git a/experimental/graphite/src/UniformCache.cpp b/experimental/graphite/src/UniformCache.cpp
index ff72d95..87b0b8c 100644
--- a/experimental/graphite/src/UniformCache.cpp
+++ b/experimental/graphite/src/UniformCache.cpp
@@ -12,11 +12,17 @@
 
 namespace skgpu {
 
-size_t UniformCache::Hash::operator()(sk_sp<UniformData> ud) const {
+size_t UniformCache::Hash::operator()(UniformData* ud) const {
+    if (!ud) {
+        return 0;
+    }
     return SkOpts::hash_fn(ud->data(), ud->dataSize(), 0);
 }
 
-bool UniformCache::Eq::operator()(sk_sp<UniformData> a, sk_sp<UniformData> b) const {
+bool UniformCache::Eq::operator()(UniformData* a, UniformData* b) const {
+    if (!a || !b) {
+        return !a && !b;
+    }
     if (a->count() != b->count() ||
         a->uniforms() != b->uniforms() ||
         a->dataSize() != b->dataSize()) {
@@ -28,28 +34,41 @@
 };
 
 UniformCache::UniformCache() {
-    // kInvalidUniformID (aka 0) is reserved
-    fUniformDataVector.push_back(nullptr);
+    // kInvalidUniformID is reserved
+    static_assert(kInvalidUniformID == 0);
+    fUniformData.push_back(nullptr);
+    fUniformDataIDs.insert({nullptr, 0});
 }
 
-sk_sp<UniformData> UniformCache::findOrCreate(sk_sp<UniformData> ud) {
+#ifdef SK_DEBUG
+void UniformCache::validate() const {
+    for (size_t i = 0; i < fUniformData.size(); ++i) {
+        auto kv = fUniformDataIDs.find(fUniformData[i].get());
+        SkASSERT(kv != fUniformDataIDs.end());
+        SkASSERT(kv->first == fUniformData[i].get());
+        SkASSERT(SkTo<uint32_t>(i) == kv->second);
+    }
+}
+#endif
 
-    auto iter = fUniformDataHash.find(ud);
-    if (iter != fUniformDataHash.end()) {
-        SkASSERT((*iter)->id() != UniformData::kInvalidUniformID);
-        return *iter;
+uint32_t UniformCache::insert(sk_sp<UniformData> data) {
+    auto kv = fUniformDataIDs.find(data.get());
+    if (kv != fUniformDataIDs.end()) {
+        return kv->second;
     }
 
-    ud->setID(fNextUniqueID++);
-    fUniformDataHash.insert(ud);
-    fUniformDataVector.push_back(ud);
-    SkASSERT(fUniformDataVector[ud->id()] == ud);
-    return ud;
+    uint32_t id = SkTo<uint32_t>(fUniformData.size());
+    SkASSERT(data && id != kInvalidUniformID);
+
+    fUniformDataIDs.insert({data.get(), id});
+    fUniformData.push_back(std::move(data));
+    this->validate();
+    return id;
 }
 
 sk_sp<UniformData> UniformCache::lookup(uint32_t uniqueID) {
-    SkASSERT(uniqueID < fUniformDataVector.size());
-    return fUniformDataVector[uniqueID];
+    SkASSERT(uniqueID < fUniformData.size());
+    return fUniformData[uniqueID];
 }
 
 } // namespace skgpu
diff --git a/experimental/graphite/src/UniformCache.h b/experimental/graphite/src/UniformCache.h
index f4b52af..6aa4ae9 100644
--- a/experimental/graphite/src/UniformCache.h
+++ b/experimental/graphite/src/UniformCache.h
@@ -8,41 +8,71 @@
 #ifndef skgpu_UniformCache_DEFINED
 #define skgpu_UniformCache_DEFINED
 
-#include <unordered_set>
-#include <vector>
 #include "include/core/SkRefCnt.h"
 
+#include <unordered_map>
+#include <vector>
+
 namespace skgpu {
 
 class UniformData;
 
 class UniformCache {
 public:
+    static constexpr uint32_t kInvalidUniformID = 0;
+
     UniformCache();
 
-    sk_sp<UniformData> findOrCreate(sk_sp<UniformData>);
+    // TODO: Revisit the UniformCache::insert and UniformData::Make APIs:
+    // 1. UniformData::Make requires knowing the data size up front, which involves two invocations
+    //    of the UniformManager. Ideally, we could align uniforms on the fly into a dynamic buffer.
+    // 2. UniformData stores the offsets for each uniform, but these aren't needed after we've
+    //    filled out the buffer. If we remember layout offsets, it should be stored per Combination
+    //    or RenderStep that defines the uniform set.
+    // 3. UniformCache's ids are only fundamentally limited by the number of draws that can be
+    //    recorded into a DrawPass, which means a very large recording with multiple passes could
+    //    exceed uint32_t across all the passes.
+    // 4. The check to know if a UniformData is present in the cache is practically the same for
+    //    checking if the data needs to be uploaded to the GPU, so UniformCache could remember the
+    //    associated BufferBindInfos as well.
+    // 5. Because UniformCache only cares about the content byte hash/equality, and can memcpy to
+    //    the GPU buffer, the cached data contents could all go into a shared byte array, instead of
+    //    needing to extend SkRefCnt.
+    // 6. insert() as a name can imply that the value is always added, so we may want a better one.
+    //    It can be a little less generic if UniformCache returns id and bind buffer info. On the
+    //    other hand unordered_map::insert has the same semantics as this insert, so maybe it's fine
+
+    // Add the block of uniform data to the cache and return a unique ID that corresponds to its
+    // contents. If an identical block of data is already in the cache, that unique ID is returned.
+    uint32_t insert(sk_sp<UniformData>);
 
     sk_sp<UniformData> lookup(uint32_t uniqueID);
 
     // The number of unique uniformdata objects in the cache
     size_t count() const {
-        SkASSERT(fUniformDataHash.size()+1 == fUniformDataVector.size());
-        return fUniformDataHash.size();
+        SkASSERT(fUniformData.size() == fUniformDataIDs.size() && fUniformData.size() > 0);
+        return fUniformData.size() - 1;
     }
 
 private:
     struct Hash {
-        size_t operator()(sk_sp<UniformData>) const;
+        // This hash operator de-references and hashes the data contents
+        size_t operator()(UniformData*) const;
     };
     struct Eq {
-        // This equality operator doesn't compare the UniformData's ids
-        bool operator()(sk_sp<UniformData>, sk_sp<UniformData>) const;
+        // This equality operator de-references and compares the actual data contents
+        bool operator()(UniformData*, UniformData*) const;
     };
 
-    std::unordered_set<sk_sp<UniformData>, Hash, Eq> fUniformDataHash;
-    std::vector<sk_sp<UniformData>> fUniformDataVector;
     // The UniformData's unique ID is only unique w/in a Recorder _not_ globally
-    uint32_t fNextUniqueID = 1;
+    std::unordered_map<UniformData*, uint32_t, Hash, Eq> fUniformDataIDs;
+    std::vector<sk_sp<UniformData>> fUniformData;
+
+#ifdef SK_DEBUG
+    void validate() const;
+#else
+    void validate() const {}
+#endif
 };
 
 } // namespace skgpu
diff --git a/tests/graphite/UniformCacheTest.cpp b/tests/graphite/UniformCacheTest.cpp
index 1ed9506..f787e72 100644
--- a/tests/graphite/UniformCacheTest.cpp
+++ b/tests/graphite/UniformCacheTest.cpp
@@ -47,15 +47,22 @@
 
     REPORTER_ASSERT(reporter, cache->count() == 0);
 
-    // Add a new unique UD
-    sk_sp<UniformData> result1;
+    // Nullptr should already be in the cache and return kInvalidUniformID
     {
-        sk_sp<UniformData> ud1 = make_ud(2, 16);
-        result1 = cache->findOrCreate(ud1);
-        REPORTER_ASSERT(reporter, result1->id() != UniformData::kInvalidUniformID);
-        REPORTER_ASSERT(reporter, ud1->id() == result1->id());
-        sk_sp<UniformData> lookup = cache->lookup(result1->id());
-        REPORTER_ASSERT(reporter, lookup->id() == result1->id());
+        uint32_t result0 = cache->insert(nullptr);
+        REPORTER_ASSERT(reporter, result0 == UniformCache::kInvalidUniformID);
+        REPORTER_ASSERT(reporter, cache->count() == 0);
+    }
+
+    // Add a new unique UD
+    sk_sp<UniformData> ud1;
+    uint32_t result1;
+    {
+        ud1 = make_ud(2, 16);
+        result1 = cache->insert(ud1);
+        REPORTER_ASSERT(reporter, result1 != UniformCache::kInvalidUniformID);
+        sk_sp<UniformData> lookup = cache->lookup(result1);
+        REPORTER_ASSERT(reporter, lookup.get() == ud1.get());
 
         REPORTER_ASSERT(reporter, cache->count() == 1);
     }
@@ -63,12 +70,12 @@
     // Try to add a duplicate UD
     {
         sk_sp<UniformData> ud2 = make_ud(2, 16);
-        sk_sp<UniformData> result2 = cache->findOrCreate(ud2);
-        REPORTER_ASSERT(reporter, result2->id() != UniformData::kInvalidUniformID);
-        REPORTER_ASSERT(reporter, ud2->id() == UniformData::kInvalidUniformID);
-        REPORTER_ASSERT(reporter, result2->id() == result1->id());
-        sk_sp<UniformData> lookup = cache->lookup(result2->id());
-        REPORTER_ASSERT(reporter, lookup->id() == result2->id());
+        uint32_t result2 = cache->insert(ud2);
+        REPORTER_ASSERT(reporter, result2 != UniformCache::kInvalidUniformID);
+        REPORTER_ASSERT(reporter, result2 == result1);
+        sk_sp<UniformData> lookup = cache->lookup(result2);
+        REPORTER_ASSERT(reporter, lookup.get() != ud2.get());
+        REPORTER_ASSERT(reporter, lookup.get() == ud1.get());
 
         REPORTER_ASSERT(reporter, cache->count() == 1);
     }
@@ -76,12 +83,12 @@
     // Add a second new unique UD
     {
         sk_sp<UniformData> ud3 = make_ud(3, 16);
-        sk_sp<UniformData> result3 = cache->findOrCreate(ud3);
-        REPORTER_ASSERT(reporter, result3->id() != UniformData::kInvalidUniformID);
-        REPORTER_ASSERT(reporter, ud3->id() == result3->id());
-        REPORTER_ASSERT(reporter, result3->id() != result1->id());
-        sk_sp<UniformData> lookup = cache->lookup(result3->id());
-        REPORTER_ASSERT(reporter, lookup->id() == result3->id());
+        uint32_t result3 = cache->insert(ud3);
+        REPORTER_ASSERT(reporter, result3 != UniformCache::kInvalidUniformID);
+        REPORTER_ASSERT(reporter, result3 != result1);
+        sk_sp<UniformData> lookup = cache->lookup(result3);
+        REPORTER_ASSERT(reporter, lookup.get() == ud3.get());
+        REPORTER_ASSERT(reporter, lookup.get() != ud1.get());
 
         REPORTER_ASSERT(reporter, cache->count() == 2);
     }
diff --git a/tests/graphite/UniformTest.cpp b/tests/graphite/UniformTest.cpp
index 96a3b43..d566a20 100644
--- a/tests/graphite/UniformTest.cpp
+++ b/tests/graphite/UniformTest.cpp
@@ -9,7 +9,6 @@
 
 #include "experimental/graphite/src/ContextUtils.h"
 #include "experimental/graphite/src/DrawList.h" // TODO: split PaintParams out into their own header
-#include "experimental/graphite/src/UniformCache.h"
 #include "include/core/SkPaint.h"
 #include "include/effects/SkGradientShader.h"
 
@@ -62,8 +61,6 @@
 DEF_GRAPHITE_TEST(UniformTest, reporter) {
     using namespace skgpu;
 
-    UniformCache cache;
-
     // Intentionally does not include ShaderType::kNone, which represents no fragment shading stage
     // and is thus not relevant to uniform extraction/caching.
     for (auto s : { ShaderCombo::ShaderType::kSolidColor,
@@ -87,13 +84,12 @@
                 expected.fBlendMode = bm;
 
                 auto [ p, expectedNumUniforms ] = create_paint(expected);
-                auto [ actual, ud] = ExtractCombo(&cache, PaintParams(p));
+                auto [ actual, ud] = ExtractCombo(PaintParams(p));
                 REPORTER_ASSERT(reporter, expected == actual);
                 REPORTER_ASSERT(reporter, expectedNumUniforms == ud->count());
                 for (int i = 0; i < ud->count(); ++i) {
                     REPORTER_ASSERT(reporter, ud->offset(i) >= 0 && ud->offset(i) < ud->dataSize());
                 }
-                REPORTER_ASSERT(reporter, ud->id() != UniformData::kInvalidUniformID);
             }
         }
     }