Support for placing Cues before Clusters

Adding support for placing Cues element before the Cluster element. We
recompute the new offsets using a recursive algorithm and update the Cues and
Seek Heads with the updates offsets.

Change-Id: I038f1a403b1defa853b9026bd3e48f4ad1006866

Makefile

@@ -2,7 +2,7 @@ CXX       := g++
 CXXFLAGS  := -W -Wall -g
 LIBWEBMA  := libwebm.a
 LIBWEBMSO := libwebm.so
-WEBMOBJS  := mkvparser.o mkvreader.o mkvmuxer.o mkvmuxerutil.o mkvwriter.o
+WEBMOBJS  := mkvparser.o mkvreader.o mkvmuxer.o mkvmuxerutil.o mkvwriter.o mkvreadablewriter.o
 OBJSA     := $(WEBMOBJS:.o=_a.o)
 OBJSSO    := $(WEBMOBJS:.o=_so.o)
 OBJECTS1  := sample.o

mkvmuxer.cpp

@@ -7,6 +7,7 @@
 // be found in the AUTHORS file in the root of the source tree.
 
 #include "mkvmuxer.hpp"
+#include "mkvreadablewriter.hpp"
 
 #include <cstdio>
 #include <cstdlib>
@@ -90,6 +91,25 @@ bool WriteEbmlHeader(IMkvWriter* writer) {
   return true;
 }
 
+bool ChunkedCopy(mkvmuxer::IMkvReadableWriter* source,
+                 mkvmuxer::IMkvWriter* dst,
+                 mkvmuxer::int64 start, int64 size) {
+  // TODO(vigneshv): Check if this is a reasonable value.
+  const uint32 kBufSize = 2048;
+  uint8* buf = new uint8[kBufSize];
+  int64 offset = start;
+  while (size > 0) {
+    const int64 read_len = (size > kBufSize) ? kBufSize : size;
+    if (source->Read(offset, read_len, buf))
+      return false;
+    dst->Write(buf, read_len);
+    offset += read_len;
+    size -= read_len;
+  }
+  delete[] buf;
+  return true;
+}
+
 ///////////////////////////////////////////////////////////////
 //
 // Frame Class
@@ -249,7 +269,7 @@ bool Cues::AddCue(CuePoint* cue) {
   return true;
 }
 
-const CuePoint* Cues::GetCueByIndex(int32 index) const {
+CuePoint* Cues::GetCueByIndex(int32 index) const {
   if (cue_entries_ == NULL)
     return NULL;
 
@@ -259,6 +279,14 @@ const CuePoint* Cues::GetCueByIndex(int32 index) const {
   return cue_entries_[index];
 }
 
+uint64 Cues::Size() {
+  uint64 size = 0;
+  for (int32 i = 0; i < cue_entries_size_; ++i)
+    size += GetCueByIndex(i)->Size();
+  size += EbmlMasterElementSize(kMkvCues, size);
+  return size;
+}
+
 bool Cues::Write(IMkvWriter* writer) const {
   if (!writer)
     return false;
@@ -1722,6 +1750,26 @@ bool SeekHead::AddSeekEntry(uint32 id, uint64 pos) {
   return false;
 }
 
+uint32 SeekHead::GetId(int index) const {
+  if (index < 0 || index >= kSeekEntryCount)
+    return -1;
+  return seek_entry_id_[index];
+}
+
+uint64 SeekHead::GetPosition(int index) const {
+  if (index < 0 || index >= kSeekEntryCount)
+    return -1;
+  return seek_entry_pos_[index];
+}
+
+bool SeekHead::SetSeekEntry(int index, uint32 id, uint64 position) {
+  if (index < 0 || index >= kSeekEntryCount)
+    return false;
+  seek_entry_id_[index] = id;
+  seek_entry_pos_[index] = position;
+  return true;
+}
+
 uint64 SeekHead::MaxEntrySize() const {
   const uint64 max_entry_payload_size =
       EbmlElementSize(kMkvSeekID, 0xffffffffULL) +
@@ -1914,6 +1962,7 @@ Segment::Segment()
       cluster_list_(NULL),
       cluster_list_capacity_(0),
       cluster_list_size_(0),
+      cues_position_(kAfterClusters),
       cues_track_(0),
       force_new_cluster_(false),
       frames_(NULL),
@@ -1973,16 +2022,84 @@ Segment::~Segment() {
   }
 }
 
+void Segment::MoveCuesBeforeClustersHelper(uint64 diff,
+                                           int32 index,
+                                           uint64* cues_size) {
+  const uint64 old_cues_size = *cues_size;
+  CuePoint* const cue_point = cues_.GetCueByIndex(index);
+  if (cue_point == NULL)
+    return;
+  const uint64 old_cue_point_size = cue_point->Size();
+  const uint64 cluster_pos = cue_point->cluster_pos() + diff;
+  cue_point->set_cluster_pos(cluster_pos);  // update the new cluster position
+  // New size of the cue is computed as follows
+  //    Let a = current size of Cues Element
+  //    Let b = Difference in Cue Point's size after this pass
+  //    Let c = Difference in length of Cues Element's size
+  //            (This is computed as CodedSize(a + b) - CodedSize(a)
+  //    Let d = a + b + c. Now d is the new size of the Cues element which is
+  //                       passed on to the next recursive call.
+  const uint64 cue_point_size_diff = cue_point->Size() - old_cue_point_size;
+  const uint64 cue_size_diff = GetCodedUIntSize(*cues_size +
+                                                cue_point_size_diff) -
+                               GetCodedUIntSize(*cues_size);
+  *cues_size += cue_point_size_diff + cue_size_diff;
+  diff = *cues_size - old_cues_size;
+  if (diff > 0) {
+    for (int32 i = 0; i < cues_.cue_entries_size(); ++i) {
+      MoveCuesBeforeClustersHelper(diff, i, cues_size);
+    }
+  }
+}
+
+void Segment::MoveCuesBeforeClusters() {
+  const uint64 current_cue_size = cues_.Size();
+  uint64 cue_size = current_cue_size;
+  for (int32 i = 0; i < cues_.cue_entries_size(); i++)
+    MoveCuesBeforeClustersHelper(current_cue_size, i, &cue_size);
+
+  // Adjust the Seek Entry to reflect the change in position
+  // of Cluster and Cues
+  int32 cluster_index;
+  int32 cues_index;
+  for (int32 i = 0; i < SeekHead::kSeekEntryCount; ++i) {
+    if (seek_head_.GetId(i) == kMkvCluster)
+      cluster_index = i;
+    if (seek_head_.GetId(i) == kMkvCues)
+      cues_index = i;
+  }
+  seek_head_.SetSeekEntry(cues_index, kMkvCues,
+                          seek_head_.GetPosition(cluster_index));
+  seek_head_.SetSeekEntry(cluster_index, kMkvCluster,
+                          cues_.Size() + seek_head_.GetPosition(cues_index));
+}
+
 bool Segment::Init(IMkvWriter* ptr_writer) {
   if (!ptr_writer) {
     return false;
   }
-  writer_cluster_ = ptr_writer;
+  writer_ = ptr_writer;
-  writer_cues_ = ptr_writer;
+  if (cues_position_ == kAfterClusters) {
-  writer_header_ = ptr_writer;
+    writer_cluster_ = writer_;
+    writer_cues_ = writer_;
+    writer_header_ = writer_;
+  } else {
+    writer_cluster_ = writer_temp_;
+    writer_cues_ = writer_temp_;
+    writer_header_ = writer_temp_;
+  }
   return segment_info_.Init();
 }
 
+bool Segment::WriteCuesBeforeClusters(IMkvReadableWriter* ptr_writer) {
+  if (!ptr_writer || writer_cluster_ || writer_cues_ || writer_header_) {
+    return false;
+  }
+  writer_temp_ = ptr_writer;
+  cues_position_ = kBeforeClusters;
+  return true;
+}
+
 bool Segment::Finalize() {
   if (WriteFramesAll() < 0)
     return false;
@@ -2007,7 +2124,6 @@ bool Segment::Finalize() {
     if (!segment_info_.Finalize(writer_header_))
       return false;
 
-    // TODO(fgalligan): Add support for putting the Cues at the front.
     if (output_cues_)
       if (!seek_head_.AddSeekEntry(kMkvCues, MaxOffset()))
         return false;
@@ -2026,6 +2142,18 @@ bool Segment::Finalize() {
         return false;
     }
 
+    const int64 current_offset = writer_cluster_->Position();
+    const int64 cluster_offset = cluster_list_[0]->size_position() -
+                                 GetUIntSize(kMkvCluster);
+    if (cues_position_ == kBeforeClusters) {
+      writer_cluster_ = writer_;
+      writer_cues_ = writer_;
+      writer_header_ = writer_;
+      ChunkedCopy(writer_temp_, writer_cluster_, 0, cluster_offset);
+      MoveCuesBeforeClusters();
+    }
+
+    // Write the seek headers and cues
     if (output_cues_)
       if (!cues_.Write(writer_cues_))
         return false;
@@ -2033,6 +2161,11 @@ bool Segment::Finalize() {
     if (!seek_head_.Finalize(writer_header_))
       return false;
 
+    if (cues_position_ == kBeforeClusters) {
+      ChunkedCopy(writer_temp_, writer_cluster_,
+                  cluster_offset, current_offset - cluster_offset);
+    }
+
     if (writer_header_->Seekable()) {
       if (size_position_ == -1)
         return false;

mkvmuxer.hpp

@@ -51,10 +51,21 @@ class IMkvWriter {
   LIBWEBM_DISALLOW_COPY_AND_ASSIGN(IMkvWriter);
 };
 
+class IMkvReadableWriter : public IMkvWriter {
+ public:
+  virtual int Read(long long position, long length,
+                   unsigned char* buffer) = 0;
+  virtual int Length(long long* total, long long* available) = 0;
+};
+
 // Writes out the EBML header for a WebM file. This function must be called
 // before any other libwebm writing functions are called.
 bool WriteEbmlHeader(IMkvWriter* writer);
 
+// Copies in Chunk from source to destination between the given byte positions
+bool ChunkedCopy(IMkvReadableWriter* source, IMkvWriter* dst,
+                 int64 start, int64 size);
+
 ///////////////////////////////////////////////////////////////
 // Class to hold data the will be written to a block.
 class Frame {
@@ -152,7 +163,10 @@ class Cues {
 
   // Returns the cue point by index. Returns NULL if there is no cue point
   // match.
-  const CuePoint* GetCueByIndex(int32 index) const;
+  CuePoint* GetCueByIndex(int32 index) const;
+
+  // Returns the total size of the Cues element
+  uint64 Size();
 
   // Output the Cues element to the writer. Returns true on success.
   bool Write(IMkvWriter* writer) const;
@@ -728,6 +742,7 @@ class Cluster {
   // Returns the size in bytes for the entire Cluster element.
   uint64 Size() const;
 
+  int64 size_position() const { return size_position_; }
   int32 blocks_added() const { return blocks_added_; }
   uint64 payload_size() const { return payload_size_; }
   int64 position_for_cues() const { return position_for_cues_; }
@@ -822,14 +837,26 @@ class SeekHead {
   // Writes out SeekHead and SeekEntry elements. Returns true on success.
   bool Finalize(IMkvWriter* writer) const;
 
+  // Returns the id of the Seek Entry at the given index. Returns -1 if index is
+  // out of range.
+  uint32 GetId(int index) const;
+
+  // Returns the position of the Seek Entry at the given index. Returns -1 if
+  // index is out of range.
+  uint64 GetPosition(int index) const;
+
+  // Sets the Seek Entry id and position at given index.
+  // Returns true on success.
+  bool SetSeekEntry(int index, uint32 id, uint64 position);
+
   // Reserves space by writing out a Void element which will be updated with
   // a SeekHead element later. Returns true on success.
   bool Write(IMkvWriter* writer);
 
- private:
+  // We are going to put a cap on the number of Seek Entries.
-   // We are going to put a cap on the number of Seek Entries.
   const static int32 kSeekEntryCount = 5;
 
+ private:
   // Returns the maximum size in bytes of one seek entry.
   uint64 MaxEntrySize() const;
 
@@ -901,6 +928,11 @@ class Segment {
     kFile = 0x2
   };
 
+  enum CuesPosition {
+    kAfterClusters = 0x0,  // Position Cues after Clusters - Default
+    kBeforeClusters = 0x1  // Position Cues before Clusters
+  };
+
   const static uint64 kDefaultMaxClusterDuration = 30000000000ULL;
 
   Segment();
@@ -910,6 +942,11 @@ class Segment {
   // |ptr_writer| is NULL.
   bool Init(IMkvWriter* ptr_writer);
 
+  // This function must be called before Init() if Cues are to be written before
+  // the Clusters. Input parameter is a IMkvReadableWriter object which supports
+  // both writing and reading.
+  bool WriteCuesBeforeClusters(IMkvReadableWriter* ptr_writer);
+
   // Adds a generic track to the segment.  Returns the newly-allocated
   // track object (which is owned by the segment) on success, NULL on
   // error. |number| is the number to use for the track.  |number|
@@ -1034,6 +1071,7 @@ class Segment {
   uint64 max_cluster_size() const { return max_cluster_size_; }
   void set_mode(Mode mode) { mode_ = mode; }
   Mode mode() const { return mode_; }
+  CuesPosition cues_position() const { return cues_position_; }
   bool output_cues() const { return output_cues_; }
   const SegmentInfo* segment_info() const { return &segment_info_; }
 
@@ -1088,6 +1126,21 @@ class Segment {
   // was necessary but creation was not successful.
   bool DoNewClusterProcessing(uint64 track_num, uint64 timestamp_ns, bool key);
 
+
+  // Adjusts Cue Point values (to place Cues before Clusters) so that they
+  // reflect the correct offsets.
+  void MoveCuesBeforeClusters();
+
+  // This function recursively computes the correct cluster offsets (this is
+  // done to move the Cues before Clusters). It recursively updates the change
+  // in size (which indicates a change in cluster offset) until no sizes change.
+  // Parameters:
+  // diff - indicates the difference in size of the Cues element that needs to
+  //        accounted for.
+  // index - index in the list of Cues which is currently being adjusted.
+  // cue_size - size of the Cues element.
+  void MoveCuesBeforeClustersHelper(uint64 diff, int index, uint64* cue_size);
+
   // Seeds the random number generator used to make UIDs.
   unsigned int seed_;
 
@@ -1132,6 +1185,9 @@ class Segment {
   // Number of clusters in the cluster list.
   int32 cluster_list_size_;
 
+  // Indicates whether Cues should be written before or after Clusters
+  CuesPosition cues_position_;
+
   // Track number that is associated with the cues element for this segment.
   uint64 cues_track_;
 
@@ -1191,6 +1247,10 @@ class Segment {
   IMkvWriter* writer_cues_;
   IMkvWriter* writer_header_;
 
+  // Pointer to actual writer object and temp writer object
+  IMkvWriter* writer_;
+  IMkvReadableWriter* writer_temp_;
+
   LIBWEBM_DISALLOW_COPY_AND_ASSIGN(Segment);
 };
 

mkvmuxerutil.hpp

@@ -22,6 +22,8 @@ const int64 kMaxBlockTimecode = 0x07FFFLL;
 int32 SerializeInt(IMkvWriter* writer, int64 value, int32 size);
 
 // Returns the size in bytes of the element.
+int32 GetUIntSize(uint64 value);
+int32 GetCodedUIntSize(uint64 value);
 uint64 EbmlMasterElementSize(uint64 type, uint64 value);
 uint64 EbmlElementSize(uint64 type, uint64 value);
 uint64 EbmlElementSize(uint64 type, float value);

sample_muxer.cpp

@@ -18,6 +18,7 @@
 
 // libwebm muxer includes
 #include "mkvmuxer.hpp"
+#include "mkvreadablewriter.hpp"
 #include "mkvwriter.hpp"
 #include "mkvmuxerutil.hpp"
 
@@ -53,6 +54,7 @@ void Usage() {
   printf("\n");
   printf("Cues options:\n");
   printf("  -output_cues_block_number <int> >0 outputs cue block number\n");
+  printf("  -cues_before_clusters <int> >0 puts Cues before Clusters\n");
   printf("\n");
   printf("Metadata options:\n");
   printf("  -webvtt-subtitles <vttfile>    "
@@ -145,6 +147,7 @@ int main(int argc, char* argv[]) {
   bool output_audio = true;
   bool live_mode = false;
   bool output_cues = true;
+  bool cues_before_clusters = false;
   bool cues_on_video_track = true;
   bool cues_on_audio_track = false;
   uint64 max_cluster_duration = 0;
@@ -182,6 +185,8 @@ int main(int argc, char* argv[]) {
       live_mode = strtol(argv[++i], &end, 10) == 0 ? false : true;
     } else if (!strcmp("-output_cues", argv[i]) && i < argc_check) {
       output_cues = strtol(argv[++i], &end, 10) == 0 ? false : true;
+    } else if (!strcmp("-cues_before_clusters", argv[i]) && i < argc_check) {
+      cues_before_clusters = strtol(argv[++i], &end, 10) == 0 ? false : true;
     } else if (!strcmp("-cues_on_video_track", argv[i]) && i < argc_check) {
       cues_on_video_track = strtol(argv[++i], &end, 10) == 0 ? false : true;
       if (cues_on_video_track)
@@ -267,6 +272,14 @@ int main(int argc, char* argv[]) {
   // Set Segment element attributes
   mkvmuxer::Segment muxer_segment;
 
+  mkvmuxer::MkvReadableWriter writer_temp;
+  if (cues_before_clusters) {
+    if (!writer_temp.Open(NULL, true)) {
+      printf("\n Filename is invalid or error while opening.\n");
+      return EXIT_FAILURE;
+    }
+    muxer_segment.WriteCuesBeforeClusters(&writer_temp);
+  }
   if (!muxer_segment.Init(&writer)) {
     printf("\n Could not initialize muxer segment!\n");
     return EXIT_FAILURE;
@@ -511,6 +524,8 @@ int main(int argc, char* argv[]) {
   delete parser_segment;
 
   writer.Close();
+  if (cues_before_clusters)
+    writer_temp.Close();
   reader.Close();
 
   return EXIT_SUCCESS;