From f71dd6e23e9d238ab0b03cdf105db733b8486778 Mon Sep 17 00:00:00 2001
From: angiebird <angiebird@google.com>
Date: Fri, 9 Oct 2020 19:03:23 -0700
Subject: [PATCH] vp9_extrc_get_encodeframe_decision()

Bug: webm:1707

Change-Id: I90a327b97d7158b65767fe3fbfd5f260030e17f5
---
 vp9/common/vp9_onyxc_int.h     |  6 ------
 vp9/encoder/vp9_encoder.c      | 11 +++++++++++
 vp9/encoder/vp9_ext_ratectrl.c | 31 +++++++++++++++++++++++++++++++
 vp9/encoder/vp9_ext_ratectrl.h |  4 ++++
 vp9/simple_encode.h            |  8 ++++----
 5 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/vp9/common/vp9_onyxc_int.h b/vp9/common/vp9_onyxc_int.h
index 6f9c6985f..ea47218b8 100644
--- a/vp9/common/vp9_onyxc_int.h
+++ b/vp9/common/vp9_onyxc_int.h
@@ -240,13 +240,11 @@ typedef struct VP9Common {
   // TODO(angiebird): current_video_frame/current_frame_coding_index into a
   // structure
   unsigned int current_video_frame;
-#if CONFIG_RATE_CTRL
   // Each show or no show frame is assigned with a coding index based on its
   // coding order (starting from zero).
 
   // Current frame's coding index.
   int current_frame_coding_index;
-#endif
   BITSTREAM_PROFILE profile;
 
   // VPX_BITS_8 in profile 0 or 1, VPX_BITS_10 or VPX_BITS_12 in profile 2 or 3.
@@ -276,9 +274,7 @@ typedef struct VP9Common {
 
 static INLINE void init_frame_indexes(VP9_COMMON *cm) {
   cm->current_video_frame = 0;
-#if CONFIG_RATE_CTRL
   cm->current_frame_coding_index = 0;
-#endif  // CONFIG_RATE_CTRL
 }
 
 static INLINE void update_frame_indexes(VP9_COMMON *cm, int show_frame) {
@@ -287,9 +283,7 @@ static INLINE void update_frame_indexes(VP9_COMMON *cm, int show_frame) {
     // update not a real frame
     ++cm->current_video_frame;
   }
-#if CONFIG_RATE_CTRL
   ++cm->current_frame_coding_index;
-#endif  // CONFIG_RATE_CTRL
 }
 
 typedef struct {
diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index 251a834d1..bf338eed8 100644
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -4480,6 +4480,14 @@ static void encode_with_recode_loop(VP9_COMP *cpi, size_t *size, uint8_t *dest
       q = cpi->encode_command.external_quantize_index;
     }
 #endif
+    if (cpi->ext_ratectrl.ready) {
+      const GF_GROUP *gf_group = &cpi->twopass.gf_group;
+      vpx_rc_encodeframe_decision_t encode_frame_decision;
+      vp9_extrc_get_encodeframe_decision(
+          &cpi->ext_ratectrl, gf_group, cm->current_video_frame,
+          cm->current_frame_coding_index, &encode_frame_decision);
+      q = encode_frame_decision.q_index;
+    }
 
     vp9_set_quantizer(cpi, q);
 
@@ -4519,6 +4527,9 @@ static void encode_with_recode_loop(VP9_COMP *cpi, size_t *size, uint8_t *dest
       if (frame_over_shoot_limit == 0) frame_over_shoot_limit = 1;
     }
 
+    if (cpi->ext_ratectrl.ready) {
+      break;
+    }
 #if CONFIG_RATE_CTRL
     // This part needs to be after save_coding_context() because
     // restore_coding_context will be called in the end of this function.
diff --git a/vp9/encoder/vp9_ext_ratectrl.c b/vp9/encoder/vp9_ext_ratectrl.c
index 51139f6ac..59b5c09b8 100644
--- a/vp9/encoder/vp9_ext_ratectrl.c
+++ b/vp9/encoder/vp9_ext_ratectrl.c
@@ -81,3 +81,34 @@ void vp9_extrc_send_firstpass_stats(EXT_RATECTRL *ext_ratectrl,
                                              rc_firstpass_stats);
   }
 }
+
+static int extrc_get_frame_type(FRAME_UPDATE_TYPE update_type) {
+  // TODO(angiebird): Add unit test to make sure this function behaves like
+  // get_frame_type_from_update_type()
+  // TODO(angiebird): Merge this function with get_frame_type_from_update_type()
+  switch (update_type) {
+    case KF_UPDATE: return 0;       // kFrameTypeKey;
+    case ARF_UPDATE: return 2;      // kFrameTypeAltRef;
+    case GF_UPDATE: return 4;       // kFrameTypeGolden;
+    case OVERLAY_UPDATE: return 3;  // kFrameTypeOverlay;
+    case LF_UPDATE: return 1;       // kFrameTypeInter;
+    default:
+      fprintf(stderr, "Unsupported update_type %d\n", update_type);
+      abort();
+      return 1;
+  }
+}
+
+void vp9_extrc_get_encodeframe_decision(
+    EXT_RATECTRL *ext_ratectrl, const GF_GROUP *gf_group, int show_index,
+    int coding_index, vpx_rc_encodeframe_decision_t *encode_frame_decision) {
+  if (ext_ratectrl->ready) {
+    FRAME_UPDATE_TYPE update_type = gf_group->update_type[gf_group->index];
+    vpx_rc_encodeframe_info_t encode_frame_info;
+    encode_frame_info.show_index = show_index;
+    encode_frame_info.coding_index = coding_index;
+    encode_frame_info.frame_type = extrc_get_frame_type(update_type);
+    ext_ratectrl->funcs.get_encodeframe_decision(
+        ext_ratectrl->model, &encode_frame_info, encode_frame_decision);
+  }
+}
diff --git a/vp9/encoder/vp9_ext_ratectrl.h b/vp9/encoder/vp9_ext_ratectrl.h
index 77d3bdd99..990735ebc 100644
--- a/vp9/encoder/vp9_ext_ratectrl.h
+++ b/vp9/encoder/vp9_ext_ratectrl.h
@@ -32,4 +32,8 @@ void vp9_extrc_delete(EXT_RATECTRL *ext_ratectrl);
 void vp9_extrc_send_firstpass_stats(EXT_RATECTRL *ext_ratectrl,
                                     const FIRST_PASS_INFO *first_pass_info);
 
+void vp9_extrc_get_encodeframe_decision(
+    EXT_RATECTRL *ext_ratectrl, const GF_GROUP *gf_group, int show_index,
+    int coding_index, vpx_rc_encodeframe_decision_t *encode_frame_decision);
+
 #endif  // VPX_VP9_ENCODER_VP9_EXT_RATECTRL_H_
diff --git a/vp9/simple_encode.h b/vp9/simple_encode.h
index ccbde9934..ce370a795 100644
--- a/vp9/simple_encode.h
+++ b/vp9/simple_encode.h
@@ -27,10 +27,10 @@ enum StatusCode {
 // TODO(angiebird): Add description for each frame type.
 enum FrameType {
   kFrameTypeKey = 0,
-  kFrameTypeInter,
-  kFrameTypeAltRef,
-  kFrameTypeOverlay,
-  kFrameTypeGolden,
+  kFrameTypeInter = 1,
+  kFrameTypeAltRef = 2,
+  kFrameTypeOverlay = 3,
+  kFrameTypeGolden = 4,
 };
 
 // TODO(angiebird): Add description for each reference frame type.
-- 
2.40.0