JPEG/R: stride support

Bug: 264715926
Test: jpegr_test.cpp, manual test with camera/display app
Change-Id: I92c1d19b8bd49060794b4a9427c67017aa6743a9
diff --git a/libs/jpegrecoverymap/recoverymapmath.cpp b/libs/jpegrecoverymap/recoverymapmath.cpp
index 2cffde3..7c0c15a 100644
--- a/libs/jpegrecoverymap/recoverymapmath.cpp
+++ b/libs/jpegrecoverymap/recoverymapmath.cpp
@@ -493,17 +493,28 @@
 }
 
 Color getP010Pixel(jr_uncompressed_ptr image, size_t x, size_t y) {
-  size_t pixel_count = image->width * image->height;
+  size_t luma_stride = image->luma_stride;
+  size_t chroma_stride = image->chroma_stride;
+  uint16_t* luma_data = reinterpret_cast<uint16_t*>(image->data);
+  uint16_t* chroma_data = reinterpret_cast<uint16_t*>(image->chroma_data);
 
-  size_t pixel_y_idx = x + y * image->width;
-  size_t pixel_uv_idx = x / 2 + (y / 2) * (image->width / 2);
+  if (luma_stride == 0) {
+    luma_stride = image->width;
+  }
+  if (chroma_stride == 0) {
+    chroma_stride = luma_stride;
+  }
+  if (chroma_data == nullptr) {
+    chroma_data = &reinterpret_cast<uint16_t*>(image->data)[image->luma_stride * image->height];
+  }
 
-  uint16_t y_uint = reinterpret_cast<uint16_t*>(image->data)[pixel_y_idx]
-                  >> 6;
-  uint16_t u_uint = reinterpret_cast<uint16_t*>(image->data)[pixel_count + pixel_uv_idx * 2]
-                  >> 6;
-  uint16_t v_uint = reinterpret_cast<uint16_t*>(image->data)[pixel_count + pixel_uv_idx * 2 + 1]
-                  >> 6;
+  size_t pixel_y_idx = y * luma_stride + x;
+  size_t pixel_u_idx = (y >> 1) * chroma_stride + (x & ~0x1);
+  size_t pixel_v_idx = pixel_u_idx + 1;
+
+  uint16_t y_uint = luma_data[pixel_y_idx] >> 6;
+  uint16_t u_uint = chroma_data[pixel_u_idx] >> 6;
+  uint16_t v_uint = chroma_data[pixel_v_idx] >> 6;
 
   // Conversions include taking narrow-range into account.
   return {{{ (static_cast<float>(y_uint) - 64.0f) / 876.0f,