YUVUtils function for splitting image into yuva planes.

Makes the code from yuv_splitter reusable and able to produce subsampled
planes.

Bug: chromium:1210557
Change-Id: Icce112658bbdb866c3ecb9dcff1a5e8d0d30135a
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/411297
Reviewed-by: Michael Ludwig <michaelludwig@google.com>
Commit-Queue: Brian Salomon <bsalomon@google.com>
diff --git a/gm/wacky_yuv_formats.cpp b/gm/wacky_yuv_formats.cpp
index d87b628..c8ea3f5 100644
--- a/gm/wacky_yuv_formats.cpp
+++ b/gm/wacky_yuv_formats.cpp
@@ -1153,29 +1153,6 @@
 #include "src/core/SkAutoPixmapStorage.h"
 #include "tools/Resources.h"
 
-static void draw_into_alpha(const SkImage* img, sk_sp<SkColorFilter> cf, const SkPixmap& dst) {
-    auto canvas = SkCanvas::MakeRasterDirect(dst.info(), dst.writable_addr(), dst.rowBytes());
-    canvas->scale(1.0f * dst.width() / img->width(), 1.0f * dst.height() / img->height());
-    SkPaint paint;
-    paint.setColorFilter(cf);
-    paint.setBlendMode(SkBlendMode::kSrc);
-    canvas->drawImage(img, 0, 0, SkSamplingOptions(SkFilterMode::kLinear), &paint);
-}
-
-static void split_into_yuv(const SkImage* img, SkYUVColorSpace cs, const SkPixmap dst[3]) {
-    float m[20];
-    SkColorMatrix_RGB2YUV(cs, m);
-
-    memcpy(m + 15, m + 0, 5 * sizeof(float));   // copy Y into A
-    draw_into_alpha(img, SkColorFilters::Matrix(m), dst[0]);
-
-    memcpy(m + 15, m + 5, 5 * sizeof(float));   // copy U into A
-    draw_into_alpha(img, SkColorFilters::Matrix(m), dst[1]);
-
-    memcpy(m + 15, m + 10, 5 * sizeof(float));   // copy V into A
-    draw_into_alpha(img, SkColorFilters::Matrix(m), dst[2]);
-}
-
 static void draw_diff(SkCanvas* canvas, SkScalar x, SkScalar y,
                       const SkImage* a, const SkImage* b) {
     auto sh = SkShaders::Blend(SkBlendMode::kDifference,
@@ -1200,9 +1177,7 @@
 // resulting (recombined) images (gpu only for now).
 //
 class YUVSplitterGM : public skiagm::GM {
-    sk_sp<SkImage>      fOrig;
-    SkAutoPixmapStorage fStorage[3];
-    SkPixmap            fPM[3];
+    sk_sp<SkImage> fOrig;
 
 public:
     YUVSplitterGM() {}
@@ -1219,27 +1194,26 @@
 
     void onOnceBeforeDraw() override {
         fOrig = GetResourceAsImage("images/mandrill_256.png");
-
-        SkImageInfo info = SkImageInfo::MakeA8(fOrig->dimensions());
-        fStorage[0].alloc(info);
-        fStorage[1].alloc(info);
-        fStorage[2].alloc(info);
-        for (int i = 0; i < 3; ++i) {
-            fPM[i] = fStorage[i];
-        }
     }
 
     void onDraw(SkCanvas* canvas) override {
         canvas->translate(fOrig->width(), 0);
         canvas->save();
-        for (auto cs : {kRec709_SkYUVColorSpace, kRec601_SkYUVColorSpace, kJPEG_SkYUVColorSpace,
+        SkYUVAInfo info;
+        std::array<sk_sp<SkImage>, SkYUVAInfo::kMaxPlanes> planes;
+        for (auto cs : {kRec709_SkYUVColorSpace,
+                        kRec601_SkYUVColorSpace,
+                        kJPEG_SkYUVColorSpace,
                         kBT2020_SkYUVColorSpace}) {
-            split_into_yuv(fOrig.get(), cs, fPM);
-            SkYUVAInfo yuvaInfo(fOrig->dimensions(),
-                                SkYUVAInfo::PlaneConfig::kY_U_V,
-                                SkYUVAInfo::Subsampling::k444,
-                                cs);
-            auto yuvaPixmaps = SkYUVAPixmaps::FromExternalPixmaps(yuvaInfo, fPM);
+            std::tie(planes, info) = sk_gpu_test::MakeYUVAPlanesAsA8(fOrig.get(),
+                                                                     cs,
+                                                                     SkYUVAInfo::Subsampling::k444,
+                                                                     /*recording context*/ nullptr);
+            SkPixmap pixmaps[4];
+            for (int i = 0; i < info.numPlanes(); ++i) {
+                planes[i]->peekPixels(&pixmaps[i]);
+            }
+            auto yuvaPixmaps = SkYUVAPixmaps::FromExternalPixmaps(info, pixmaps);
             auto img = SkImage::MakeFromYUVAPixmaps(canvas->recordingContext(),
                                                     yuvaPixmaps,
                                                     GrMipMapped::kNo,
@@ -1253,11 +1227,11 @@
         }
         canvas->restore();
         canvas->translate(-fOrig->width(), 0);
-
-        canvas->drawImage(SkImage::MakeRasterCopy(fPM[0]), 0, 0);
-        canvas->drawImage(SkImage::MakeRasterCopy(fPM[1]), 0, fPM[0].height());
-        canvas->drawImage(SkImage::MakeRasterCopy(fPM[2]),
-                          0, fPM[0].height() + fPM[1].height());
+        int y = 0;
+        for (int i = 0; i < info.numPlanes(); ++i) {
+            canvas->drawImage(planes[i], 0, y);
+            y += planes[i]->height();
+        }
     }
 
 private:
diff --git a/tools/gpu/YUVUtils.cpp b/tools/gpu/YUVUtils.cpp
index 9344998..b74a9b0 100644
--- a/tools/gpu/YUVUtils.cpp
+++ b/tools/gpu/YUVUtils.cpp
@@ -7,8 +7,11 @@
 
 #include "tools/gpu/YUVUtils.h"
 
+#include "include/core/SkCanvas.h"
+#include "include/core/SkColorFilter.h"
 #include "include/core/SkColorPriv.h"
 #include "include/core/SkData.h"
+#include "include/core/SkSurface.h"
 #include "include/gpu/GrRecordingContext.h"
 #include "include/gpu/GrYUVABackendTextures.h"
 #include "src/codec/SkCodecImageGenerator.h"
@@ -134,6 +137,55 @@
 
 namespace sk_gpu_test {
 
+std::tuple<std::array<sk_sp<SkImage>, SkYUVAInfo::kMaxPlanes>, SkYUVAInfo>
+MakeYUVAPlanesAsA8(SkImage* src,
+                   SkYUVColorSpace cs,
+                   SkYUVAInfo::Subsampling ss,
+                   GrRecordingContext* rContext) {
+    float rgbToYUV[20];
+    SkColorMatrix_RGB2YUV(cs, rgbToYUV);
+
+    SkYUVAInfo::PlaneConfig config = src->isOpaque() ? SkYUVAInfo::PlaneConfig::kY_U_V
+                                                     : SkYUVAInfo::PlaneConfig::kY_U_V_A;
+    SkISize dims[SkYUVAInfo::kMaxPlanes];
+    int n = SkYUVAInfo::PlaneDimensions(src->dimensions(),
+                                        config,
+                                        ss,
+                                        kTopLeft_SkEncodedOrigin,
+                                        dims);
+    std::array<sk_sp<SkImage>, 4> planes;
+    for (int i = 0; i < n; ++i) {
+        SkImageInfo info = SkImageInfo::MakeA8(dims[i]);
+        sk_sp<SkSurface> surf;
+        if (rContext) {
+            surf = SkSurface::MakeRenderTarget(rContext, SkBudgeted::kYes, info, 1, nullptr);
+        } else {
+            surf = SkSurface::MakeRaster(info);
+        }
+        if (!surf) {
+            return {};
+        }
+
+        SkPaint paint;
+        paint.setBlendMode(SkBlendMode::kSrc);
+
+        // Make a matrix with the ith row of rgbToYUV copied to the A row since we're drawing to A8.
+        float m[20] = {};
+        std::copy_n(rgbToYUV + 5*i, 5, m + 15);
+        paint.setColorFilter(SkColorFilters::Matrix(m));
+        surf->getCanvas()->drawImageRect(src,
+                                         SkRect::Make(dims[i]),
+                                         SkSamplingOptions(SkFilterMode::kLinear),
+                                         &paint);
+        planes[i] = surf->makeImageSnapshot();
+        if (!planes[i]) {
+            return {};
+        }
+    }
+    SkYUVAInfo info(src->dimensions(), config, ss, cs);
+    return {planes, info};
+}
+
 std::unique_ptr<LazyYUVImage> LazyYUVImage::Make(sk_sp<SkData> data,
                                                  GrMipmapped mipmapped,
                                                  sk_sp<SkColorSpace> cs) {
diff --git a/tools/gpu/YUVUtils.h b/tools/gpu/YUVUtils.h
index 1a1ebbf..7762a62 100644
--- a/tools/gpu/YUVUtils.h
+++ b/tools/gpu/YUVUtils.h
@@ -13,10 +13,22 @@
 #include "include/gpu/GrBackendSurface.h"
 #include "src/core/SkAutoMalloc.h"
 
+#include <tuple>
+
 class SkData;
 
 namespace sk_gpu_test {
 
+// Splits an input image into A8 YUV[A] planes using the passed subsampling and YUV color space. If
+// the src image is opaque there will be three planes (Y, U, and V) and if not there will be a
+// fourth A plane. The planes are returned along with a SkYUVAInfo describing the resulting planar
+// image. Images are made as textures if GrRecordingContext is not null, otherwise as cpu images.
+std::tuple<std::array<sk_sp<SkImage>, SkYUVAInfo::kMaxPlanes>, SkYUVAInfo>
+MakeYUVAPlanesAsA8(SkImage*,
+                   SkYUVColorSpace,
+                   SkYUVAInfo::Subsampling,
+                   GrRecordingContext*);
+
 // Utility that decodes a JPEG but preserves the YUVA8 planes in the image, and uses
 // MakeFromYUVAPixmaps to create a GPU multiplane YUVA image for a context. It extracts the planar
 // data once, and lazily creates the actual SkImage when the GrContext is provided (and refreshes