Add IO for TPL stats

Overload TempOutFile constructor to allow IO mode.

Bug: b/281563704

Change-Id: I1f4f5b29db0e331941b6795e478eeeab51f625ad

7 files changed, 212 insertions, 8 deletions

diff --git a/libs.mk b/libs.mk
index ea5cc15a1..f6f6cc94c 100644
--- a/libs.mk
+++ b/libs.mk
@@ -178,7 +178,7 @@ INSTALL-LIBS-yes += include/vpx/vpx_image.h
 INSTALL-LIBS-yes += include/vpx/vpx_integer.h
 INSTALL-LIBS-$(CONFIG_DECODERS) += include/vpx/vpx_decoder.h
 INSTALL-LIBS-$(CONFIG_ENCODERS) += include/vpx/vpx_encoder.h
-INSTALL-LIBS-$(CONFIG_VP9_ENCODER) += include/vpx/vpx_tpl.h
+INSTALL-LIBS-$(CONFIG_ENCODERS) += include/vpx/vpx_tpl.h
 ifeq ($(CONFIG_EXTERNAL_BUILD),yes)
 ifeq ($(CONFIG_MSVS),yes)
 INSTALL-LIBS-yes                  += $(foreach p,$(VS_PLATFORMS),$(LIBSUBDIR)/$(p)/$(CODEC_LIB).lib)
diff --git a/test/encode_api_test.cc b/test/encode_api_test.cc
index af98ad5dd..e8a044ae1 100644
--- a/test/encode_api_test.cc
+++ b/test/encode_api_test.cc
@@ -11,6 +11,7 @@
 #include <climits>
 #include <cstring>
 #include <initializer_list>
+#include <new>
 
 #include "third_party/googletest/src/include/gtest/gtest.h"
 #include "test/codec_factory.h"
@@ -368,7 +369,7 @@ class EncodeApiGetTplStatsTest
     : public ::libvpx_test::EncoderTest,
       public ::testing::TestWithParam<const libvpx_test::CodecFactory *> {
  public:
-  EncodeApiGetTplStatsTest() : EncoderTest(GetParam()) {}
+  EncodeApiGetTplStatsTest() : EncoderTest(GetParam()), test_io_(false) {}
   ~EncodeApiGetTplStatsTest() override {}
 
  protected:
@@ -396,6 +397,34 @@ class EncodeApiGetTplStatsTest
     return VPX_CODEC_OK;
   }
 
+  void CompareTplGopStats(const VpxTplGopStats &ref_gop_stats,
+                          const VpxTplGopStats &test_gop_stats) {
+    ASSERT_EQ(ref_gop_stats.size, test_gop_stats.size);
+    for (int frame = 0; frame < ref_gop_stats.size; frame++) {
+      const VpxTplFrameStats &ref_frame_stats =
+          ref_gop_stats.frame_stats_list[frame];
+      const VpxTplFrameStats &test_frame_stats =
+          test_gop_stats.frame_stats_list[frame];
+      ASSERT_EQ(ref_frame_stats.num_blocks, test_frame_stats.num_blocks);
+      ASSERT_EQ(ref_frame_stats.frame_width, test_frame_stats.frame_width);
+      ASSERT_EQ(ref_frame_stats.frame_height, test_frame_stats.frame_height);
+      for (int block = 0; block < ref_frame_stats.num_blocks; block++) {
+        const VpxTplBlockStats &ref_block_stats =
+            ref_frame_stats.block_stats_list[block];
+        const VpxTplBlockStats &test_block_stats =
+            test_frame_stats.block_stats_list[block];
+        ASSERT_EQ(ref_block_stats.inter_cost, test_block_stats.inter_cost);
+        ASSERT_EQ(ref_block_stats.intra_cost, test_block_stats.intra_cost);
+        ASSERT_EQ(ref_block_stats.mv_c, test_block_stats.mv_c);
+        ASSERT_EQ(ref_block_stats.mv_r, test_block_stats.mv_r);
+        ASSERT_EQ(ref_block_stats.recrf_dist, test_block_stats.recrf_dist);
+        ASSERT_EQ(ref_block_stats.recrf_rate, test_block_stats.recrf_rate);
+        ASSERT_EQ(ref_block_stats.ref_frame_index,
+                  test_block_stats.ref_frame_index);
+      }
+    }
+  }
+
   void PostEncodeFrameHook(::libvpx_test::Encoder *encoder) override {
     ::libvpx_test::CxDataIterator iter = encoder->GetCxData();
     while (const vpx_codec_cx_pkt_t *pkt = iter.Next()) {
@@ -416,7 +445,21 @@ class EncodeApiGetTplStatsTest
             }
           }
           ASSERT_TRUE(stats_not_all_zero);
-          // Free the memory right away now as this is only a test.
+          if (test_io_ && tpl_stats.size > 0) {
+            libvpx_test::TempOutFile *temp_out_file =
+                new (std::nothrow) libvpx_test::TempOutFile("w+");
+            ASSERT_NE(temp_out_file, nullptr);
+            ASSERT_NE(temp_out_file->file(), nullptr);
+            vpx_write_tpl_gop_stats(temp_out_file->file(), &tpl_stats);
+            rewind(temp_out_file->file());
+            VpxTplGopStats gop_stats_io;
+            ASSERT_EQ(
+                vpx_read_tpl_gop_stats(temp_out_file->file(), &gop_stats_io),
+                VPX_CODEC_OK);
+            CompareTplGopStats(gop_stats_io, tpl_stats);
+            vpx_free_tpl_gop_stats(&gop_stats_io);
+            delete temp_out_file;
+          }
           free(tpl_stats.frame_stats_list);
           break;
         }
@@ -427,6 +470,7 @@ class EncodeApiGetTplStatsTest
 
   int width_;
   int height_;
+  bool test_io_;
 };
 
 TEST_P(EncodeApiGetTplStatsTest, GetTplStats) {
@@ -438,6 +482,16 @@ TEST_P(EncodeApiGetTplStatsTest, GetTplStats) {
   ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
 }
 
+TEST_P(EncodeApiGetTplStatsTest, GetTplStatsIO) {
+  cfg_.g_lag_in_frames = 25;
+  width_ = 352;
+  height_ = 288;
+  test_io_ = true;
+  ::libvpx_test::I420VideoSource video("hantro_collage_w352h288.yuv", width_,
+                                       height_, 30, 1, 0, 50);
+  ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     VP9, EncodeApiGetTplStatsTest,
     ::testing::Values(
diff --git a/test/video_source.h b/test/video_source.h
index a10ff6fb0..5ed99d063 100644
--- a/test/video_source.h
+++ b/test/video_source.h
@@ -64,7 +64,7 @@ inline FILE *OpenTestDataFile(const std::string &file_name) {
   return fopen(path_to_source.c_str(), "rb");
 }
 
-static FILE *GetTempOutFile(std::string *file_name) {
+static FILE *GetTempOutFile(std::string *file_name, const char *io_mode) {
   file_name->clear();
 #if defined(_WIN32)
   char fname[MAX_PATH];
@@ -73,7 +73,7 @@ static FILE *GetTempOutFile(std::string *file_name) {
     // Assume for now that the filename generated is unique per process
     if (GetTempFileNameA(tmppath, "lvx", 0, fname)) {
       file_name->assign(fname);
-      return fopen(fname, "wb+");
+      return fopen(fname, io_mode);
     }
   }
   return nullptr;
@@ -94,13 +94,16 @@ static FILE *GetTempOutFile(std::string *file_name) {
   const int fd = mkstemp(temp_file_name.get());
   if (fd == -1) return nullptr;
   *file_name = temp_file_name.get();
-  return fdopen(fd, "wb+");
+  return fdopen(fd, io_mode);
 #endif
 }
 
 class TempOutFile {
  public:
-  TempOutFile() { file_ = GetTempOutFile(&file_name_); }
+  TempOutFile() { file_ = GetTempOutFile(&file_name_, "wb+"); }
+  TempOutFile(const char *io_mode) {
+    file_ = GetTempOutFile(&file_name_, io_mode);
+  }
   ~TempOutFile() {
     CloseFile();
     if (!file_name_.empty()) {
diff --git a/vpx/exports_com b/vpx/exports_com
index 2ab05099f..f0b46aa17 100644
--- a/vpx/exports_com
+++ b/vpx/exports_com
@@ -14,3 +14,6 @@ text vpx_img_flip
 text vpx_img_free
 text vpx_img_set_rect
 text vpx_img_wrap
+text vpx_free_tpl_gop_stats
+text vpx_read_tpl_gop_stats
+text vpx_write_tpl_gop_stats
diff --git a/vpx/src/vpx_tpl.c b/vpx/src/vpx_tpl.c
new file mode 100644
index 000000000..9cdb4a0a0
--- /dev/null
+++ b/vpx/src/vpx_tpl.c
@@ -0,0 +1,107 @@
+/*
+ *  Copyright (c) 2023 The WebM project authors. All Rights Reserved.
+ *
+ *  Use of this source code is governed by a BSD-style license
+ *  that can be found in the LICENSE file in the root of the source
+ *  tree. An additional intellectual property rights grant can be found
+ *  in the file PATENTS.  All contributing project authors may
+ *  be found in the AUTHORS file in the root of the source tree.
+ */
+
+#include <stdlib.h>
+
+#include "vpx/vpx_codec.h"
+#include "vpx/vpx_tpl.h"
+#include "vpx_mem/vpx_mem.h"
+
+#define CHECK_FPRINTF_ERROR(expr) \
+  do {                            \
+    if (expr < 0) {               \
+      return VPX_CODEC_ERROR;     \
+    }                             \
+  } while (0)
+
+#define CHECK_FSCANF_ERROR(expr, expected_value) \
+  do {                                           \
+    if (expr != expected_value) {                \
+      return VPX_CODEC_ERROR;                    \
+    }                                            \
+  } while (0)
+
+vpx_codec_err_t vpx_write_tpl_gop_stats(FILE *tpl_file,
+                                        const VpxTplGopStats *tpl_gop_stats) {
+  int i;
+  if (tpl_file == NULL || tpl_gop_stats == NULL) return VPX_CODEC_INVALID_PARAM;
+  CHECK_FPRINTF_ERROR(fprintf(tpl_file, "%d\n", tpl_gop_stats->size));
+
+  for (i = 0; i < tpl_gop_stats->size; i++) {
+    VpxTplFrameStats frame_stats = tpl_gop_stats->frame_stats_list[i];
+    const int num_blocks = frame_stats.num_blocks;
+    int block;
+    CHECK_FPRINTF_ERROR(fprintf(tpl_file, "%d %d %d\n", frame_stats.frame_width,
+                                frame_stats.frame_height, num_blocks));
+    for (block = 0; block < num_blocks; block++) {
+      VpxTplBlockStats block_stats = frame_stats.block_stats_list[block];
+      CHECK_FPRINTF_ERROR(
+          fprintf(tpl_file,
+                  "%" PRId64 " %" PRId64 " %" PRId16 " %" PRId16 " %" PRId64
+                  " %" PRId64 " %d\n",
+                  block_stats.inter_cost, block_stats.intra_cost,
+                  block_stats.mv_c, block_stats.mv_r, block_stats.recrf_dist,
+                  block_stats.recrf_rate, block_stats.ref_frame_index));
+    }
+  }
+
+  return VPX_CODEC_OK;
+}
+
+vpx_codec_err_t vpx_read_tpl_gop_stats(FILE *tpl_file,
+                                       VpxTplGopStats *tpl_gop_stats) {
+  int i, frame_list_size;
+  if (tpl_file == NULL || tpl_gop_stats == NULL) return VPX_CODEC_INVALID_PARAM;
+  CHECK_FSCANF_ERROR(fscanf(tpl_file, "%d\n", &frame_list_size), 1);
+  tpl_gop_stats->size = frame_list_size;
+  tpl_gop_stats->frame_stats_list = (VpxTplFrameStats *)vpx_calloc(
+      frame_list_size, sizeof(tpl_gop_stats->frame_stats_list[0]));
+  if (tpl_gop_stats->frame_stats_list == NULL) {
+    return VPX_CODEC_MEM_ERROR;
+  }
+  for (i = 0; i < frame_list_size; i++) {
+    VpxTplFrameStats *frame_stats = &tpl_gop_stats->frame_stats_list[i];
+    int num_blocks, width, height, block;
+    CHECK_FSCANF_ERROR(
+        fscanf(tpl_file, "%d %d %d\n", &width, &height, &num_blocks), 3);
+    frame_stats->num_blocks = num_blocks;
+    frame_stats->frame_width = width;
+    frame_stats->frame_height = height;
+    frame_stats->block_stats_list = (VpxTplBlockStats *)vpx_calloc(
+        num_blocks, sizeof(frame_stats->block_stats_list[0]));
+    if (frame_stats->block_stats_list == NULL) {
+      vpx_free_tpl_gop_stats(tpl_gop_stats);
+      return VPX_CODEC_MEM_ERROR;
+    }
+    for (block = 0; block < num_blocks; block++) {
+      VpxTplBlockStats *block_stats = &frame_stats->block_stats_list[block];
+      CHECK_FSCANF_ERROR(
+          fscanf(tpl_file,
+                 "%" SCNd64 " %" SCNd64 " %" SCNd16 " %" SCNd16 " %" SCNd64
+                 " %" SCNd64 " %d\n",
+                 &block_stats->inter_cost, &block_stats->intra_cost,
+                 &block_stats->mv_c, &block_stats->mv_r,
+                 &block_stats->recrf_dist, &block_stats->recrf_rate,
+                 &block_stats->ref_frame_index),
+          7);
+    }
+  }
+
+  return VPX_CODEC_OK;
+}
+
+void vpx_free_tpl_gop_stats(VpxTplGopStats *data) {
+  int frame;
+  if (data == NULL) return;
+  for (frame = 0; frame < data->size; frame++) {
+    vpx_free(data->frame_stats_list[frame].block_stats_list);
+  }
+  vpx_free(data->frame_stats_list);
+}
diff --git a/vpx/vpx_codec.mk b/vpx/vpx_codec.mk
index 4aec88b30..25c815ef5 100644
--- a/vpx/vpx_codec.mk
+++ b/vpx/vpx_codec.mk
@@ -37,6 +37,7 @@ API_SRCS-yes += internal/vpx_codec_internal.h
 API_SRCS-yes += internal/vpx_ratectrl_rtc.h
 API_SRCS-yes += src/vpx_codec.c
 API_SRCS-yes += src/vpx_image.c
+API_SRCS-yes += src/vpx_tpl.c
 API_SRCS-yes += vpx_codec.h
 API_SRCS-yes += vpx_codec.mk
 API_SRCS-yes += vpx_frame_buffer.h
diff --git a/vpx/vpx_tpl.h b/vpx/vpx_tpl.h
index 689fa9692..50aec49eb 100644
--- a/vpx/vpx_tpl.h
+++ b/vpx/vpx_tpl.h
@@ -15,6 +15,8 @@
 #ifndef VPX_VPX_VPX_TPL_H_
 #define VPX_VPX_VPX_TPL_H_
 
+#include <stdio.h>
+
 #include "./vpx_integer.h"
 
 #ifdef __cplusplus
@@ -29,7 +31,7 @@ extern "C" {
  * types, removing or reassigning enums, adding/removing/rearranging
  * fields to structures
  */
-#define VPX_TPL_ABI_VERSION (0) /**<\hideinitializer*/
+#define VPX_TPL_ABI_VERSION (1) /**<\hideinitializer*/
 
 /*!\brief Temporal dependency model stats for each block before propagation */
 typedef struct VpxTplBlockStats {
@@ -56,6 +58,40 @@ typedef struct VpxTplGopStats {
   VpxTplFrameStats *frame_stats_list; /**< List of tpl stats for each frame */
 } VpxTplGopStats;
 
+/*!\brief Write VpxTplGopStats to file
+ *
+ * Accepts an opened file handle and writes \p tpl_gop_stats.
+ *
+ * \param[in]    tpl_file       A FILE pointer that's already been opened.
+ * \param[in]    tpl_gop_stats  VpxTplGopStats that contains TPL stats for the
+ *                              whole GOP.
+ *
+ * \return VPX_CODEC_OK if TPL stats are successfully written.
+ */
+vpx_codec_err_t vpx_write_tpl_gop_stats(FILE *tpl_file,
+                                        const VpxTplGopStats *tpl_gop_stats);
+
+/*!\brief Read VpxTplGopStats from file
+ *
+ * Accepts an opened file handle and reads TPL stats and stores them into
+ * \p tpl_gop_stats. Allocates memory for TPL stats.
+ *
+ * \param[in]     tpl_file       A FILE pointer that's already been opened.
+ * \param[out]    tpl_gop_stats  VpxTplGopStats that contains TPL stats for the
+ *                               whole GOP.
+ *
+ * \return VPX_CODEC_OK if TPL stats are successfully read from file.
+ */
+vpx_codec_err_t vpx_read_tpl_gop_stats(FILE *tpl_file,
+                                       VpxTplGopStats *tpl_gop_stats);
+
+/*!\brief Free the memory allocated for VpxTplGopStats
+ *
+ * \param[in]    tpl_gop_stats  VpxTplGopStats that contains TPL stats for the
+ *                              whole GOP.
+ */
+void vpx_free_tpl_gop_stats(VpxTplGopStats *tpl_gop_stats);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif