Widen internal API to support more complex YUV formats

Bug: skia:7901
Change-Id: I46fec08711b8b483cf58ccae733e4dc2a9689231
Reviewed-on: https://skia-review.googlesource.com/c/162280
Commit-Queue: Jim Van Verth <jvanverth@google.com>
Reviewed-by: Brian Salomon <bsalomon@google.com>
diff --git a/tests/ImageGeneratorTest.cpp b/tests/ImageGeneratorTest.cpp
index 5c8864e..907efc0 100644
--- a/tests/ImageGeneratorTest.cpp
+++ b/tests/ImageGeneratorTest.cpp
@@ -7,9 +7,11 @@
 
 #include "SkData.h"
 #include "SkCanvas.h"
+
 #include "SkGraphics.h"
 #include "SkImageGenerator.h"
 #include "SkImageInfoPriv.h"
+#include "SkYUVAIndex.h"
 #include "Test.h"
 
 #if defined(SK_BUILD_FOR_MAC) || defined(SK_BUILD_FOR_IOS)
@@ -61,24 +63,32 @@
 DEF_TEST(ImageGenerator, reporter) {
     MyImageGenerator ig;
     SkYUVSizeInfo sizeInfo;
-    sizeInfo.fSizes[SkYUVSizeInfo::kY] = SkISize::Make(200, 200);
-    sizeInfo.fSizes[SkYUVSizeInfo::kU] = SkISize::Make(100, 100);
-    sizeInfo.fSizes[SkYUVSizeInfo::kV] = SkISize::Make( 50,  50);
-    sizeInfo.fWidthBytes[SkYUVSizeInfo::kY] = 0;
-    sizeInfo.fWidthBytes[SkYUVSizeInfo::kU] = 0;
-    sizeInfo.fWidthBytes[SkYUVSizeInfo::kV] = 0;
-    void* planes[3] = { nullptr };
+    sizeInfo.fSizes[0] = SkISize::Make(200, 200);
+    sizeInfo.fSizes[1] = SkISize::Make(100, 100);
+    sizeInfo.fSizes[2] = SkISize::Make( 50,  50);
+    sizeInfo.fSizes[3] = SkISize::Make( 25,  25);
+    sizeInfo.fWidthBytes[0] = 0;
+    sizeInfo.fWidthBytes[1] = 0;
+    sizeInfo.fWidthBytes[2] = 0;
+    sizeInfo.fWidthBytes[3] = 0;
+    void* planes[4] = { nullptr };
+    SkYUVAIndex yuvaIndices[4];
     SkYUVColorSpace colorSpace;
 
     // Check that the YUV decoding API does not cause any crashes
-    ig.queryYUV8(&sizeInfo, nullptr);
-    ig.queryYUV8(&sizeInfo, &colorSpace);
-    sizeInfo.fWidthBytes[SkYUVSizeInfo::kY] = 250;
-    sizeInfo.fWidthBytes[SkYUVSizeInfo::kU] = 250;
-    sizeInfo.fWidthBytes[SkYUVSizeInfo::kV] = 250;
+    ig.queryYUVA8(&sizeInfo, yuvaIndices, nullptr);
+    ig.queryYUVA8(&sizeInfo, yuvaIndices, &colorSpace);
+    sizeInfo.fWidthBytes[0] = 250;
+    sizeInfo.fWidthBytes[1] = 250;
+    sizeInfo.fWidthBytes[2] = 250;
+    sizeInfo.fWidthBytes[3] = 250;
+    yuvaIndices[0] = { 0, SkColorChannel::kR };
+    yuvaIndices[1] = { 1, SkColorChannel::kR };
+    yuvaIndices[2] = { 2, SkColorChannel::kR };
+    yuvaIndices[3] = { 3, SkColorChannel::kR };
     int dummy;
-    planes[SkYUVSizeInfo::kY] = planes[SkYUVSizeInfo::kU] = planes[SkYUVSizeInfo::kV] = &dummy;
-    ig.getYUV8Planes(sizeInfo, planes);
+    planes[0] = planes[1] = planes[2] = planes[3] = &dummy;
+    ig.getYUVA8Planes(sizeInfo, yuvaIndices, planes);
 
     // Suppressed due to https://code.google.com/p/skia/issues/detail?id=4339
     if (false) {
diff --git a/tests/YUVCacheTest.cpp b/tests/YUVCacheTest.cpp
index b34cf06..0410fdb 100644
--- a/tests/YUVCacheTest.cpp
+++ b/tests/YUVCacheTest.cpp
@@ -32,10 +32,16 @@
     SkResourceCache cache(1024);
 
     SkYUVPlanesCache::Info yuvInfo;
-    for (int i = 0; i < 3; i++) {
-        yuvInfo.fSizeInfo.fSizes[i].fWidth = 20 * i;
-        yuvInfo.fSizeInfo.fSizes[i].fHeight = 10 * i;
-        yuvInfo.fSizeInfo.fWidthBytes[i] = 80 * i;
+    for (int i = 0; i < SkYUVSizeInfo::kMaxCount; i++) {
+        yuvInfo.fSizeInfo.fColorTypes[i] = kAlpha_8_SkColorType;
+        yuvInfo.fSizeInfo.fSizes[i].fWidth = 20 * (i + 1);
+        yuvInfo.fSizeInfo.fSizes[i].fHeight = 10 * (i + 1);
+        yuvInfo.fSizeInfo.fWidthBytes[i] = 80 * (i + 1);
+    }
+
+    for (int i = 0; i < SkYUVAIndex::kIndexCount; ++i) {
+        yuvInfo.fYUVAIndices[i].fIndex = -1;
+        yuvInfo.fYUVAIndices[i].fChannel = SkColorChannel::kR;
     }
     yuvInfo.fColorSpace = kRec601_SkYUVColorSpace;
 
@@ -59,14 +65,12 @@
 
     REPORTER_ASSERT(reporter, data);
     REPORTER_ASSERT(reporter, data->size() == size);
-    for (int i = 0; i < 3; ++i) {
-        REPORTER_ASSERT(reporter, yuvInfo.fSizeInfo.fSizes[i].fWidth ==
-                yuvInfoRead.fSizeInfo.fSizes[i].fWidth);
-        REPORTER_ASSERT(reporter, yuvInfo.fSizeInfo.fSizes[i].fHeight ==
-                yuvInfoRead.fSizeInfo.fSizes[i].fHeight);
-        REPORTER_ASSERT(reporter, yuvInfo.fSizeInfo.fWidthBytes[i] ==
-                yuvInfoRead.fSizeInfo.fWidthBytes[i]);
+    REPORTER_ASSERT(reporter, yuvInfo.fSizeInfo == yuvInfoRead.fSizeInfo);
+
+    for (int i = 0; i < SkYUVAIndex::kIndexCount; ++i) {
+        REPORTER_ASSERT(reporter, yuvInfo.fYUVAIndices[i] == yuvInfoRead.fYUVAIndices[i]);
     }
+
     REPORTER_ASSERT(reporter, yuvInfo.fColorSpace == yuvInfoRead.fColorSpace);
 
     check_data(reporter, data, 2, kInCache, kLocked);
diff --git a/tests/YUVTest.cpp b/tests/YUVTest.cpp
index 3862952..50e36fd 100644
--- a/tests/YUVTest.cpp
+++ b/tests/YUVTest.cpp
@@ -14,8 +14,8 @@
 #include "Test.h"
 
 static void codec_yuv(skiatest::Reporter* reporter,
-                  const char path[],
-                  SkISize expectedSizes[3]) {
+                      const char path[],
+                      SkISize expectedSizes[4]) {
     std::unique_ptr<SkStream> stream(GetResourceAsStream(path));
     if (!stream) {
         return;
@@ -28,59 +28,59 @@
 
     // Test queryYUV8()
     SkYUVSizeInfo info;
-    bool success = codec->queryYUV8(nullptr, nullptr);
-    REPORTER_ASSERT(reporter, !success);
-    success = codec->queryYUV8(&info, nullptr);
-    REPORTER_ASSERT(reporter, (expectedSizes == nullptr) == !success);
-    if (!success) {
-        return;
+
+    {
+        bool success = codec->queryYUV8(nullptr, nullptr);
+        REPORTER_ASSERT(reporter, !success);
+        success = codec->queryYUV8(&info, nullptr);
+        REPORTER_ASSERT(reporter, (expectedSizes == nullptr) == !success);
+        if (!success) {
+            return;
+        }
+
+        for (int i = 0; i < SkYUVSizeInfo::kMaxCount; ++i) {
+            REPORTER_ASSERT(reporter, info.fSizes[i] == expectedSizes[i]);
+            REPORTER_ASSERT(reporter,
+                            info.fWidthBytes[i] == (uint32_t) SkAlign8(info.fSizes[i].width()));
+        }
     }
-    REPORTER_ASSERT(reporter,
-            0 == memcmp((const void*) &info, (const void*) expectedSizes, 3 * sizeof(SkISize)));
-    REPORTER_ASSERT(reporter, info.fWidthBytes[SkYUVSizeInfo::kY] ==
-            (uint32_t) SkAlign8(info.fSizes[SkYUVSizeInfo::kY].width()));
-    REPORTER_ASSERT(reporter, info.fWidthBytes[SkYUVSizeInfo::kU] ==
-            (uint32_t) SkAlign8(info.fSizes[SkYUVSizeInfo::kU].width()));
-    REPORTER_ASSERT(reporter, info.fWidthBytes[SkYUVSizeInfo::kV] ==
-            (uint32_t) SkAlign8(info.fSizes[SkYUVSizeInfo::kV].width()));
-    SkYUVColorSpace colorSpace;
-    success = codec->queryYUV8(&info, &colorSpace);
-    REPORTER_ASSERT(reporter,
-            0 == memcmp((const void*) &info, (const void*) expectedSizes, 3 * sizeof(SkISize)));
-    REPORTER_ASSERT(reporter, info.fWidthBytes[SkYUVSizeInfo::kY] ==
-            (uint32_t) SkAlign8(info.fSizes[SkYUVSizeInfo::kY].width()));
-    REPORTER_ASSERT(reporter, info.fWidthBytes[SkYUVSizeInfo::kU] ==
-            (uint32_t) SkAlign8(info.fSizes[SkYUVSizeInfo::kU].width()));
-    REPORTER_ASSERT(reporter, info.fWidthBytes[SkYUVSizeInfo::kV] ==
-            (uint32_t) SkAlign8(info.fSizes[SkYUVSizeInfo::kV].width()));
-    REPORTER_ASSERT(reporter, kJPEG_SkYUVColorSpace == colorSpace);
+
+    {
+        SkYUVColorSpace colorSpace;
+        bool success = codec->queryYUV8(&info, &colorSpace);
+        REPORTER_ASSERT(reporter, (expectedSizes == nullptr) == !success);
+        if (!success) {
+            return;
+        }
+
+        for (int i = 0; i < SkYUVSizeInfo::kMaxCount; ++i) {
+            REPORTER_ASSERT(reporter, info.fSizes[i] == expectedSizes[i]);
+            REPORTER_ASSERT(reporter,
+                            info.fWidthBytes[i] == (uint32_t) SkAlign8(info.fSizes[i].width()));
+        }
+        REPORTER_ASSERT(reporter, kJPEG_SkYUVColorSpace == colorSpace);
+    }
 
     // Allocate the memory for the YUV decode
-    size_t totalBytes =
-            info.fWidthBytes[SkYUVSizeInfo::kY] * info.fSizes[SkYUVSizeInfo::kY].height() +
-            info.fWidthBytes[SkYUVSizeInfo::kU] * info.fSizes[SkYUVSizeInfo::kU].height() +
-            info.fWidthBytes[SkYUVSizeInfo::kV] * info.fSizes[SkYUVSizeInfo::kV].height();
+    size_t totalBytes = info.computeTotalBytes();
+
     SkAutoMalloc storage(totalBytes);
-    void* planes[3];
-    planes[0] = storage.get();
-    planes[1] = SkTAddOffset<void>(planes[0],
-            info.fWidthBytes[SkYUVSizeInfo::kY] * info.fSizes[SkYUVSizeInfo::kY].height());
-    planes[2] = SkTAddOffset<void>(planes[1],
-            info.fWidthBytes[SkYUVSizeInfo::kU] * info.fSizes[SkYUVSizeInfo::kU].height());
+    void* planes[SkYUVSizeInfo::kMaxCount];
+
+    info.computePlanes(storage.get(), planes);
 
     // Test getYUV8Planes()
-    REPORTER_ASSERT(reporter, SkCodec::kInvalidInput ==
-            codec->getYUV8Planes(info, nullptr));
-    REPORTER_ASSERT(reporter, SkCodec::kSuccess ==
-            codec->getYUV8Planes(info, planes));
+    REPORTER_ASSERT(reporter, SkCodec::kInvalidInput == codec->getYUV8Planes(info, nullptr));
+    REPORTER_ASSERT(reporter, SkCodec::kSuccess == codec->getYUV8Planes(info, planes));
 }
 
 DEF_TEST(Jpeg_YUV_Codec, r) {
-    SkISize sizes[3];
+    SkISize sizes[4];
 
     sizes[0].set(128, 128);
     sizes[1].set(64, 64);
     sizes[2].set(64, 64);
+    sizes[3].set(0, 0);
     codec_yuv(r, "images/color_wheel.jpg", sizes);
 
     // H2V2