From 91c0477b4332b9e7880ce48332b1f77d5255a839 Mon Sep 17 00:00:00 2001
From: Emil Lundmark <lndmrk@webrtc.org>
Date: Tue, 25 Aug 2020 11:43:25 +0200
Subject: [PATCH] Add AV1 support to IVF reader and writer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Bug: None
Change-Id: I2ce36d282596bba54b7924e838b19b06eaacf0fd
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/182503
Reviewed-by: Erik Språng <sprang@webrtc.org>
Commit-Queue: Emil Lundmark <lndmrk@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#31998}
---
 modules/video_coding/utility/ivf_file_reader.cc   |  4 ++++
 .../utility/ivf_file_reader_unittest.cc           | 15 +++++++++++++++
 modules/video_coding/utility/ivf_file_writer.cc   |  6 ++++++
 .../utility/ivf_file_writer_unittest.cc           | 10 ++++++++++
 4 files changed, 35 insertions(+)

diff --git a/modules/video_coding/utility/ivf_file_reader.cc b/modules/video_coding/utility/ivf_file_reader.cc
index 9667bb7cec..27af52fe36 100644
--- a/modules/video_coding/utility/ivf_file_reader.cc
+++ b/modules/video_coding/utility/ivf_file_reader.cc
@@ -27,6 +27,7 @@ constexpr int kCodecTypeBytesCount = 4;
 constexpr uint8_t kFileHeaderStart[kCodecTypeBytesCount] = {'D', 'K', 'I', 'F'};
 constexpr uint8_t kVp8Header[kCodecTypeBytesCount] = {'V', 'P', '8', '0'};
 constexpr uint8_t kVp9Header[kCodecTypeBytesCount] = {'V', 'P', '9', '0'};
+constexpr uint8_t kAv1Header[kCodecTypeBytesCount] = {'A', 'V', '0', '1'};
 constexpr uint8_t kH264Header[kCodecTypeBytesCount] = {'H', '2', '6', '4'};
 
 }  // namespace
@@ -191,6 +192,9 @@ absl::optional<VideoCodecType> IvfFileReader::ParseCodecType(uint8_t* buffer,
   if (memcmp(&buffer[start_pos], kVp9Header, kCodecTypeBytesCount) == 0) {
     return VideoCodecType::kVideoCodecVP9;
   }
+  if (memcmp(&buffer[start_pos], kAv1Header, kCodecTypeBytesCount) == 0) {
+    return VideoCodecType::kVideoCodecAV1;
+  }
   if (memcmp(&buffer[start_pos], kH264Header, kCodecTypeBytesCount) == 0) {
     return VideoCodecType::kVideoCodecH264;
   }
diff --git a/modules/video_coding/utility/ivf_file_reader_unittest.cc b/modules/video_coding/utility/ivf_file_reader_unittest.cc
index 6ff580511b..58a808840d 100644
--- a/modules/video_coding/utility/ivf_file_reader_unittest.cc
+++ b/modules/video_coding/utility/ivf_file_reader_unittest.cc
@@ -145,6 +145,16 @@ TEST_F(IvfFileReaderTest, BasicVP9FileMsTimestamp) {
   ValidateContent(kVideoCodecVP9, true, 1);
 }
 
+TEST_F(IvfFileReaderTest, BasicAv1FileNtpTimestamp) {
+  CreateTestFile(kVideoCodecAV1, false, 1);
+  ValidateContent(kVideoCodecAV1, false, 1);
+}
+
+TEST_F(IvfFileReaderTest, BasicAv1FileMsTimestamp) {
+  CreateTestFile(kVideoCodecAV1, true, 1);
+  ValidateContent(kVideoCodecAV1, true, 1);
+}
+
 TEST_F(IvfFileReaderTest, BasicH264FileNtpTimestamp) {
   CreateTestFile(kVideoCodecH264, false, 1);
   ValidateContent(kVideoCodecH264, false, 1);
@@ -165,6 +175,11 @@ TEST_F(IvfFileReaderTest, MultilayerVP9FileNtpTimestamp) {
   ValidateContent(kVideoCodecVP9, false, 3);
 }
 
+TEST_F(IvfFileReaderTest, MultilayerAv1FileNtpTimestamp) {
+  CreateTestFile(kVideoCodecAV1, false, 3);
+  ValidateContent(kVideoCodecAV1, false, 3);
+}
+
 TEST_F(IvfFileReaderTest, MultilayerH264FileNtpTimestamp) {
   CreateTestFile(kVideoCodecH264, false, 3);
   ValidateContent(kVideoCodecH264, false, 3);
diff --git a/modules/video_coding/utility/ivf_file_writer.cc b/modules/video_coding/utility/ivf_file_writer.cc
index 46b8e87ba3..496da894a1 100644
--- a/modules/video_coding/utility/ivf_file_writer.cc
+++ b/modules/video_coding/utility/ivf_file_writer.cc
@@ -75,6 +75,12 @@ bool IvfFileWriter::WriteHeader() {
       ivf_header[10] = '9';
       ivf_header[11] = '0';
       break;
+    case kVideoCodecAV1:
+      ivf_header[8] = 'A';
+      ivf_header[9] = 'V';
+      ivf_header[10] = '0';
+      ivf_header[11] = '1';
+      break;
     case kVideoCodecH264:
       ivf_header[8] = 'H';
       ivf_header[9] = '2';
diff --git a/modules/video_coding/utility/ivf_file_writer_unittest.cc b/modules/video_coding/utility/ivf_file_writer_unittest.cc
index 49e0459ba6..8e781a7b22 100644
--- a/modules/video_coding/utility/ivf_file_writer_unittest.cc
+++ b/modules/video_coding/utility/ivf_file_writer_unittest.cc
@@ -147,6 +147,16 @@ TEST_F(IvfFileWriterTest, WritesBasicVP9FileMsTimestamp) {
   RunBasicFileStructureTest(kVideoCodecVP9, fourcc, true);
 }
 
+TEST_F(IvfFileWriterTest, WritesBasicAv1FileNtpTimestamp) {
+  const uint8_t fourcc[4] = {'A', 'V', '0', '1'};
+  RunBasicFileStructureTest(kVideoCodecAV1, fourcc, false);
+}
+
+TEST_F(IvfFileWriterTest, WritesBasicAv1FileMsTimestamp) {
+  const uint8_t fourcc[4] = {'A', 'V', '0', '1'};
+  RunBasicFileStructureTest(kVideoCodecAV1, fourcc, true);
+}
+
 TEST_F(IvfFileWriterTest, WritesBasicH264FileNtpTimestamp) {
   const uint8_t fourcc[4] = {'H', '2', '6', '4'};
   RunBasicFileStructureTest(kVideoCodecH264, fourcc, false);