From 2e5261647f0440e78081f9bfaaaa5ef9d10140d6 Mon Sep 17 00:00:00 2001
From: Jerome Jiang <jianj@google.com>
Date: Thu, 4 May 2023 10:59:46 -0400
Subject: [PATCH] Add num_blocks to VpxTplFrameStats

I realized the calculation of the size of the list of VpxTplBlockStats
is non-trivial. So it's better to add the field for the size.

Bug: b/273736974
Change-Id: Ic1b50597c1f89a8f866b5669ca676407be6dc9d8
---
 test/encode_api_test.cc     | 1 +
 vp9/encoder/vp9_tpl_model.c | 1 +
 vp9/vp9_cx_iface.c          | 1 +
 vpx/vpx_encoder.h           | 2 +-
 4 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/test/encode_api_test.cc b/test/encode_api_test.cc
index 0514cd8..e435ed8 100644
--- a/test/encode_api_test.cc
+++ b/test/encode_api_test.cc
@@ -406,6 +406,7 @@ class EncodeApiGetTplStatsTest
             if (tpl_stats[i].frame_width != 0) {
               ASSERT_EQ(tpl_stats[i].frame_width, width_);
               ASSERT_EQ(tpl_stats[i].frame_height, height_);
+              ASSERT_GT(tpl_stats[i].num_blocks, 0);
               ASSERT_NE(tpl_stats[i].block_stats_list, nullptr);
               stats_not_all_zero = true;
             }
diff --git a/vp9/encoder/vp9_tpl_model.c b/vp9/encoder/vp9_tpl_model.c
index ea5d61e..ed771dc 100644
--- a/vp9/encoder/vp9_tpl_model.c
+++ b/vp9/encoder/vp9_tpl_model.c
@@ -1354,6 +1354,7 @@ void vp9_init_tpl_buffer(VP9_COMP *cpi) {
         cm, cpi->tpl_frame_stats[frame].block_stats_list,
         vpx_calloc(mi_rows * mi_cols,
                    sizeof(*cpi->tpl_frame_stats[frame].block_stats_list)));
+    cpi->tpl_frame_stats[frame].num_blocks = mi_rows * mi_cols;
     cpi->tpl_stats[frame].is_valid = 0;
     cpi->tpl_stats[frame].width = mi_cols;
     cpi->tpl_stats[frame].height = mi_rows;
diff --git a/vp9/vp9_cx_iface.c b/vp9/vp9_cx_iface.c
index 66efba1..e264ae9 100644
--- a/vp9/vp9_cx_iface.c
+++ b/vp9/vp9_cx_iface.c
@@ -1799,6 +1799,7 @@ static vpx_codec_err_t ctrl_get_tpl_stats(vpx_codec_alg_priv_t *ctx,
   for (i = 0; i < MAX_ARF_GOP_SIZE; i++) {
     data[i].frame_width = cpi->tpl_frame_stats[i].frame_width;
     data[i].frame_height = cpi->tpl_frame_stats[i].frame_height;
+    data[i].num_blocks = cpi->tpl_frame_stats[i].num_blocks;
     data[i].block_stats_list = cpi->tpl_frame_stats[i].block_stats_list;
   }
 
diff --git a/vpx/vpx_encoder.h b/vpx/vpx_encoder.h
index 66c5a68..a7f1552 100644
--- a/vpx/vpx_encoder.h
+++ b/vpx/vpx_encoder.h
@@ -267,7 +267,7 @@ typedef struct VpxTplBlockStats {
 typedef struct VpxTplFrameStats {
   int frame_width;  /**< Frame width */
   int frame_height; /**< Frame height */
-  // Size of the list can be calculated from frame_width and frame_height.
+  int num_blocks;   /**< Number of blocks. Size of block_stats_list */
   VpxTplBlockStats *block_stats_list; /**< List of tpl stats for each block */
 } VpxTplFrameStats;
 
-- 
2.7.4