matroskadec_haali: defer parsing of the full MKV header until required

This allows quickly scanning the file header for the segment UID instead
of parsing the full headers for this information.

3 files changed, 165 insertions, 6 deletions

diff --git a/libavformat/MatroskaParser.c b/libavformat/MatroskaParser.c
index da640837b0..d938fac8c9 100644
--- a/libavformat/MatroskaParser.c
+++ b/libavformat/MatroskaParser.c
@@ -2813,6 +2813,63 @@ segment:
   mf->tcCluster = mf->firstTimecode;
 }
 
+static void parseFileSparse(MatroskaFile *mf) {
+  ulonglong len = filepos(mf), adjust;
+  unsigned  i;
+  int id = readID(mf);
+  int found_uid = 0;
+
+  if (id==EOF)
+    errorjmp(mf,"Unexpected EOF at start of file");
+
+  // files with multiple concatenated segments can have only
+  // one EBML prolog
+  if (len > 0 && id == 0x18538067)
+    goto segment;
+
+  if (id!=0x1a45dfa3)
+    errorjmp(mf,"First element in file is not EBML");
+
+  parseEBML(mf,readSize(mf));
+
+  // next we need to find the first segment
+  for (;;) {
+    id = readID(mf);
+    if (id==EOF)
+      errorjmp(mf,"No segments found in the file");
+segment:
+    len = readSizeUnspec(mf);
+    if (id == 0x18538067) // Segment
+      break;
+    if (len == MAXU64)
+      errorjmp(mf,"No segments found in the file");
+    skipbytes(mf,len);
+  }
+
+  // found it
+  mf->pSegment = filepos(mf);
+
+  // we want to read data until we find a seekhead or a trackinfo
+  FOREACH2(mf,len,0x1f43b675)
+    case 0x1549a966: // SegmentInfo
+      mf->pSegmentInfo = cur;
+      FOREACH(mf,len)
+        case 0x73a4: // SegmentUID
+          if (len != sizeof(mf->Seg.UID))
+            errorjmp(mf,"SegmentUID size is not %d bytes",mf->Seg.UID);
+          readbytes(mf,mf->Seg.UID,sizeof(mf->Seg.UID));
+          return;
+      ENDFOR(mf);
+      break;
+  ENDFOR1(mf);
+    // if we found our segment info
+    if (mf->pSegmentInfo)
+      break;
+  ENDFOR2();
+
+  errorjmp(mf,"Couldn't find SegmentInfo");
+}
+
 static void DeleteChapter(MatroskaFile *mf,struct Chapter *ch) {
   unsigned i,j;
 
@@ -2865,6 +2922,33 @@ MatroskaFile  *mkv_OpenEx(InputStream *io,
   return mf;
 }
 
+MatroskaFile  *mkv_OpenSparse(InputStream *io,
+                              char *err_msg,unsigned msgsize)
+{
+  MatroskaFile        *mf = io->memalloc(io,sizeof(*mf));
+  if (mf == NULL) {
+    mystrlcpy(err_msg,"Out of memory",msgsize);
+    return NULL;
+  }
+
+  memset(mf,0,sizeof(*mf));
+
+  mf->cache = io;
+  mf->flags = MKVF_AVOID_SEEKS;
+  io->progress(io,0,0);
+
+  if (setjmp(mf->jb)==0) {
+    seek(mf,0);
+    parseFileSparse(mf);
+  } else { // parser error
+    mystrlcpy(err_msg,mf->errmsg,msgsize);
+    mkv_Close(mf);
+    return NULL;
+  }
+
+  return mf;
+}
+
 MatroskaFile  *mkv_Open(InputStream *io,
                         char *err_msg,unsigned msgsize)
 {
diff --git a/libavformat/MatroskaParser.h b/libavformat/MatroskaParser.h
index 93eda2d7e9..f09eba49fc 100644
--- a/libavformat/MatroskaParser.h
+++ b/libavformat/MatroskaParser.h
@@ -294,6 +294,11 @@ X MatroskaFile  *mkv_OpenEx(/* in */  InputStream *io,
 			  /* out */ char *err_msg,
 			  /* in */  unsigned msgsize);
 
+/* Open the file and only parse enough information to find the segment uid */
+X MatroskaFile  *mkv_OpenSparse(/* in */ InputStream *io,
+        /* out */ char *err_msg,
+        /* in */  unsigned msgsize);
+
 /* Close and deallocate mf
  * NULL pointer is ok and is simply ignored
  */
diff --git a/libavformat/matroskadec_haali.c b/libavformat/matroskadec_haali.c
index 67f8cbc124..060af64676 100644
--- a/libavformat/matroskadec_haali.c
+++ b/libavformat/matroskadec_haali.c
@@ -65,7 +65,9 @@ typedef struct MatroskaSegment {
   AVIOStream   *iostream;
   MatroskaFile *matroska;
   SegmentInfo  *info;
-  int free_avio;
+  char          UID[16];
+  int           free_avio;
+  int           failed;
 } MatroskaSegment;
 
 typedef struct VirtualTimelineEntry {
@@ -307,6 +309,55 @@ static MatroskaSegment* mkv_open_segment(AVFormatContext *s, AVIOContext *pb, ul
   }
 
   segment->info = mkv_GetFileInfo(segment->matroska);
+  memcpy(segment->UID, segment->info->UID, 16);
+
+  av_dynarray_add(&ctx->segments, &ctx->num_segments, segment);
+  return segment;
+}
+
+static void mkv_reopen_segment(AVFormatContext *s, MatroskaSegment *segment)
+{
+  MatroskaDemuxContext *ctx = (MatroskaDemuxContext *)s->priv_data;
+  char ErrorMessage[256];
+
+  /* reset packet size */
+  segment->iostream->pb->max_packet_size = 0;
+  ffio_set_buf_size(segment->iostream->pb, IO_BUFFER_SIZE * 4);
+
+  segment->matroska = mkv_OpenEx(&segment->iostream->base, 0, 0, ErrorMessage, sizeof(ErrorMessage));
+  if (!segment->matroska) {
+    av_log(s, AV_LOG_ERROR, "mkv_OpenEx returned error: %s\n", ErrorMessage);
+    segment->failed = 1;
+  }
+
+  segment->info = mkv_GetFileInfo(segment->matroska);
+}
+
+static MatroskaSegment* mkv_discover_segment(AVFormatContext *s, AVIOContext *pb)
+{
+  MatroskaDemuxContext *ctx = (MatroskaDemuxContext *)s->priv_data;
+  char ErrorMessage[256];
+  SegmentInfo *info;
+
+  MatroskaSegment *segment = av_mallocz(sizeof(*segment));
+  segment->index    = ctx->num_segments;
+  segment->iostream = aviostream_create(pb);
+  pb->max_packet_size = IO_BUFFER_SIZE;
+
+  segment->matroska = mkv_OpenSparse(&segment->iostream->base, ErrorMessage, sizeof(ErrorMessage));
+
+  if (!segment->matroska) {
+    av_log(s, AV_LOG_ERROR, "mkv_OpenEx returned error: %s\n", ErrorMessage);
+    av_freep(&segment->iostream);
+    av_freep(&segment);
+    return NULL;
+  }
+
+  info = mkv_GetFileInfo(segment->matroska);
+  memcpy(segment->UID, info->UID, 16);
+
+  mkv_Close(segment->matroska);
+  segment->matroska = NULL;
 
   av_dynarray_add(&ctx->segments, &ctx->num_segments, segment);
   return segment;
@@ -318,13 +369,16 @@ static int mkv_find_segment_avio(AVFormatContext *s, AVIOContext *pb, ulonglong
 
   av_log(s, AV_LOG_INFO, "Scanning for Segment at %I64d\n", base);
 
-  segment = mkv_open_segment(s, pb, base);
+  if (base == 0)
+    segment = mkv_discover_segment(s, pb);
+  else
+    segment = mkv_open_segment(s, pb, base);
 
   if (!segment)
     return 0;
 
   av_log(s, AV_LOG_INFO, "Found Segment with UID: %08x%08x%08x%08x\n",
-    *(unsigned int*)&segment->info->UID[0], *(unsigned int*)&segment->info->UID[4], *(unsigned int*)&segment->info->UID[8], *(unsigned int*)&segment->info->UID[12]);
+    *(unsigned int*)&segment->UID[0], *(unsigned int*)&segment->UID[4], *(unsigned int*)&segment->UID[8], *(unsigned int*)&segment->UID[12]);
 
   if (base == 0) {
     segment->free_avio = 1;
@@ -395,7 +449,7 @@ static MatroskaSegment* mkv_get_segment(AVFormatContext *s, char uid[16])
   MatroskaDemuxContext *ctx = (MatroskaDemuxContext *)s->priv_data;
   int i;
 
-  if (mkv_uid_zero(uid) || mkv_uid_compare(ctx->segments[0]->info->UID, uid))
+  if (mkv_uid_zero(uid) || mkv_uid_compare(ctx->segments[0]->UID, uid))
     return ctx->segments[0];
 
   if (!ctx->segments_scanned) {
@@ -418,7 +472,12 @@ static MatroskaSegment* mkv_get_segment(AVFormatContext *s, char uid[16])
   }
 
   for (i = 1; i < ctx->num_segments; i++) {
-    if (mkv_uid_compare(ctx->segments[i]->info->UID, uid)) {
+    if (!ctx->segments[i]->failed && mkv_uid_compare(ctx->segments[i]->UID, uid)) {
+      if (!ctx->segments[i]->matroska) {
+        mkv_reopen_segment(s, ctx->segments[i]);
+        if (ctx->segments[i]->failed)
+          break;
+      }
       return ctx->segments[i];
     }
   }
@@ -1130,7 +1189,8 @@ static int mkv_read_header(AVFormatContext *s)
   }
 
   for (i = 0; i < ctx->num_segments; i++) {
-    mkv_process_attachments(s, ctx->segments[i]);
+    if (ctx->segments[i]->matroska)
+      mkv_process_attachments(s, ctx->segments[i]);
   }
 
   if (tagCount > 0 && tags) {
@@ -1140,6 +1200,16 @@ static int mkv_read_header(AVFormatContext *s)
   /* Can only build the index after tracks are loaded */
   mkv_build_index(s);
 
+  /* close segments which were not needed for the virtual timeline */
+  for (i = 0; i < ctx->num_segments; i++) {
+    if (!ctx->segments[i]->matroska) {
+      if (ctx->segments[i]->free_avio)
+        avio_closep(&ctx->segments[i]->iostream->pb);
+      ctx->segments[i]->free_avio = 0;
+      av_freep(&ctx->segments[i]->iostream);
+    }
+  }
+
   return 0;
 }