Add removeUserDefinedSnippet method to ShaderCodeDictionary.

We need to purge out user snippets when they are no longer needed, to
avoid unbounded memory usage.

Change-Id: Ib33909ec8c94cd6272f1a28e52a7ab92b27dfb0d
Bug: skia:13405
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/552577
Reviewed-by: Michael Ludwig <michaelludwig@google.com>
Auto-Submit: John Stiles <johnstiles@google.com>
Commit-Queue: Michael Ludwig <michaelludwig@google.com>
diff --git a/src/core/SkShaderCodeDictionary.cpp b/src/core/SkShaderCodeDictionary.cpp
index 9257b4a..ba5b089 100644
--- a/src/core/SkShaderCodeDictionary.cpp
+++ b/src/core/SkShaderCodeDictionary.cpp
@@ -633,8 +633,12 @@
         return true;
     }
 
-    int userDefinedCodeSnippetID = snippetID - kBuiltInCodeSnippetIDCount;
-    return userDefinedCodeSnippetID < SkTo<int>(fUserDefinedCodeSnippets.size());
+    int index = snippetID - kBuiltInCodeSnippetIDCount;
+    if (index >= SkTo<int>(fUserDefinedCodeSnippets.size())) {
+        return false;
+    }
+
+    return fUserDefinedCodeSnippets[index] != nullptr;
 }
 
 static constexpr int kNoChildren = 0;
@@ -688,6 +692,20 @@
     return SkBlenderID(kBuiltInCodeSnippetIDCount + fUserDefinedCodeSnippets.size() - 1);
 }
 
+void SkShaderCodeDictionary::removeUserDefinedSnippet(int codeSnippetID) {
+    SkASSERT(codeSnippetID >= kBuiltInCodeSnippetIDCount);
+    SkASSERT(this->isValidID(codeSnippetID));
+
+    int index = codeSnippetID - kBuiltInCodeSnippetIDCount;
+    fUserDefinedCodeSnippets[index] = nullptr;
+
+    // Reclaim unused code snippet IDs at the end of the snippet list.
+    // (For now, we don't make any attempt to reclaim any gaps in the middle.)
+    while (fUserDefinedCodeSnippets.back() == nullptr) {
+        fUserDefinedCodeSnippets.pop_back();
+    }
+}
+
 SkShaderCodeDictionary::SkShaderCodeDictionary() {
     // The 0th index is reserved as invalid
     fEntryVector.push_back(nullptr);
diff --git a/src/core/SkShaderCodeDictionary.h b/src/core/SkShaderCodeDictionary.h
index 5c61ce8..2075502 100644
--- a/src/core/SkShaderCodeDictionary.h
+++ b/src/core/SkShaderCodeDictionary.h
@@ -206,6 +206,8 @@
 
     SkBlenderID addUserDefinedBlender(sk_sp<SkRuntimeEffect>);
 
+    void removeUserDefinedSnippet(int codeSnippetID);
+
 private:
 #ifdef SK_GRAPHITE_ENABLED
     Entry* makeEntry(const SkPaintParamsKey&, const skgpu::BlendInfo&);
diff --git a/tests/graphite/KeyTest.cpp b/tests/graphite/KeyTest.cpp
index 88c1b25..2c552a2 100644
--- a/tests/graphite/KeyTest.cpp
+++ b/tests/graphite/KeyTest.cpp
@@ -72,6 +72,8 @@
     };
 
     int userSnippetID = dict->addUserDefinedSnippet("keyAlmostTooBig", kDataFields);
+    REPORTER_ASSERT(reporter, dict->isValidID(userSnippetID));
+
     SkPaintParamsKey key = create_key(&builder, userSnippetID, kMaxBlockDataSize);
 
     // Key is created successfully.
@@ -90,13 +92,41 @@
             {"data", SkPaintParamsKey::DataPayloadType::kByte, kBlockDataSize},
     };
 
+    // Snippet creation succeeds
     int userSnippetID = dict->addUserDefinedSnippet("keyTooBig", kDataFields);
+    REPORTER_ASSERT(reporter, dict->isValidID(userSnippetID));
+
     SkPaintParamsKey key = create_key(&builder, userSnippetID, kBlockDataSize);
 
     // Key creation fails.
     REPORTER_ASSERT(reporter, key.isErrorKey());
 }
 
+DEF_GRAPHITE_TEST_FOR_CONTEXTS(CanRemoveUserDefinedSnippet, reporter, context) {
+
+    SkShaderCodeDictionary* dict = context->priv().shaderCodeDictionary();
+    SkPaintParamsKeyBuilder builder(dict, SkBackend::kGraphite);
+
+    static constexpr int kDataSize = 1;
+    static constexpr SkPaintParamsKey::DataPayloadField kDataFields[] = {
+            {"Data", SkPaintParamsKey::DataPayloadType::kByte, kDataSize},
+    };
+
+    // After adding a user defined snippet, it has a valid ID and can be used to create keys.
+    int userSnippetID = dict->addUserDefinedSnippet("userSnippet", kDataFields);
+    REPORTER_ASSERT(reporter, dict->isValidID(userSnippetID));
+
+    SkPaintParamsKey key = create_key(&builder, userSnippetID, kDataSize);
+    REPORTER_ASSERT(reporter, !key.isErrorKey());
+
+    // After removing the snippet, its ID becomes invalid and it can't create keys anymore.
+    dict->removeUserDefinedSnippet(userSnippetID);
+    REPORTER_ASSERT(reporter, !dict->isValidID(userSnippetID));
+
+    SkPaintParamsKey key2 = create_key(&builder, userSnippetID, kDataSize);
+    REPORTER_ASSERT(reporter, key2.isErrorKey());
+}
+
 DEF_GRAPHITE_TEST_FOR_CONTEXTS(KeyEqualityChecksSnippetID, reporter, context) {
 
     SkShaderCodeDictionary* dict = context->priv().shaderCodeDictionary();