Extract most of the mutable state of SkShader into a separate Context object.

src/core/SkDraw.cpp

Index: src/core/SkDraw.cpp
diff --git a/src/core/SkDraw.cpp b/src/core/SkDraw.cpp
index 05b06453dbf200df8cebccfc3aee4c69825d5143..957c19639d18ba9c75e7a63f70089acd7071787c 100644
--- a/src/core/SkDraw.cpp
+++ b/src/core/SkDraw.cpp
@@ -2355,13 +2355,33 @@ static bool texture_to_matrix(const VertState& state, const SkPoint verts[],
     return matrix->setPolyToPoly(src, dst, 3);
 }
 
+/**
+ *  This class is NOT immutable due to the setup() method being called repeatedly after
+ *  the shader has been associated with a paint. This is okay because we only ever
+ *  associate this shader with temporary paints.
+ */
 class SkTriColorShader : public SkShader {
 public:
     SkTriColorShader() {}
 
-    bool setup(const SkPoint pts[], const SkColor colors[], int, int, int);
+    // TODO(dominikg): Feels like this should be part of the context, but we don't have
+    //                 access to it where we need it.
+    bool setup(const SkPoint pts[], const SkColor colors[], int, int, int, const SkMatrix&);
+
+    virtual SkShader::Context* createContext(
+            const SkBitmap&, const SkPaint&, const SkMatrix&, void*) const SK_OVERRIDE;
+    virtual size_t contextSize() const SK_OVERRIDE;
+
+    class TriColorShaderContext : public SkShader::Context {
+    public:
+        TriColorShaderContext(const SkTriColorShader& shader, const SkBitmap& device,
+                              const SkPaint& paint, const SkMatrix& matrix);
+        virtual ~TriColorShaderContext();
+
+        virtual void shadeSpan(int x, int y, SkPMColor dstC[], int count) SK_OVERRIDE;
 
-    virtual void shadeSpan(int x, int y, SkPMColor dstC[], int count) SK_OVERRIDE;
+        typedef SkShader::Context INHERITED;
+    };
 
     SK_DEVELOPER_TO_STRING()
     SK_DECLARE_PUBLIC_FLATTENABLE_DESERIALIZATION_PROCS(SkTriColorShader)
@@ -2376,8 +2396,17 @@ private:
     typedef SkShader INHERITED;
 };
 
+SkShader::Context* SkTriColorShader::createContext(const SkBitmap& device, const SkPaint& paint,
+                                                   const SkMatrix& matrix, void* storage) const {
+    if (!this->validContext(device, paint, matrix)) {
+        return NULL;
+    }
+
+    return SkNEW_PLACEMENT_ARGS(storage, TriColorShaderContext, (*this, device, paint, matrix));
+}
+
 bool SkTriColorShader::setup(const SkPoint pts[], const SkColor colors[],
-                             int index0, int index1, int index2) {
+                             int index0, int index1, int index2, const SkMatrix& matrix) {
 
     fColors[0] = SkPreMultiplyColor(colors[index0]);
     fColors[1] = SkPreMultiplyColor(colors[index1]);
@@ -2394,7 +2423,9 @@ bool SkTriColorShader::setup(const SkPoint pts[], const SkColor colors[],
     if (!m.invert(&im)) {
         return false;
     }
-    return fDstToUnit.setConcat(im, this->getTotalInverse());
+    SkMatrix inverse;
+    SkASSERT(matrix.invert(&inverse));
+    return fDstToUnit.setConcat(im, inverse);
 }
 
 #include "SkColorPriv.h"
@@ -2411,11 +2442,24 @@ static int ScalarTo256(SkScalar v) {
     return SkAlpha255To256(scale);
 }
 
-void SkTriColorShader::shadeSpan(int x, int y, SkPMColor dstC[], int count) {
+
+SkTriColorShader::TriColorShaderContext::TriColorShaderContext(
+        const SkTriColorShader& shader, const SkBitmap& device,
+        const SkPaint& paint, const SkMatrix& matrix)
+    : INHERITED(shader, device, paint, matrix) {}
+
+SkTriColorShader::TriColorShaderContext::~TriColorShaderContext() {}
+
+size_t SkTriColorShader::contextSize() const {
+    return sizeof(TriColorShaderContext);
+}
+void SkTriColorShader::TriColorShaderContext::shadeSpan(int x, int y, SkPMColor dstC[], int count) {
+    const SkTriColorShader& triColorShader = static_cast<const SkTriColorShader&>(fShader);
+
     SkPoint src;
 
     for (int i = 0; i < count; i++) {
-        fDstToUnit.mapXY(SkIntToScalar(x), SkIntToScalar(y), &src);
+        triColorShader.fDstToUnit.mapXY(SkIntToScalar(x), SkIntToScalar(y), &src);
         x += 1;
 
         int scale1 = ScalarTo256(src.fX);
@@ -2430,9 +2474,9 @@ void SkTriColorShader::shadeSpan(int x, int y, SkPMColor dstC[], int count) {
             scale0 = 0;
         }
 
-        dstC[i] = SkAlphaMulQ(fColors[0], scale0) +
-        SkAlphaMulQ(fColors[1], scale1) +
-        SkAlphaMulQ(fColors[2], scale2);
+        dstC[i] = SkAlphaMulQ(triColorShader.fColors[0], scale0) +
+        SkAlphaMulQ(triColorShader.fColors[1], scale1) +
+        SkAlphaMulQ(triColorShader.fColors[2], scale2);
     }
 }
 
@@ -2518,8 +2562,8 @@ void SkDraw::drawVertices(SkCanvas::VertexMode vmode, int count,
     }
 
     SkAutoBlitterChoose blitter(*fBitmap, *fMatrix, p);
-    // important that we abort early, as below we may manipulate the shader
-    // and that is only valid if the shader returned true from setContext.
+    // Important that we abort early, as below we may manipulate the shader context
+    // and that is only valid if a context was successfully created from the shader.
     // If it returned false, then our blitter will be the NullBlitter.
     if (blitter->isNullBlitter()) {
         return;
@@ -2536,30 +2580,19 @@ void SkDraw::drawVertices(SkCanvas::VertexMode vmode, int count,
             savedLocalM = shader->getLocalMatrix();
         }
 
-        // setContext has already been called and verified to return true
-        // by the constructor of SkAutoBlitterChoose
-        bool prevContextSuccess = true;
         while (vertProc(&state)) {
             if (NULL != textures) {
                 if (texture_to_matrix(state, vertices, textures, &tempM)) {
                     tempM.postConcat(savedLocalM);
                     shader->setLocalMatrix(tempM);
-                    // Need to recall setContext since we changed the local matrix.
-                    // However, we also need to balance the calls this with a
-                    // call to endContext which requires tracking the result of
-                    // the previous call to setContext.
-                    if (prevContextSuccess) {
-                        shader->endContext();
-                    }
-                    prevContextSuccess = shader->setContext(*fBitmap, p, *fMatrix);
-                    if (!prevContextSuccess) {
+                    if (!blitter->resetShaderContext(*fBitmap, p, *fMatrix)) {
                         continue;
                     }
                 }
             }
             if (NULL != colors) {
                 if (!triShader.setup(vertices, colors,
-                                     state.f0, state.f1, state.f2)) {
+                                     state.f0, state.f1, state.f2, *fMatrix)) {
                     continue;
                 }
             }
@@ -2574,13 +2607,6 @@ void SkDraw::drawVertices(SkCanvas::VertexMode vmode, int count,
         if (NULL != shader) {
             shader->setLocalMatrix(savedLocalM);
         }
-
-        // If the final call to setContext fails we must make it suceed so that the
-        // call to endContext in the destructor for SkAutoBlitterChoose is balanced.
-        if (!prevContextSuccess) {
-            prevContextSuccess = shader->setContext(*fBitmap, paint, SkMatrix::I());
-            SkASSERT(prevContextSuccess);
-        }
     } else {
         // no colors[] and no texture
         HairProc hairProc = ChooseHairProc(paint.isAntiAlias());