4 files changed, 72 insertions, 43 deletions

diff --git a/third_party/aom/common/tools_common.c b/third_party/aom/common/tools_common.c
index 359ec7341..21cd80026 100644
--- a/third_party/aom/common/tools_common.c
+++ b/third_party/aom/common/tools_common.c
@@ -236,7 +236,7 @@ double sse_to_psnr(double samples, double peak, double sse) {
 }
 
 // TODO(debargha): Consolidate the functions below into a separate file.
-static void highbd_img_upshift(aom_image_t *dst, aom_image_t *src,
+static void highbd_img_upshift(aom_image_t *dst, const aom_image_t *src,
                                int input_shift) {
   // Note the offset is 1 less than half.
   const int offset = input_shift > 0 ? (1 << (input_shift - 1)) - 1 : 0;
@@ -262,8 +262,8 @@ static void highbd_img_upshift(aom_image_t *dst, aom_image_t *src,
       h = (h + src->y_chroma_shift) >> src->y_chroma_shift;
     }
     for (y = 0; y < h; y++) {
-      uint16_t *p_src =
-          (uint16_t *)(src->planes[plane] + y * src->stride[plane]);
+      const uint16_t *p_src =
+          (const uint16_t *)(src->planes[plane] + y * src->stride[plane]);
       uint16_t *p_dst =
           (uint16_t *)(dst->planes[plane] + y * dst->stride[plane]);
       for (x = 0; x < w; x++) *p_dst++ = (*p_src++ << input_shift) + offset;
@@ -271,7 +271,7 @@ static void highbd_img_upshift(aom_image_t *dst, aom_image_t *src,
   }
 }
 
-static void lowbd_img_upshift(aom_image_t *dst, aom_image_t *src,
+static void lowbd_img_upshift(aom_image_t *dst, const aom_image_t *src,
                               int input_shift) {
   // Note the offset is 1 less than half.
   const int offset = input_shift > 0 ? (1 << (input_shift - 1)) - 1 : 0;
@@ -297,7 +297,7 @@ static void lowbd_img_upshift(aom_image_t *dst, aom_image_t *src,
       h = (h + src->y_chroma_shift) >> src->y_chroma_shift;
     }
     for (y = 0; y < h; y++) {
-      uint8_t *p_src = src->planes[plane] + y * src->stride[plane];
+      const uint8_t *p_src = src->planes[plane] + y * src->stride[plane];
       uint16_t *p_dst =
           (uint16_t *)(dst->planes[plane] + y * dst->stride[plane]);
       for (x = 0; x < w; x++) {
@@ -307,7 +307,8 @@ static void lowbd_img_upshift(aom_image_t *dst, aom_image_t *src,
   }
 }
 
-void aom_img_upshift(aom_image_t *dst, aom_image_t *src, int input_shift) {
+void aom_img_upshift(aom_image_t *dst, const aom_image_t *src,
+                     int input_shift) {
   if (src->fmt & AOM_IMG_FMT_HIGHBITDEPTH) {
     highbd_img_upshift(dst, src, input_shift);
   } else {
@@ -315,7 +316,7 @@ void aom_img_upshift(aom_image_t *dst, aom_image_t *src, int input_shift) {
   }
 }
 
-void aom_img_truncate_16_to_8(aom_image_t *dst, aom_image_t *src) {
+void aom_img_truncate_16_to_8(aom_image_t *dst, const aom_image_t *src) {
   int plane;
   if (dst->fmt + AOM_IMG_FMT_HIGHBITDEPTH != src->fmt || dst->d_w != src->d_w ||
       dst->d_h != src->d_h || dst->x_chroma_shift != src->x_chroma_shift ||
@@ -337,8 +338,8 @@ void aom_img_truncate_16_to_8(aom_image_t *dst, aom_image_t *src) {
       h = (h + src->y_chroma_shift) >> src->y_chroma_shift;
     }
     for (y = 0; y < h; y++) {
-      uint16_t *p_src =
-          (uint16_t *)(src->planes[plane] + y * src->stride[plane]);
+      const uint16_t *p_src =
+          (const uint16_t *)(src->planes[plane] + y * src->stride[plane]);
       uint8_t *p_dst = dst->planes[plane] + y * dst->stride[plane];
       for (x = 0; x < w; x++) {
         *p_dst++ = (uint8_t)(*p_src++);
@@ -347,7 +348,7 @@ void aom_img_truncate_16_to_8(aom_image_t *dst, aom_image_t *src) {
   }
 }
 
-static void highbd_img_downshift(aom_image_t *dst, aom_image_t *src,
+static void highbd_img_downshift(aom_image_t *dst, const aom_image_t *src,
                                  int down_shift) {
   int plane;
   if (dst->d_w != src->d_w || dst->d_h != src->d_h ||
@@ -371,8 +372,8 @@ static void highbd_img_downshift(aom_image_t *dst, aom_image_t *src,
       h = (h + src->y_chroma_shift) >> src->y_chroma_shift;
     }
     for (y = 0; y < h; y++) {
-      uint16_t *p_src =
-          (uint16_t *)(src->planes[plane] + y * src->stride[plane]);
+      const uint16_t *p_src =
+          (const uint16_t *)(src->planes[plane] + y * src->stride[plane]);
       uint16_t *p_dst =
           (uint16_t *)(dst->planes[plane] + y * dst->stride[plane]);
       for (x = 0; x < w; x++) *p_dst++ = *p_src++ >> down_shift;
@@ -380,7 +381,7 @@ static void highbd_img_downshift(aom_image_t *dst, aom_image_t *src,
   }
 }
 
-static void lowbd_img_downshift(aom_image_t *dst, aom_image_t *src,
+static void lowbd_img_downshift(aom_image_t *dst, const aom_image_t *src,
                                 int down_shift) {
   int plane;
   if (dst->d_w != src->d_w || dst->d_h != src->d_h ||
@@ -404,8 +405,8 @@ static void lowbd_img_downshift(aom_image_t *dst, aom_image_t *src,
       h = (h + src->y_chroma_shift) >> src->y_chroma_shift;
     }
     for (y = 0; y < h; y++) {
-      uint16_t *p_src =
-          (uint16_t *)(src->planes[plane] + y * src->stride[plane]);
+      const uint16_t *p_src =
+          (const uint16_t *)(src->planes[plane] + y * src->stride[plane]);
       uint8_t *p_dst = dst->planes[plane] + y * dst->stride[plane];
       for (x = 0; x < w; x++) {
         *p_dst++ = *p_src++ >> down_shift;
@@ -414,7 +415,8 @@ static void lowbd_img_downshift(aom_image_t *dst, aom_image_t *src,
   }
 }
 
-void aom_img_downshift(aom_image_t *dst, aom_image_t *src, int down_shift) {
+void aom_img_downshift(aom_image_t *dst, const aom_image_t *src,
+                       int down_shift) {
   if (dst->fmt & AOM_IMG_FMT_HIGHBITDEPTH) {
     highbd_img_downshift(dst, src, down_shift);
   } else {
diff --git a/third_party/aom/common/tools_common.h b/third_party/aom/common/tools_common.h
index abee4ea63..587903650 100644
--- a/third_party/aom/common/tools_common.h
+++ b/third_party/aom/common/tools_common.h
@@ -152,9 +152,10 @@ void aom_img_write(const aom_image_t *img, FILE *file);
 int aom_img_read(aom_image_t *img, FILE *file);
 
 double sse_to_psnr(double samples, double peak, double mse);
-void aom_img_upshift(aom_image_t *dst, aom_image_t *src, int input_shift);
-void aom_img_downshift(aom_image_t *dst, aom_image_t *src, int down_shift);
-void aom_img_truncate_16_to_8(aom_image_t *dst, aom_image_t *src);
+void aom_img_upshift(aom_image_t *dst, const aom_image_t *src, int input_shift);
+void aom_img_downshift(aom_image_t *dst, const aom_image_t *src,
+                       int down_shift);
+void aom_img_truncate_16_to_8(aom_image_t *dst, const aom_image_t *src);
 
 #ifdef __cplusplus
 } /* extern "C" */
diff --git a/third_party/aom/common/video_common.h b/third_party/aom/common/video_common.h
index f96af4b7e..965038d39 100644
--- a/third_party/aom/common/video_common.h
+++ b/third_party/aom/common/video_common.h
@@ -19,6 +19,7 @@ typedef struct {
   int frame_width;
   int frame_height;
   struct AvxRational time_base;
+  unsigned int is_annexb;
 } AvxVideoInfo;
 
 #endif  // VIDEO_COMMON_H_
diff --git a/third_party/aom/common/video_reader.c b/third_party/aom/common/video_reader.c
index f5327c928..b54c250c2 100644
--- a/third_party/aom/common/video_reader.c
+++ b/third_party/aom/common/video_reader.c
@@ -8,19 +8,20 @@
  * Media Patent License 1.0 was not distributed with this source code in the
  * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
  */
-#include "common/video_reader.h"
-
 #include <stdlib.h>
 #include <string.h>
+#include <assert.h>
 
 #include "aom_ports/mem_ops.h"
 #include "common/ivfdec.h"
-
-static const char *const kIVFSignature = "DKIF";
+#include "common/obudec.h"
+#include "common/tools_common.h"
+#include "common/video_reader.h"
 
 struct AvxVideoReaderStruct {
   AvxVideoInfo info;
-  FILE *file;
+  struct AvxInputContext input_ctx;
+  struct ObuDecInputContext obu_ctx;
   uint8_t *buffer;
   size_t buffer_size;
   size_t frame_size;
@@ -28,42 +29,64 @@ struct AvxVideoReaderStruct {
 };
 
 AvxVideoReader *aom_video_reader_open(const char *filename) {
-  char header[32];
   AvxVideoReader *reader = NULL;
   FILE *const file = fopen(filename, "rb");
   if (!file) return NULL;  // Can't open file
 
-  if (fread(header, 1, 32, file) != 32) return NULL;  // Can't read file header
-
-  if (memcmp(kIVFSignature, header, 4) != 0)
-    return NULL;  // Wrong IVF signature
-
-  if (mem_get_le16(header + 4) != 0) return NULL;  // Wrong IVF version
-
   reader = (AvxVideoReader *)calloc(1, sizeof(*reader));
-  if (!reader) return NULL;  // Can't allocate AvxVideoReader
+  if (!reader) {
+    fclose(file);
+    return NULL;  // Can't allocate AvxVideoReader
+  }
 
-  reader->file = file;
-  reader->info.codec_fourcc = mem_get_le32(header + 8);
-  reader->info.frame_width = mem_get_le16(header + 12);
-  reader->info.frame_height = mem_get_le16(header + 14);
-  reader->info.time_base.numerator = mem_get_le32(header + 16);
-  reader->info.time_base.denominator = mem_get_le32(header + 20);
+  reader->input_ctx.filename = filename;
+  reader->input_ctx.file = file;
+  reader->obu_ctx.avx_ctx = &reader->input_ctx;
+  reader->obu_ctx.is_annexb = 1;
+
+  if (file_is_ivf(&reader->input_ctx)) {
+    reader->input_ctx.file_type = FILE_TYPE_IVF;
+    reader->info.codec_fourcc = reader->input_ctx.fourcc;
+    reader->info.frame_width = reader->input_ctx.width;
+    reader->info.frame_height = reader->input_ctx.height;
+  } else if (file_is_obu(&reader->obu_ctx)) {
+    reader->input_ctx.file_type = FILE_TYPE_OBU;
+    // assume AV1
+    reader->info.codec_fourcc = AV1_FOURCC;
+    reader->info.is_annexb = reader->obu_ctx.is_annexb;
+  } else {
+    fclose(file);
+    free(reader);
+    return NULL;  // Unknown file type
+  }
 
   return reader;
 }
 
 void aom_video_reader_close(AvxVideoReader *reader) {
   if (reader) {
-    fclose(reader->file);
+    fclose(reader->input_ctx.file);
+    if (reader->input_ctx.file_type == FILE_TYPE_OBU) {
+      obudec_free(&reader->obu_ctx);
+    }
     free(reader->buffer);
     free(reader);
   }
 }
 
 int aom_video_reader_read_frame(AvxVideoReader *reader) {
-  return !ivf_read_frame(reader->file, &reader->buffer, &reader->frame_size,
-                         &reader->buffer_size, &reader->pts);
+  if (reader->input_ctx.file_type == FILE_TYPE_IVF) {
+    return !ivf_read_frame(reader->input_ctx.file, &reader->buffer,
+                           &reader->frame_size, &reader->buffer_size,
+                           &reader->pts);
+  } else if (reader->input_ctx.file_type == FILE_TYPE_OBU) {
+    return !obudec_read_temporal_unit(&reader->obu_ctx, &reader->buffer,
+                                      &reader->frame_size,
+                                      &reader->buffer_size);
+  } else {
+    assert(0);
+    return 0;
+  }
 }
 
 const uint8_t *aom_video_reader_get_frame(AvxVideoReader *reader,
@@ -77,7 +100,9 @@ int64_t aom_video_reader_get_frame_pts(AvxVideoReader *reader) {
   return (int64_t)reader->pts;
 }
 
-FILE *aom_video_reader_get_file(AvxVideoReader *reader) { return reader->file; }
+FILE *aom_video_reader_get_file(AvxVideoReader *reader) {
+  return reader->input_ctx.file;
+}
 
 const AvxVideoInfo *aom_video_reader_get_info(AvxVideoReader *reader) {
   return &reader->info;