1 files changed, 37 insertions, 69 deletions

diff --git a/third_party/jpeg-xl/lib/jxl/enc_linalg_test.cc b/third_party/jpeg-xl/lib/jxl/enc_linalg_test.cc
index 967b9a3afb..c02f009ca7 100644
--- a/third_party/jpeg-xl/lib/jxl/enc_linalg_test.cc
+++ b/third_party/jpeg-xl/lib/jxl/enc_linalg_test.cc
@@ -5,110 +5,78 @@
 
 #include "lib/jxl/enc_linalg.h"
 
-#include "lib/jxl/image_test_utils.h"
+#include "lib/jxl/base/random.h"
 #include "lib/jxl/testing.h"
 
 namespace jxl {
 namespace {
 
-ImageD Identity(const size_t N) {
-  ImageD out(N, N);
-  for (size_t i = 0; i < N; ++i) {
-    double* JXL_RESTRICT row = out.Row(i);
-    std::fill(row, row + N, 0);
-    row[i] = 1.0;
-  }
-  return out;
-}
+Matrix2x2 Diagonal(const Vector2& d) { return {{{d[0], 0.0}, {0.0, d[1]}}}; }
 
-ImageD Diagonal(const ImageD& d) {
-  JXL_ASSERT(d.ysize() == 1);
-  ImageD out(d.xsize(), d.xsize());
-  const double* JXL_RESTRICT row_diag = d.Row(0);
-  for (size_t k = 0; k < d.xsize(); ++k) {
-    double* JXL_RESTRICT row_out = out.Row(k);
-    std::fill(row_out, row_out + d.xsize(), 0.0);
-    row_out[k] = row_diag[k];
-  }
-  return out;
-}
+Matrix2x2 Identity() { return Diagonal({1.0, 1.0}); }
 
-ImageD MatMul(const ImageD& A, const ImageD& B) {
-  JXL_ASSERT(A.ysize() == B.xsize());
-  ImageD out(A.xsize(), B.ysize());
-  for (size_t y = 0; y < B.ysize(); ++y) {
-    const double* const JXL_RESTRICT row_b = B.Row(y);
-    double* const JXL_RESTRICT row_out = out.Row(y);
-    for (size_t x = 0; x < A.xsize(); ++x) {
-      row_out[x] = 0.0;
-      for (size_t k = 0; k < B.xsize(); ++k) {
-        row_out[x] += A.Row(k)[x] * row_b[k];
-      }
+Matrix2x2 MatMul(const Matrix2x2& A, const Matrix2x2& B) {
+  Matrix2x2 out;
+  for (size_t y = 0; y < 2; ++y) {
+    for (size_t x = 0; x < 2; ++x) {
+      out[y][x] = A[0][x] * B[y][0] + A[1][x] * B[y][1];
     }
   }
   return out;
 }
 
-ImageD Transpose(const ImageD& A) {
-  ImageD out(A.ysize(), A.xsize());
-  for (size_t x = 0; x < A.xsize(); ++x) {
-    double* const JXL_RESTRICT row_out = out.Row(x);
-    for (size_t y = 0; y < A.ysize(); ++y) {
-      row_out[y] = A.Row(y)[x];
-    }
-  }
-  return out;
+Matrix2x2 Transpose(const Matrix2x2& A) {
+  return {{{A[0][0], A[1][0]}, {A[0][1], A[1][1]}}};
 }
 
-ImageD RandomSymmetricMatrix(const size_t N, Rng& rng, const double vmin,
-                             const double vmax) {
-  ImageD A(N, N);
-  GenerateImage(rng, &A, vmin, vmax);
-  for (size_t i = 0; i < N; ++i) {
-    for (size_t j = 0; j < i; ++j) {
-      A.Row(j)[i] = A.Row(i)[j];
-    }
-  }
+Matrix2x2 RandomSymmetricMatrix(Rng& rng, const double vmin,
+                                const double vmax) {
+  Matrix2x2 A;
+  A[0][0] = rng.UniformF(vmin, vmax);
+  A[0][1] = A[1][0] = rng.UniformF(vmin, vmax);
+  A[1][1] = rng.UniformF(vmin, vmax);
   return A;
 }
 
-void VerifyMatrixEqual(const ImageD& A, const ImageD& B, const double eps) {
-  ASSERT_EQ(A.xsize(), B.xsize());
-  ASSERT_EQ(A.ysize(), B.ysize());
-  for (size_t y = 0; y < A.ysize(); ++y) {
-    for (size_t x = 0; x < A.xsize(); ++x) {
-      ASSERT_NEAR(A.Row(y)[x], B.Row(y)[x], eps);
+void VerifyMatrixEqual(const Matrix2x2& A, const Matrix2x2& B,
+                       const double eps) {
+  for (size_t y = 0; y < 2; ++y) {
+    for (size_t x = 0; x < 2; ++x) {
+      ASSERT_NEAR(A[y][x], B[y][x], eps);
     }
   }
 }
 
-void VerifyOrthogonal(const ImageD& A, const double eps) {
-  VerifyMatrixEqual(Identity(A.xsize()), MatMul(Transpose(A), A), eps);
+void VerifyOrthogonal(const Matrix2x2& A, const double eps) {
+  VerifyMatrixEqual(Identity(), MatMul(Transpose(A), A), eps);
 }
 
 TEST(LinAlgTest, ConvertToDiagonal) {
   {
-    ImageD I = Identity(2);
-    ImageD U(2, 2), d(2, 1);
-    ConvertToDiagonal(I, &d, &U);
+    Matrix2x2 I = Identity();
+    Matrix2x2 U;
+    Vector2 d;
+    ConvertToDiagonal(I, d, U);
     VerifyMatrixEqual(I, U, 1e-15);
     for (size_t k = 0; k < 2; ++k) {
-      ASSERT_NEAR(d.Row(0)[k], 1.0, 1e-15);
+      ASSERT_NEAR(d[k], 1.0, 1e-15);
     }
   }
   {
-    ImageD A = Identity(2);
-    A.Row(0)[1] = A.Row(1)[0] = 2.0;
-    ImageD U(2, 2), d(2, 1);
-    ConvertToDiagonal(A, &d, &U);
+    Matrix2x2 A = Identity();
+    A[0][1] = A[1][0] = 2.0;
+    Matrix2x2 U;
+    Vector2 d;
+    ConvertToDiagonal(A, d, U);
     VerifyOrthogonal(U, 1e-12);
     VerifyMatrixEqual(A, MatMul(U, MatMul(Diagonal(d), Transpose(U))), 1e-12);
   }
   Rng rng(0);
   for (size_t i = 0; i < 100; ++i) {
-    ImageD A = RandomSymmetricMatrix(2, rng, -1.0, 1.0);
-    ImageD U(2, 2), d(2, 1);
-    ConvertToDiagonal(A, &d, &U);
+    Matrix2x2 A = RandomSymmetricMatrix(rng, -1.0, 1.0);
+    Matrix2x2 U;
+    Vector2 d;
+    ConvertToDiagonal(A, d, U);
     VerifyOrthogonal(U, 1e-12);
     VerifyMatrixEqual(A, MatMul(U, MatMul(Diagonal(d), Transpose(U))), 1e-12);
   }