Fix particle bug where uniforms are allocated too late

Also adds an example particle system that uses uniforms, and an instance
of the particles GM that draws it.

Change-Id: I64e2514fd17c8ce615b4e13b9f82424f80b8424e
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/356840
Commit-Queue: Brian Osman <brianosman@google.com>
Reviewed-by: John Stiles <johnstiles@google.com>

gm/particles.cpp

@@ -14,12 +14,27 @@
 #include "modules/particles/include/SkParticleEffect.h"
 #include "modules/particles/include/SkParticleSerialization.h"
 #include "modules/skresources/include/SkResources.h"
+#include "src/sksl/SkSLVMGenerator.h"
 #include "tools/Resources.h"
 
+struct UniformValue {
+    const char*        fName;
+    std::vector<float> fData;
+};
+
 class ParticlesGM : public skiagm::GM {
 public:
-    ParticlesGM(const char* name, double startTime, SkISize size, SkPoint origin)
-            : GM(SK_ColorBLACK), fName(name), fStartTime(startTime), fSize(size), fOrigin(origin) {}
+    ParticlesGM(const char* name,
+                double startTime,
+                SkISize size,
+                SkPoint origin,
+                std::vector<UniformValue> uniforms = {})
+            : GM(SK_ColorBLACK)
+            , fName(name)
+            , fStartTime(startTime)
+            , fSize(size)
+            , fOrigin(origin)
+            , fUniforms(std::move(uniforms)) {}
 
     SkString onShortName() override { return SkStringPrintf("particles_%s", fName); }
     SkISize onISize() override { return fSize; }
@@ -37,6 +52,10 @@ public:
         effectParams->prepare(resourceProvider.get());
 
         fEffect = sk_make_sp<SkParticleEffect>(effectParams);
+        for (const auto& val : fUniforms) {
+            SkAssertResult(fEffect->setUniform(val.fName, val.fData.data(), val.fData.size()));
+        }
+
         fEffect->start(/*now=*/0.0, /*looping=*/true);
 
         // Fast-forward (in 30 fps time-slices) to the requested time
@@ -60,11 +79,12 @@ public:
     }
 
 protected:
-    const char*             fName;
-    const double            fStartTime;
-    const SkISize           fSize;
-    const SkPoint           fOrigin;
-    sk_sp<SkParticleEffect> fEffect;
+    const char*               fName;
+    const double              fStartTime;
+    const SkISize             fSize;
+    const SkPoint             fOrigin;
+    sk_sp<SkParticleEffect>   fEffect;
+    std::vector<UniformValue> fUniforms;
 };
 
 DEF_GM(return new ParticlesGM("confetti",     1.0, {400, 400}, {200, 200});)
@@ -74,5 +94,9 @@ DEF_GM(return new ParticlesGM("mandrill",     1.0, {250, 250}, { 25,  25});)
 DEF_GM(return new ParticlesGM("spiral",       2.0, {250, 250}, {125, 125});)
 DEF_GM(return new ParticlesGM("sprite_frame", 1.0, {200, 200}, {100, 100});)
 DEF_GM(return new ParticlesGM("text",         1.0, {250, 110}, { 10, 100});)
+DEF_GM(return new ParticlesGM("uniforms",     2.0, {250, 250}, {125, 125},
+                                              {{"rate",  {2.0f}},
+                                               {"spin",  {4.0f}},
+                                               {"color", {0.25f, 0.75f, 0.75f}}});)
 
 #endif  // SK_BUILD_FOR_GOOGLE3

modules/particles/include/SkParticleEffect.h

@@ -175,10 +175,16 @@ public:
     const SkSL::UniformInfo* uniformInfo() const;
     float* uniformData() { return fUniforms.data(); }
 
+    // Sets named uniform to the data in 'val'. 'count' must be equal to the total number of floats
+    // in the uniform (eg, the number of elements in a vector). Returns false if the uniform isn't
+    // found, or if count is incorrect. Returns true if the value is changed successfully.
+    bool setUniform(const char* name, const float* val, int count);
+
     static void RegisterParticleTypes();
 
 private:
     void setCapacity(int capacity);
+    void updateStorage();
 
     // Helpers to break down update
     void advanceTime(double now);
@@ -223,7 +229,7 @@ private:
     SkAutoTMalloc<float> fStableRandoms;
 
     // Cached
-    int fCapacity;
+    int fCapacity = 0;
     SkTArray<float, true> fUniforms;
 
     friend struct SkParticleProgram;

modules/particles/src/SkParticleEffect.cpp

@@ -220,7 +220,43 @@ SkParticleEffect::SkParticleEffect(sk_sp<SkParticleEffectParams> params)
         , fLastTime(-1.0)
         , fSpawnRemainder(0.0f) {
     fState.fAge = -1.0f;
-    this->setCapacity(fParams->fMaxCount);
+    this->updateStorage();
+}
+
+void SkParticleEffect::updateStorage() {
+    // Handle user edits to fMaxCount
+    if (fParams->fMaxCount != fCapacity) {
+        this->setCapacity(fParams->fMaxCount);
+    }
+
+    // Ensure our storage block for uniforms is large enough
+    if (this->uniformInfo()) {
+        int newCount = this->uniformInfo()->fUniformSlotCount;
+        if (newCount > fUniforms.count()) {
+            fUniforms.push_back_n(newCount - fUniforms.count(), 0.0f);
+        } else {
+            fUniforms.resize(newCount);
+        }
+    }
+}
+
+bool SkParticleEffect::setUniform(const char* name, const float* val, int count) {
+    const SkSL::UniformInfo* info = this->uniformInfo();
+    if (!info) {
+        return false;
+    }
+
+    auto it = std::find_if(info->fUniforms.begin(), info->fUniforms.end(),
+                           [name](const auto& u) { return u.fName == name; });
+    if (it == info->fUniforms.end()) {
+        return false;
+    }
+    if (it->fRows * it->fColumns != count) {
+        return false;
+    }
+
+    std::copy(val, val + count, this->uniformData() + it->fSlot);
+    return true;
 }
 
 void SkParticleEffect::start(double now, bool looping, SkPoint position, SkVector heading,
@@ -317,23 +353,8 @@ void SkParticleEffect::advanceTime(double now) {
     }
     fLastTime = now;
 
-    // Handle user edits to fMaxCount
-    if (fParams->fMaxCount != fCapacity) {
-        this->setCapacity(fParams->fMaxCount);
-    }
-
-    // Ensure our storage block for uniforms are large enough
-    auto resizeWithZero = [](SkTArray<float, true>* uniforms, const SkSL::UniformInfo* info) {
-        if (info) {
-            int newCount = info->fUniformSlotCount;
-            if (newCount > uniforms->count()) {
-                uniforms->push_back_n(newCount - uniforms->count(), 0.0f);
-            } else {
-                uniforms->resize(newCount);
-            }
-        }
-    };
-    resizeWithZero(&fUniforms, this->uniformInfo());
+    // Possibly re-allocate cached storage, if our params have changed
+    this->updateStorage();
 
     // Copy known values into the uniform blocks
     if (fParams->fProgram) {

resources/particles/uniforms.json

@@ -0,0 +1,32 @@
+{
+   "MaxCount": 800,
+   "Drawable": {
+      "Type": "SkCircleDrawable",
+      "Radius": 2
+   },
+   "Code": [
+      "uniform float rate;",
+      "uniform float spin;",
+      "uniform float3 color;",
+      "",
+      "void effectSpawn(inout Effect effect) {",
+      "  effect.lifetime = 4;",
+      "}",
+      "",
+      "void effectUpdate(inout Effect effect) {",
+      "  effect.rate = 100 * rate;",
+      "  effect.spin = spin;",
+      "}",
+      "",
+      "void spawn(inout Particle p) {",
+      "  p.lifetime = 2 + rand(p.seed);",
+      "  p.vel = p.dir * mix(50, 60, rand(p.seed));",
+      "}",
+      "",
+      "void update(inout Particle p) {",
+      "  p.color.rgb = color;",
+      "}",
+      ""
+   ],
+   "Bindings": []
+}