From 3571f4ab4c54d2f39f3faf043d86a7932ef4a5d6 Mon Sep 17 00:00:00 2001
From: Moritz Bunkus <moritz@bunkus.org>
Date: Sun, 7 Oct 2018 22:58:32 +0200
Subject: [PATCH] extract: add support for extracting AV1 to IVF

Part of the implementation of #2261.
---
 src/extract/xtr_base.cpp |  2 +-
 src/extract/xtr_ivf.cpp  | 32 ++++++++++++++++++++++++++++----
 src/extract/xtr_ivf.h    |  8 ++++++--
 tests/results.txt        |  1 +
 tests/test-658X_av1.rb   | 11 +++++++++++
 5 files changed, 47 insertions(+), 7 deletions(-)
 create mode 100755 tests/test-658X_av1.rb

diff --git a/src/extract/xtr_base.cpp b/src/extract/xtr_base.cpp
index bff9e04ef..047ec3184 100644
--- a/src/extract/xtr_base.cpp
+++ b/src/extract/xtr_base.cpp
@@ -184,7 +184,7 @@ xtr_base_c::create_extractor(const std::string &new_codec_id,
     result.reset(new xtr_mpeg1_2_video_c(new_codec_id, new_tid, tspec));
   else if (new_codec_id == MKV_V_THEORA)
     result.reset(new xtr_oggtheora_c(new_codec_id, new_tid, tspec));
-  else if ((new_codec_id == MKV_V_VP8) || (new_codec_id == MKV_V_VP9))
+  else if ((new_codec_id == MKV_V_VP8) || (new_codec_id == MKV_V_VP9) || (new_codec_id == MKV_V_AV1))
     result.reset(new xtr_ivf_c(new_codec_id, new_tid, tspec));
 
   // Subtitle formats
diff --git a/src/extract/xtr_ivf.cpp b/src/extract/xtr_ivf.cpp
index 418058e19..bb45cb145 100644
--- a/src/extract/xtr_ivf.cpp
+++ b/src/extract/xtr_ivf.cpp
@@ -12,6 +12,7 @@
 
 #include "common/common_pch.h"
 
+#include "common/av1.h"
 #include "common/codec.h"
 #include "common/ebml.h"
 #include "common/endian.h"
@@ -21,9 +22,7 @@ xtr_ivf_c::xtr_ivf_c(const std::string &codec_id,
                      int64_t tid,
                      track_spec_t &tspec)
   : xtr_base_c(codec_id, tid, tspec)
-  , m_frame_rate_num(0)
-  , m_frame_rate_den(0)
-  , m_frame_count(0)
+  , m_is_av1{codec_id == MKV_V_AV1}
 {
 }
 
@@ -52,7 +51,9 @@ xtr_ivf_c::create_file(xtr_base_c *master,
                             "track %4% with the CodecID '%5%' is already being written to the same file.\n"))
             % m_tid % m_codec_id % m_file_name % master->m_tid % master->m_codec_id);
 
-  auto fourcc = m_codec_id == MKV_V_VP8 ? "VP80" : "VP90";
+  auto fourcc = m_is_av1                ? "AV01"
+              : m_codec_id == MKV_V_VP8 ? "VP80"
+              :                           "VP90";
 
   memcpy(m_file_header.file_magic, "DKIF", 4);
   memcpy(m_file_header.fourcc,     fourcc, 4);
@@ -74,6 +75,8 @@ xtr_ivf_c::handle_frame(xtr_frame_t &f) {
          % f.timestamp % m_frame_rate_num % m_frame_rate_den % frame_number
          % (frame_number * 1000000000ull * m_frame_rate_den / m_frame_rate_num));
 
+  av1_prepend_temporal_delimiter_obu_if_needed(*f.frame);
+
   ivf::frame_header_t frame_header;
   put_uint32_le(&frame_header.frame_size, f.frame->get_size());
   put_uint32_le(&frame_header.timestamp,  frame_number);
@@ -91,3 +94,24 @@ xtr_ivf_c::finish_file() {
   m_out->setFilePointer(0);
   m_out->write(&m_file_header, sizeof(m_file_header));
 }
+
+void
+xtr_ivf_c::av1_prepend_temporal_delimiter_obu_if_needed(memory_c &frame) {
+  if (!m_is_av1 || !frame.get_size())
+    return;
+
+  auto type = (frame.get_buffer()[0] & 0x74) >> 3;
+
+  if (type == mtx::av1::OBU_TEMPORAL_DELIMITER)
+    return;
+
+  auto old_size = frame.get_size();
+
+  frame.resize(old_size + 2);
+
+  auto buffer = frame.get_buffer();
+
+  std::memmove(&buffer[2], &buffer[0], old_size);
+  buffer[0] = 0x12;             // type = OBU_TEMPORAL_DELIMITER, extension not present, OBU size field present
+  buffer[1] = 0x00;             // OBU size
+}
diff --git a/src/extract/xtr_ivf.h b/src/extract/xtr_ivf.h
index 32fc34a6c..3fe6b85e9 100644
--- a/src/extract/xtr_ivf.h
+++ b/src/extract/xtr_ivf.h
@@ -19,9 +19,10 @@
 
 class xtr_ivf_c: public xtr_base_c {
 public:
-  uint64_t m_frame_rate_num, m_frame_rate_den;
-  uint32_t m_frame_count;
+  uint64_t m_frame_rate_num{}, m_frame_rate_den{};
+  uint32_t m_frame_count{};
   ivf::file_header_t m_file_header;
+  bool m_is_av1{};
 
 public:
   xtr_ivf_c(const std::string &codec_id, int64_t tid, track_spec_t &tspec);
@@ -33,4 +34,7 @@ public:
   virtual const char *get_container_name() {
     return "IVF";
   };
+
+protected:
+  virtual void av1_prepend_temporal_delimiter_obu_if_needed(memory_c &frame);
 };
diff --git a/tests/results.txt b/tests/results.txt
index 879ed501b..c0f4480ab 100644
--- a/tests/results.txt
+++ b/tests/results.txt
@@ -503,3 +503,4 @@ T_654text_subtitles_without_duration:ada9aaa1ae07ff6b20fb0618304d0820:passed:201
 T_655mpeg_ts_teletext_subs_long_gap_until_end_of_display:6027ec285e5c03ec7fa84afea1d57f56:passed:20181005-212927:0.513727609
 T_656mpeg_ts_bad_utf8_in_service_names2:544606778297772ea61146bdbd8f1186:passed:20181006-122753:0.020246407
 T_657av1_from_mp4:08895840fd959294a31f7cbba8b2d9dd:passed:20181007-220646:0.040201275
+T_658X_av1:9d2beba51aa2d3d61ddb2b82f2c94937+076ba34e5ca17057f32ad54ef12ef42c:passed:20181007-225807:0.019174721
diff --git a/tests/test-658X_av1.rb b/tests/test-658X_av1.rb
new file mode 100755
index 000000000..2997302b3
--- /dev/null
+++ b/tests/test-658X_av1.rb
@@ -0,0 +1,11 @@
+#!/usr/bin/ruby -w
+
+# T_658X_av1
+describe "mkvextract / AV1 to IVF"
+
+test "extraction to IVF" do
+  extract "data/av1/av1.webm", 0 => "#{tmp}-1"
+  merge "#{tmp}-1", :output => "#{tmp}-2"
+
+  (1..2).map { |idx| hash_file("#{tmp}-#{idx}") }.join('+')
+end