Add codec control to get GOP info

Change-Id: I820c3cc05efcbb8ec235711eedb85b2163cd5dcd

aom/aomcx.h

diff --git a/aom/aomcx.h b/aom/aomcx.h
index e7bfe96..4b39f34 100644
--- a/aom/aomcx.h
+++ b/aom/aomcx.h
@@ -1624,6 +1624,12 @@
    */
   AV1E_SET_EXTERNAL_RATE_CONTROL = 173,
 
+  /*!\brief Codec control function to get GOP structure from the encoder.
+   *
+   * args: a pointer to aom_gop_info_t
+   */
+  AV1E_GET_GOP_INFO,
+
   // Any new encoder control IDs should be added above.
   // Maximum allowed encoder control ID is 229.
   // No encoder control ID should be added below.
@@ -1860,6 +1866,20 @@
   AOM_FULL_SUPERFRAME_DROP, /**< Only full superframe can drop. */
 } AOM_SVC_FRAME_DROP_MODE;
 
+/*!\brief The GOP structure information determined by the encoder.
+ * 250 is MAX_STATIC_GF_GROUP_LENGTH defined in av1/firstpass.h.
+ * This is a subset of GF_GROUP. More fields can be added if needed.
+ */
+typedef struct aom_gop_info {
+  int gop_size; /**< The number of frames of this GOP */
+  /*! The coding index for each entry in the gop */
+  int coding_index[250];
+  /*! The display index for each entry in the gop */
+  int display_index[250];
+  /*! The layer depth for each entry in the gop */
+  int layer_depth[250];
+} aom_gop_info_t;
+
 /*!\cond */
 /*!\brief Encoder control function parameter type
  *
@@ -2370,6 +2390,9 @@
 AOM_CTRL_USE_TYPE(AV1E_SET_EXTERNAL_RATE_CONTROL, aom_rc_funcs_t *)
 #define AOM_CTRL_AV1E_SET_EXTERNAL_RATE_CONTROL
 
+AOM_CTRL_USE_TYPE(AV1E_GET_GOP_INFO, aom_gop_info_t *)
+#define AOM_CTRL_AV1E_GET_GOP_INFO
+
 /*!\endcond */
 /*! @} - end defgroup aom_encoder */
 #ifdef __cplusplus

av1/av1_cx_iface.c

diff --git a/av1/av1_cx_iface.c b/av1/av1_cx_iface.c
index fc4c431..a1ee702 100644
--- a/av1/av1_cx_iface.c
+++ b/av1/av1_cx_iface.c
@@ -1721,6 +1721,20 @@
   return res;
 }
 
+static aom_codec_err_t ctrl_get_gop_info(aom_codec_alg_priv_t *ctx,
+                                         va_list args) {
+  aom_gop_info_t *const gop_info = va_arg(args, aom_gop_info_t *);
+  if (gop_info == NULL) return AOM_CODEC_INVALID_PARAM;
+  const GF_GROUP *const gf_group = &ctx->ppi->gf_group;
+  gop_info->gop_size = gf_group->size;
+  for (int i = 0; i < gf_group->size; ++i) {
+    gop_info->coding_index[i] = i;
+    gop_info->display_index[i] = gf_group->display_idx[i];
+    gop_info->layer_depth[i] = gf_group->layer_depth[i];
+  }
+  return AOM_CODEC_OK;
+}
+
 static aom_codec_err_t ctrl_set_cpuused(aom_codec_alg_priv_t *ctx,
                                         va_list args) {
   struct av1_extracfg extra_cfg = ctx->extra_cfg;
@@ -4961,6 +4975,7 @@
   { AV1E_GET_LUMA_CDEF_STRENGTH, ctrl_get_luma_cdef_strength },
   { AV1E_GET_HIGH_MOTION_CONTENT_SCREEN_RTC,
     ctrl_get_high_motion_content_screen_rtc },
+  { AV1E_GET_GOP_INFO, ctrl_get_gop_info },
 
   CTRL_MAP_END,
 };

test/encode_api_test.cc

diff --git a/test/encode_api_test.cc b/test/encode_api_test.cc
index 60a2575..71bb049 100644
--- a/test/encode_api_test.cc
+++ b/test/encode_api_test.cc
@@ -24,6 +24,11 @@
 #include "aom/aom_encoder.h"
 #include "aom/aom_image.h"
 
+#include "test/codec_factory.h"
+#include "test/encode_test_driver.h"
+#include "test/util.h"
+#include "test/video_source.h"
+
 namespace {
 
 #if CONFIG_REALTIME_ONLY
@@ -1729,4 +1734,49 @@
   ASSERT_EQ(aom_codec_destroy(&enc), AOM_CODEC_OK);
 }
 
+class GetGopInfoTest : public ::libaom_test::EncoderTest,
+                       public ::testing::Test {
+ protected:
+  GetGopInfoTest() : EncoderTest(&::libaom_test::kAV1) {}
+  ~GetGopInfoTest() override = default;
+
+  void SetUp() override {
+    InitializeConfig(::libaom_test::kTwoPassGood);
+    cfg_.g_w = 176;
+    cfg_.g_h = 144;
+    cfg_.rc_target_bitrate = 200;
+    cfg_.g_lag_in_frames = 25;
+    cfg_.g_limit = kFrameLimit;
+  }
+
+  void PreEncodeFrameHook(::libaom_test::VideoSource *video,
+                          ::libaom_test::Encoder *encoder) override {
+    if (video->frame() == 0) {
+      encoder->Control(AOME_SET_CPUUSED, 3);
+    }
+  }
+
+  void PostEncodeFrameHook(::libaom_test::Encoder *encoder) override {
+    if (cfg_.g_pass == AOM_RC_FIRST_PASS) return;
+    encoder->Control(AV1E_GET_GOP_INFO, &gop_info_);
+  }
+
+  void FramePktHook(const aom_codec_cx_pkt_t * /*pkt*/) override {
+    // This is verified here (not in PostEncodeFrameHook) because
+    // PostEncodeFrameHook is also called when encoder is reading frames into
+    // lookahead buffer, when GOP structure hasn't been determined.
+    ASSERT_GT(gop_info_.gop_size, 0);
+  }
+
+  aom_gop_info_t gop_info_;
+  static constexpr int kFrameLimit = 10;
+};
+
+TEST_F(GetGopInfoTest, GetGopInfo) {
+  ::libaom_test::RandomVideoSource video;
+  video.SetSize(cfg_.g_w, cfg_.g_h);
+  video.set_limit(kFrameLimit);
+  ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+}
+
 }  // namespace

test/encode_test_driver.h

diff --git a/test/encode_test_driver.h b/test/encode_test_driver.h
index 739e9a0..a05d645 100644
--- a/test/encode_test_driver.h
+++ b/test/encode_test_driver.h
@@ -163,7 +163,12 @@
     const aom_codec_err_t res = aom_codec_control(&encoder_, ctrl_id, arg);
     ASSERT_EQ(AOM_CODEC_OK, res) << EncoderError();
   }
-#endif
+
+  void Control(int ctrl_id, aom_gop_info_t *arg) {
+    const aom_codec_err_t res = aom_codec_control(&encoder_, ctrl_id, arg);
+    ASSERT_EQ(AOM_CODEC_OK, res) << EncoderError();
+  }
+#endif  // CONFIG_AV1_ENCODER
 
   void SetOption(const char *name, const char *value) {
     const aom_codec_err_t res = aom_codec_set_option(&encoder_, name, value);