RC: Implement using GOP from RC

Bug: b:450252793
Change-Id: I87446d3a2f54ab9011a2847ffc06c1a578a3eb35
diff --git a/aom/aom_ext_ratectrl.h b/aom/aom_ext_ratectrl.h
index edc5aca..c2bbf2c 100644
--- a/aom/aom_ext_ratectrl.h
+++ b/aom/aom_ext_ratectrl.h
@@ -43,6 +43,12 @@
  */
 #define AOM_RC_MAX_REF_FRAMES 8
 
+/*!\brief Max layer depth.
+ *
+ * Corresponds to MAX_ARF_LAYERS defined in ratectrl.h.
+ */
+#define AOM_RC_MAX_ARF_LAYERS 6
+
 /*!\brief The type of the external rate control.
  *
  * This controls what encoder parameters are determined by the external rate
diff --git a/av1/encoder/av1_ext_ratectrl.c b/av1/encoder/av1_ext_ratectrl.c
index 2b048c8..417c87e 100644
--- a/av1/encoder/av1_ext_ratectrl.c
+++ b/av1/encoder/av1_ext_ratectrl.c
@@ -119,6 +119,22 @@
   return AOM_CODEC_OK;
 }
 
+aom_codec_err_t av1_extrc_get_gop_decision(
+    AOM_EXT_RATECTRL *ext_ratectrl, aom_rc_gop_decision_t *gop_decision) {
+  aom_rc_status_t rc_status;
+  assert(ext_ratectrl != NULL);
+  assert(gop_decision != NULL);
+  if ((ext_ratectrl->funcs.rc_type & AOM_RC_GOP) == 0) {
+    return AOM_CODEC_INVALID_PARAM;
+  }
+  rc_status =
+      ext_ratectrl->funcs.get_gop_decision(ext_ratectrl->model, gop_decision);
+  if (rc_status == AOM_RC_ERROR) {
+    return AOM_CODEC_ERROR;
+  }
+  return AOM_CODEC_OK;
+}
+
 aom_codec_err_t av1_extrc_get_encodeframe_decision(
     AOM_EXT_RATECTRL *ext_ratectrl, int gop_index,
     aom_rc_encodeframe_decision_t *encode_frame_decision) {
diff --git a/av1/encoder/gop_structure.c b/av1/encoder/gop_structure.c
index 15ec0eb..a366d1f 100644
--- a/av1/encoder/gop_structure.c
+++ b/av1/encoder/gop_structure.c
@@ -17,10 +17,12 @@
 
 #include "aom/aom_codec.h"
 #include "aom/aom_encoder.h"
+#include "aom/aom_ext_ratectrl.h"
 
 #include "av1/common/av1_common_int.h"
 
 #include "av1/encoder/encoder.h"
+#include "av1/encoder/av1_ext_ratectrl.h"
 #include "av1/encoder/firstpass.h"
 #include "av1/encoder/gop_structure.h"
 #include "av1/encoder/pass2_strategy.h"
@@ -839,6 +841,38 @@
   gf_group->max_layer_depth = AOMMIN(log_gop_length, MAX_ARF_LAYERS);
 }
 
+static void construct_gop_structure_from_rc(
+    GF_GROUP *gf_group, aom_rc_gop_decision_t *rc_gop_decision) {
+  gf_group->size = rc_gop_decision->gop_frame_count;
+  for (int frame_index = 0; frame_index < gf_group->size; ++frame_index) {
+    aom_rc_gop_frame_t *gop_frame_rc =
+        &rc_gop_decision->gop_frame_list[frame_index];
+    gf_group->update_type[frame_index] = gop_frame_rc->update_type;
+    gf_group->layer_depth[frame_index] = gop_frame_rc->layer_depth;
+    gf_group->display_idx[frame_index] =
+        gop_frame_rc->display_idx + rc_gop_decision->global_order_idx_offset;
+    gf_group->update_ref_idx[frame_index] = gop_frame_rc->update_ref_idx;
+    gf_group->cur_frame_idx[frame_index] = gop_frame_rc->display_idx;
+    switch (gf_group->update_type[frame_index]) {
+      case LF_UPDATE:
+      case INTNL_OVERLAY_UPDATE:
+        gf_group->arf_src_offset[frame_index] = 0;
+        break;
+      case ARF_UPDATE:
+      case INTNL_ARF_UPDATE:
+        gf_group->arf_src_offset[frame_index] =
+            gop_frame_rc->display_idx - frame_index;
+        break;
+      default: gf_group->arf_src_offset[frame_index] = 0;
+    }
+    gf_group->cur_frame_idx[frame_index] = frame_index;
+    gf_group->frame_type[frame_index] =
+        gop_frame_rc->is_key_frame ? KEY_FRAME : INTER_FRAME;
+    gf_group->refbuf_state[frame_index] =
+        gop_frame_rc->is_key_frame ? REFBUF_RESET : REFBUF_UPDATE;
+  }
+}
+
 void av1_gop_setup_structure(AV1_COMP *cpi) {
   RATE_CONTROL *const rc = &cpi->rc;
   PRIMARY_RATE_CONTROL *const p_rc = &cpi->ppi->p_rc;
@@ -848,26 +882,39 @@
   const int key_frame = rc->frames_since_key == 0;
   FRAME_UPDATE_TYPE first_frame_update_type = ARF_UPDATE;
 
-  if (key_frame) {
-    first_frame_update_type = KF_UPDATE;
-    if (cpi->oxcf.kf_max_pyr_height != -1) {
-      gf_group->max_layer_depth_allowed = AOMMIN(
-          cpi->oxcf.kf_max_pyr_height, gf_group->max_layer_depth_allowed);
+  if (cpi->ext_ratectrl.ready &&
+      (cpi->ext_ratectrl.funcs.rc_type & AOM_RC_GOP) != 0 &&
+      cpi->ext_ratectrl.funcs.get_gop_decision != NULL) {
+    aom_rc_gop_decision_t gop_decision;
+    aom_codec_err_t codec_status =
+        av1_extrc_get_gop_decision(&cpi->ext_ratectrl, &gop_decision);
+    if (codec_status != AOM_CODEC_OK) {
+      aom_internal_error(cpi->common.error, codec_status,
+                         "av1_extrc_get_gop_decision() failed");
     }
-  } else if (!cpi->ppi->gf_state.arf_gf_boost_lst) {
-    first_frame_update_type = GF_UPDATE;
+    construct_gop_structure_from_rc(gf_group, &gop_decision);
+  } else {
+    if (key_frame) {
+      first_frame_update_type = KF_UPDATE;
+      if (cpi->oxcf.kf_max_pyr_height != -1) {
+        gf_group->max_layer_depth_allowed = AOMMIN(
+            cpi->oxcf.kf_max_pyr_height, gf_group->max_layer_depth_allowed);
+      }
+    } else if (!cpi->ppi->gf_state.arf_gf_boost_lst) {
+      first_frame_update_type = GF_UPDATE;
+    }
+
+    if (cpi->oxcf.algo_cfg.sharpness == 3)
+      gf_group->max_layer_depth_allowed =
+          AOMMIN(gf_group->max_layer_depth_allowed, 2);
+
+    gf_group->size = construct_multi_layer_gf_structure(
+        cpi, twopass, gf_group, rc, frame_info, p_rc->baseline_gf_interval,
+        first_frame_update_type);
+
+    if (gf_group->max_layer_depth_allowed == 0)
+      set_ld_layer_depth(gf_group, p_rc->baseline_gf_interval);
   }
-
-  if (cpi->oxcf.algo_cfg.sharpness == 3)
-    gf_group->max_layer_depth_allowed =
-        AOMMIN(gf_group->max_layer_depth_allowed, 2);
-
-  gf_group->size = construct_multi_layer_gf_structure(
-      cpi, twopass, gf_group, rc, frame_info, p_rc->baseline_gf_interval,
-      first_frame_update_type);
-
-  if (gf_group->max_layer_depth_allowed == 0)
-    set_ld_layer_depth(gf_group, p_rc->baseline_gf_interval);
 }
 
 int av1_gop_check_forward_keyframe(const GF_GROUP *gf_group,
diff --git a/test/ext_ratectrl_test.cc b/test/ext_ratectrl_test.cc
index fac4077..c031f4a 100644
--- a/test/ext_ratectrl_test.cc
+++ b/test/ext_ratectrl_test.cc
@@ -39,6 +39,51 @@
 // A flag to indicate if send_tpl_gop_stats() is called.
 bool is_send_extrc_tpl_gop_stats_called = false;
 
+// A flag to indicate if get_gop_decision() is called.
+bool is_get_gop_decision_called = false;
+
+// Construct a single ALTREF GOP.
+const int kGopFrameCount = kFrameNum + 1;
+aom_rc_gop_frame_t gop_frame_list[kGopFrameCount];
+
+aom_rc_status_t mock_get_gop_decision(aom_rc_model_t /*ratectrl_model*/,
+                                      aom_rc_gop_decision_t *gop_decision) {
+  gop_decision->gop_frame_count = kGopFrameCount;
+  gop_decision->gop_frame_list = gop_frame_list;
+  for (int i = 0; i < kGopFrameCount; ++i) {
+    auto current_gop_frame = &gop_decision->gop_frame_list[i];
+    current_gop_frame->coding_idx = i;
+    if (i == 0) {
+      // Key frame
+      current_gop_frame->is_key_frame = true;
+      current_gop_frame->update_type = AOM_RC_KF_UPDATE;
+      current_gop_frame->layer_depth = 0;
+      current_gop_frame->display_idx = 0;
+    } else if (i == 1) {
+      // ALTREF
+      current_gop_frame->is_key_frame = false;
+      current_gop_frame->update_type = AOM_RC_ARF_UPDATE;
+      current_gop_frame->layer_depth = 1;
+      current_gop_frame->display_idx = 4;
+    } else if (i == 5) {
+      // Overlay frame
+      current_gop_frame->is_key_frame = false;
+      current_gop_frame->update_type = AOM_RC_OVERLAY_UPDATE;
+      current_gop_frame->layer_depth = AOM_RC_MAX_ARF_LAYERS - 1;
+      current_gop_frame->display_idx = 4;
+    } else {
+      // Leaf frames
+      current_gop_frame->is_key_frame = false;
+      current_gop_frame->update_type = AOM_RC_LF_UPDATE;
+      current_gop_frame->layer_depth = AOM_RC_MAX_ARF_LAYERS - 1;
+      current_gop_frame->display_idx = i - 2;  // Key and ARF in the front
+    }
+  }
+  gop_decision->global_order_idx_offset = 0;
+  is_get_gop_decision_called = true;
+  return AOM_RC_OK;
+}
+
 aom_rc_status_t mock_create_model(void *priv,
                                   const aom_rc_config_t *ratectrl_config,
                                   aom_rc_model_t *ratectrl_model) {
@@ -89,6 +134,7 @@
     rc_funcs->delete_model = mock_delete_model;
     rc_funcs->send_firstpass_stats = mock_send_firstpass_stats;
     rc_funcs->send_tpl_gop_stats = mock_send_extrc_tpl_gop_stats;
+    rc_funcs->get_gop_decision = nullptr;
     rc_funcs->get_encodeframe_decision = nullptr;
     rc_funcs->update_encodeframe_result = nullptr;
   }
@@ -102,6 +148,7 @@
     is_delete_model_called = false;
     is_send_firstpass_stats_called = false;
     is_send_extrc_tpl_gop_stats_called = false;
+    is_get_gop_decision_called = false;
   }
 
   void PreEncodeFrameHook(::libaom_test::VideoSource *video,
@@ -129,7 +176,7 @@
 
 AV1_INSTANTIATE_TEST_SUITE(ExtRateCtrlTest,
                            ::testing::Values(::libaom_test::kTwoPassGood),
-                           ::testing::Values(0));
+                           ::testing::Values(3));
 
 // Corresponds to QP 43 for 8-bit video.
 const int kFrameQindex = 172;
@@ -178,5 +225,40 @@
 
 AV1_INSTANTIATE_TEST_SUITE(ExtRateCtrlQpTest,
                            ::testing::Values(::libaom_test::kTwoPassGood),
-                           ::testing::Values(0));
+                           ::testing::Values(3));
+
+class ExtRateCtrlGopTest : public ExtRateCtrlTest {
+ protected:
+  ExtRateCtrlGopTest() {
+    rc_funcs_.rc_type = AOM_RC_GOP;
+    rc_funcs_.get_gop_decision = mock_get_gop_decision;
+  }
+
+  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_EQ(gop_info_.gop_size, kGopFrameCount);
+  }
+
+  ~ExtRateCtrlGopTest() override = default;
+  aom_gop_info_t gop_info_;
+};
+
+TEST_P(ExtRateCtrlGopTest, TestExternalRateCtrlGop) {
+  ::libaom_test::Y4mVideoSource video("screendata.y4m", 0, kFrameNum);
+  ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+  EXPECT_TRUE(is_create_model_called);
+  EXPECT_TRUE(is_get_gop_decision_called);
+  EXPECT_TRUE(is_delete_model_called);
+}
+
+AV1_INSTANTIATE_TEST_SUITE(ExtRateCtrlGopTest,
+                           ::testing::Values(::libaom_test::kTwoPassGood),
+                           ::testing::Values(3));
 }  // namespace