Modify the decoder AV1_SET_REFERENCE API

Modified the decoder AV1_SET_REFERENCE API so that it allowed to set a
reference frame by copying the buffer poniters instead of copying the
whole frame data.

Use use_external_refernce_buffers flag to identify whether y_buffer,
u_buffer, and v_buffer points to the internally allocated memory or
external buffers, so that these buffer pointers are restored after
the usage of external refernce buffers.

Modified the example code to use the modified API.

Change-Id: I2393e132258028501983c09b0bfadb553298a596
diff --git a/aom/aom.h b/aom/aom.h
index fecbeaf..72cc1df 100644
--- a/aom/aom.h
+++ b/aom/aom.h
@@ -104,8 +104,9 @@
  * Define the data struct to access av1 reference frames.
  */
 typedef struct av1_ref_frame {
-  int idx;         /**< frame index to get (input) */
-  aom_image_t img; /**< img structure to populate (output) */
+  int idx;              /**< frame index to get (input) */
+  int use_external_ref; /**< Directly use external ref buffer(decoder only) */
+  aom_image_t img;      /**< img structure to populate (output) */
 } av1_ref_frame_t;
 
 /*!\cond */
diff --git a/aom/src/aom_image.c b/aom/src/aom_image.c
index 2bcd020..437f024 100644
--- a/aom/src/aom_image.c
+++ b/aom/src/aom_image.c
@@ -188,29 +188,30 @@
       if (img->fmt & AOM_IMG_FMT_HAS_ALPHA) {
         img->planes[AOM_PLANE_ALPHA] =
             data + x * bytes_per_sample + y * img->stride[AOM_PLANE_ALPHA];
-        data += img->h * img->stride[AOM_PLANE_ALPHA];
+        data += (img->h + 2 * border) * img->stride[AOM_PLANE_ALPHA];
       }
 
       img->planes[AOM_PLANE_Y] =
           data + x * bytes_per_sample + y * img->stride[AOM_PLANE_Y];
-      data += img->h * img->stride[AOM_PLANE_Y];
+      data += (img->h + 2 * border) * img->stride[AOM_PLANE_Y];
 
+      unsigned int uv_border_h = border >> img->y_chroma_shift;
+      unsigned int uv_x = x >> img->x_chroma_shift;
+      unsigned int uv_y = y >> img->y_chroma_shift;
       if (!(img->fmt & AOM_IMG_FMT_UV_FLIP)) {
         img->planes[AOM_PLANE_U] =
-            data + (x >> img->x_chroma_shift) * bytes_per_sample +
-            (y >> img->y_chroma_shift) * img->stride[AOM_PLANE_U];
-        data += (img->h >> img->y_chroma_shift) * img->stride[AOM_PLANE_U];
+            data + uv_x * bytes_per_sample + uv_y * img->stride[AOM_PLANE_U];
+        data += ((img->h >> img->y_chroma_shift) + 2 * uv_border_h) *
+                img->stride[AOM_PLANE_U];
         img->planes[AOM_PLANE_V] =
-            data + (x >> img->x_chroma_shift) * bytes_per_sample +
-            (y >> img->y_chroma_shift) * img->stride[AOM_PLANE_V];
+            data + uv_x * bytes_per_sample + uv_y * img->stride[AOM_PLANE_V];
       } else {
         img->planes[AOM_PLANE_V] =
-            data + (x >> img->x_chroma_shift) * bytes_per_sample +
-            (y >> img->y_chroma_shift) * img->stride[AOM_PLANE_V];
-        data += (img->h >> img->y_chroma_shift) * img->stride[AOM_PLANE_V];
+            data + uv_x * bytes_per_sample + uv_y * img->stride[AOM_PLANE_V];
+        data += ((img->h >> img->y_chroma_shift) + 2 * uv_border_h) *
+                img->stride[AOM_PLANE_V];
         img->planes[AOM_PLANE_U] =
-            data + (x >> img->x_chroma_shift) * bytes_per_sample +
-            (y >> img->y_chroma_shift) * img->stride[AOM_PLANE_U];
+            data + uv_x * bytes_per_sample + uv_y * img->stride[AOM_PLANE_U];
       }
     }
     return 0;
diff --git a/aom_scale/generic/yv12config.c b/aom_scale/generic/yv12config.c
index 4cb2125..e6662a2 100644
--- a/aom_scale/generic/yv12config.c
+++ b/aom_scale/generic/yv12config.c
@@ -155,6 +155,8 @@
                                        (uv_border_h * uv_stride) + uv_border_w,
                                    aom_byte_align);
 
+    ybf->use_external_refernce_buffers = 0;
+
     if (use_highbitdepth) {
       if (ybf->y_buffer_8bit) aom_free(ybf->y_buffer_8bit);
       ybf->y_buffer_8bit = (uint8_t *)aom_memalign(32, (size_t)yplane_size);
diff --git a/aom_scale/yv12config.h b/aom_scale/yv12config.h
index 13f84f0..ca63fed 100644
--- a/aom_scale/yv12config.h
+++ b/aom_scale/yv12config.h
@@ -78,6 +78,14 @@
     uint8_t *buffers[4];
   };
 
+  // Indicate whether y_buffer, u_buffer, and v_buffer points to the internally
+  // allocated memory or external buffers.
+  int use_external_refernce_buffers;
+  // This is needed to store y_buffer, u_buffer, and v_buffer when set reference
+  // uses an external refernece, and restore those buffer pointers after the
+  // external reference frame is no longer used.
+  uint8_t *store_buf_adr[3];
+
   // If the frame is stored in a 16-bit buffer, this stores an 8-bit version
   // for use in global motion detection. It is allocated on-demand.
   uint8_t *y_buffer_8bit;
diff --git a/av1/av1_dx_iface.c b/av1/av1_dx_iface.c
index 9f292d1..266382f 100644
--- a/av1/av1_dx_iface.c
+++ b/av1/av1_dx_iface.c
@@ -718,7 +718,7 @@
     FrameWorkerData *const frame_worker_data = (FrameWorkerData *)worker->data1;
     image2yuvconfig(&frame->img, &sd);
     return av1_set_reference_dec(&frame_worker_data->pbi->common, frame->idx,
-                                 &sd);
+                                 frame->use_external_ref, &sd);
   } else {
     return AOM_CODEC_INVALID_PARAM;
   }
diff --git a/av1/common/onyxc_int.h b/av1/common/onyxc_int.h
index 5192ead..237c11e 100644
--- a/av1/common/onyxc_int.h
+++ b/av1/common/onyxc_int.h
@@ -595,6 +595,17 @@
     if (frame_bufs[i].ref_count == 0) break;
 
   if (i != FRAME_BUFFERS) {
+    if (frame_bufs[i].buf.use_external_refernce_buffers) {
+      // If this frame buffer's y_buffer, u_buffer, and v_buffer point to the
+      // external reference buffers. Restore the buffer pointers to point to the
+      // internally allocated memory.
+      YV12_BUFFER_CONFIG *ybf = &frame_bufs[i].buf;
+      ybf->y_buffer = ybf->store_buf_adr[0];
+      ybf->u_buffer = ybf->store_buf_adr[1];
+      ybf->v_buffer = ybf->store_buf_adr[2];
+      ybf->use_external_refernce_buffers = 0;
+    }
+
     frame_bufs[i].ref_count = 1;
   } else {
     // Reset i to be INVALID_IDX to indicate no free buffer found.
diff --git a/av1/decoder/decoder.c b/av1/decoder/decoder.c
index ecc5944..dc7d0df 100644
--- a/av1/decoder/decoder.c
+++ b/av1/decoder/decoder.c
@@ -186,7 +186,18 @@
   return cm->error.error_code;
 }
 
+static int equal_dimensions_and_border(const YV12_BUFFER_CONFIG *a,
+                                       const YV12_BUFFER_CONFIG *b) {
+  return a->y_height == b->y_height && a->y_width == b->y_width &&
+         a->uv_height == b->uv_height && a->uv_width == b->uv_width &&
+         a->y_stride == b->y_stride && a->uv_stride == b->uv_stride &&
+         a->border == b->border &&
+         (a->flags & YV12_FLAG_HIGHBITDEPTH) ==
+             (b->flags & YV12_FLAG_HIGHBITDEPTH);
+}
+
 aom_codec_err_t av1_set_reference_dec(AV1_COMMON *cm, int idx,
+                                      int use_external_ref,
                                       YV12_BUFFER_CONFIG *sd) {
   const int num_planes = av1_num_planes(cm);
   YV12_BUFFER_CONFIG *ref_buf = NULL;
@@ -199,12 +210,32 @@
     return AOM_CODEC_ERROR;
   }
 
-  if (!equal_dimensions(ref_buf, sd)) {
-    aom_internal_error(&cm->error, AOM_CODEC_ERROR,
-                       "Incorrect buffer dimensions");
+  if (!use_external_ref) {
+    if (!equal_dimensions(ref_buf, sd)) {
+      aom_internal_error(&cm->error, AOM_CODEC_ERROR,
+                         "Incorrect buffer dimensions");
+    } else {
+      // Overwrite the reference frame buffer.
+      aom_yv12_copy_frame(sd, ref_buf, num_planes);
+    }
   } else {
-    // Overwrite the reference frame buffer.
-    aom_yv12_copy_frame(sd, ref_buf, num_planes);
+    if (!equal_dimensions_and_border(ref_buf, sd)) {
+      aom_internal_error(&cm->error, AOM_CODEC_ERROR,
+                         "Incorrect buffer dimensions");
+    } else {
+      // Overwrite the reference frame buffer pointers.
+      // Once we no longer need the external reference buffer, these pointers
+      // are restored.
+      ref_buf->store_buf_adr[0] = ref_buf->y_buffer;
+      ref_buf->store_buf_adr[1] = ref_buf->u_buffer;
+      ref_buf->store_buf_adr[2] = ref_buf->v_buffer;
+      ref_buf->y_buffer = sd->y_buffer;
+      ref_buf->u_buffer = sd->u_buffer;
+      ref_buf->v_buffer = sd->v_buffer;
+      ref_buf->use_external_refernce_buffers = 1;
+      // TODO(yunqing): This will be removed later.
+      aom_yv12_extend_frame_borders_c(ref_buf, num_planes);
+    }
   }
 
   return cm->error.error_code;
diff --git a/av1/decoder/decoder.h b/av1/decoder/decoder.h
index ef66a65..486aad0 100644
--- a/av1/decoder/decoder.h
+++ b/av1/decoder/decoder.h
@@ -115,6 +115,7 @@
                                        YV12_BUFFER_CONFIG *sd);
 
 aom_codec_err_t av1_set_reference_dec(AV1_COMMON *cm, int idx,
+                                      int use_external_ref,
                                       YV12_BUFFER_CONFIG *sd);
 
 struct AV1Decoder *av1_decoder_create(BufferPool *const pool);
diff --git a/examples/aom_cx_set_ref.c b/examples/aom_cx_set_ref.c
index 7616e7b..c10f37d 100644
--- a/examples/aom_cx_set_ref.c
+++ b/examples/aom_cx_set_ref.c
@@ -73,25 +73,41 @@
 static void testing_decode(aom_codec_ctx_t *encoder, aom_codec_ctx_t *decoder,
                            unsigned int frame_out, int *mismatch_seen) {
   aom_image_t enc_img, dec_img;
-  struct av1_ref_frame ref_enc, ref_dec;
 
   if (*mismatch_seen) return;
 
-  ref_enc.idx = 0;
-  ref_dec.idx = 0;
-  if (aom_codec_control(encoder, AV1_GET_REFERENCE, &ref_enc))
+  /* Get the internal reference frame */
+  if (aom_codec_control(encoder, AV1_GET_NEW_FRAME_IMAGE, &enc_img))
     die_codec(encoder, "Failed to get encoder reference frame");
-  enc_img = ref_enc.img;
-  if (aom_codec_control(decoder, AV1_GET_REFERENCE, &ref_dec))
+  if (aom_codec_control(decoder, AV1_GET_NEW_FRAME_IMAGE, &dec_img))
     die_codec(decoder, "Failed to get decoder reference frame");
-  dec_img = ref_dec.img;
+
+  if ((enc_img.fmt & AOM_IMG_FMT_HIGHBITDEPTH) !=
+      (dec_img.fmt & AOM_IMG_FMT_HIGHBITDEPTH)) {
+    if (enc_img.fmt & AOM_IMG_FMT_HIGHBITDEPTH) {
+      aom_image_t enc_hbd_img;
+      aom_img_alloc(&enc_hbd_img, enc_img.fmt - AOM_IMG_FMT_HIGHBITDEPTH,
+                    enc_img.d_w, enc_img.d_h, 16);
+      aom_img_truncate_16_to_8(&enc_hbd_img, &enc_img);
+      enc_img = enc_hbd_img;
+    }
+    if (dec_img.fmt & AOM_IMG_FMT_HIGHBITDEPTH) {
+      aom_image_t dec_hbd_img;
+      aom_img_alloc(&dec_hbd_img, dec_img.fmt - AOM_IMG_FMT_HIGHBITDEPTH,
+                    dec_img.d_w, dec_img.d_h, 16);
+      aom_img_truncate_16_to_8(&dec_hbd_img, &dec_img);
+      dec_img = dec_hbd_img;
+    }
+  }
 
   if (!aom_compare_img(&enc_img, &dec_img)) {
     int y[4], u[4], v[4];
+    if (enc_img.fmt & AOM_IMG_FMT_HIGHBITDEPTH) {
+      aom_find_mismatch_high(&enc_img, &dec_img, y, u, v);
+    } else {
+      aom_find_mismatch(&enc_img, &dec_img, y, u, v);
+    }
 
-    *mismatch_seen = 1;
-
-    aom_find_mismatch(&enc_img, &dec_img, y, u, v);
     printf(
         "Encode/decode mismatch on frame %d at"
         " Y[%d, %d] {%d/%d},"
@@ -99,6 +115,7 @@
         " V[%d, %d] {%d/%d}",
         frame_out, y[0], y[1], y[2], y[3], u[0], u[1], u[2], u[3], v[0], v[1],
         v[2], v[3]);
+    *mismatch_seen = 1;
   }
 
   aom_img_free(&enc_img);
@@ -161,6 +178,7 @@
   aom_codec_enc_cfg_t cfg;
   unsigned int frame_in = 0;
   aom_image_t raw;
+  aom_image_t ext_ref;
   aom_codec_err_t res;
   AvxVideoInfo info;
   AvxVideoWriter *writer = NULL;
@@ -230,8 +248,12 @@
     die("Invalid frame size: %dx%d", info.frame_width, info.frame_height);
   }
 
-  // Allocate memory with the border so that it can be used as a reference.
   if (!aom_img_alloc_with_border(&raw, AOM_IMG_FMT_I420, info.frame_width,
+                                 info.frame_height, 1, 8, 0)) {
+    die("Failed to allocate image.");
+  }
+  // Allocate memory with the border so that it can be used as a reference.
+  if (!aom_img_alloc_with_border(&ext_ref, AOM_IMG_FMT_I420, info.frame_width,
                                  info.frame_height, 32, 8,
                                  AOM_BORDER_IN_PIXELS)) {
     die("Failed to allocate image.");
@@ -269,12 +291,14 @@
   }
 
   // Encode frames.
-  while (aom_img_read(&raw, infile)) {
+  aom_image_t *img_frm = &ext_ref;
+  while (aom_img_read(img_frm, infile)) {
     if (limit && frame_in >= limit) break;
     if (update_frame_num > 1 && frame_out + 1 == update_frame_num) {
       av1_ref_frame_t ref;
       ref.idx = 0;
-      ref.img = raw;
+      ref.use_external_ref = 0;
+      ref.img = *img_frm;
       // Set reference frame in encoder.
       if (aom_codec_control(&ecodec, AV1_SET_REFERENCE, &ref))
         die_codec(&ecodec, "Failed to set reference frame");
@@ -283,15 +307,17 @@
       // If set_reference in decoder is commented out, the enc/dec mismatch
       // would be seen.
       if (test_decode) {
+        ref.use_external_ref = 1;
         if (aom_codec_control(&dcodec, AV1_SET_REFERENCE, &ref))
           die_codec(&dcodec, "Failed to set reference frame");
       }
     }
 
-    encode_frame(&ecodec, &raw, frame_in, writer, test_decode, &dcodec,
+    encode_frame(&ecodec, img_frm, frame_in, writer, test_decode, &dcodec,
                  &frame_out, &mismatch_seen);
     frame_in++;
     if (mismatch_seen) break;
+    if (frame_out >= update_frame_num) img_frm = &raw;
   }
 
   // Flush encoder.