input/ivf: Add seeking capability

7 files changed, 72 insertions, 19 deletions

diff --git a/tools/dav1d.c b/tools/dav1d.c
index 907af3f..bb1f4fd 100644
--- a/tools/dav1d.c
+++ b/tools/dav1d.c
@@ -197,7 +197,6 @@ int main(const int argc, char *const *const argv) {
                     seq_skip);
     }
 
-    //getc(stdin);
     if (cli_settings.limit != 0 && cli_settings.limit < total)
         total = cli_settings.limit;
 
diff --git a/tools/input/annexb.c b/tools/input/annexb.c
index 032480d..7c7d4e3 100644
--- a/tools/input/annexb.c
+++ b/tools/input/annexb.c
@@ -191,5 +191,6 @@ const Demuxer annexb_demuxer = {
     .probe_sz = PROBE_SIZE,
     .open = annexb_open,
     .read = annexb_read,
+    .seek = NULL,
     .close = annexb_close,
 };
diff --git a/tools/input/demuxer.h b/tools/input/demuxer.h
index c2b88e1..354db76 100644
--- a/tools/input/demuxer.h
+++ b/tools/input/demuxer.h
@@ -39,6 +39,7 @@ typedef struct Demuxer {
     int (*open)(DemuxerPriv *ctx, const char *filename,
                 unsigned fps[2], unsigned *num_frames, unsigned timebase[2]);
     int (*read)(DemuxerPriv *ctx, Dav1dData *data);
+    int (*seek)(DemuxerPriv *ctx, uint64_t pts);
     void (*close)(DemuxerPriv *ctx);
 } Demuxer;
 
diff --git a/tools/input/input.c b/tools/input/input.c
index 3ed6983..4756648 100644
--- a/tools/input/input.c
+++ b/tools/input/input.c
@@ -128,6 +128,10 @@ int input_read(DemuxerContext *const ctx, Dav1dData *const data) {
     return ctx->impl->read(ctx->data, data);
 }
 
+int input_seek(DemuxerContext *const ctx, const uint64_t pts) {
+    return ctx->impl->seek ? ctx->impl->seek(ctx->data, pts) : -1;
+}
+
 void input_close(DemuxerContext *const ctx) {
     ctx->impl->close(ctx->data);
     free(ctx);
diff --git a/tools/input/input.h b/tools/input/input.h
index 7b2fdc9..098219c 100644
--- a/tools/input/input.h
+++ b/tools/input/input.h
@@ -36,6 +36,7 @@ int input_open(DemuxerContext **const c_out,
                const char *const name, const char *const filename,
                unsigned fps[2], unsigned *num_frames, unsigned timebase[2]);
 int input_read(DemuxerContext *ctx, Dav1dData *data);
+int input_seek(DemuxerContext *ctx, uint64_t pts);
 void input_close(DemuxerContext *ctx);
 
 #endif /* DAV1D_INPUT_INPUT_H */
diff --git a/tools/input/ivf.c b/tools/input/ivf.c
index 7b572ee..e37fe99 100644
--- a/tools/input/ivf.c
+++ b/tools/input/ivf.c
@@ -29,6 +29,7 @@
 
 #include <errno.h>
 #include <limits.h>
+#include <math.h>
 #include <stdio.h>
 #include <stdint.h>
 #include <stdlib.h>
@@ -38,6 +39,10 @@
 
 typedef struct DemuxerPriv {
     FILE *f;
+    int broken;
+    double timebase;
+    uint64_t last_ts;
+    uint64_t step;
 } IvfInputContext;
 
 static const uint8_t probe_data[] = {
@@ -87,11 +92,17 @@ static int ivf_open(IvfInputContext *const c, const char *const file,
     timebase[1] = rl32(&hdr[20]);
     const unsigned duration = rl32(&hdr[24]);
 
-    uint8_t data[4];
+    uint8_t data[8];
+    c->broken = 0;
     for (*num_frames = 0;; (*num_frames)++) {
-        if ((res = fread(data, 4, 1, c->f)) != 1)
-            break; // EOF
-        fseeko(c->f, rl32(data) + 8, SEEK_CUR);
+        if ((res = fread(data, 4, 1, c->f)) != 1) break; // EOF
+        size_t sz = rl32(data);
+        if ((res = fread(data, 8, 1, c->f)) != 1) break; // EOF
+        const uint64_t ts = rl64(data);
+        if (*num_frames && ts <= c->last_ts)
+            c->broken = 1;
+        c->last_ts = ts;
+        fseeko(c->f, sz, SEEK_CUR);
     }
 
     uint64_t fps_num = (uint64_t) timebase[0] * *num_frames;
@@ -113,34 +124,68 @@ static int ivf_open(IvfInputContext *const c, const char *const file,
     } else {
         fps[0] = fps[1] = 0;
     }
+    c->timebase = (double)timebase[0] / timebase[1];
+    c->step = duration / *num_frames;
 
     fseeko(c->f, 32, SEEK_SET);
+    c->last_ts = 0;
 
     return 0;
 }
 
-static int ivf_read(IvfInputContext *const c, Dav1dData *const buf) {
+static inline int ivf_read_header(IvfInputContext *const c, ptrdiff_t *const sz,
+                                  int64_t *const off_, uint64_t *const ts)
+{
     uint8_t data[8];
-    uint8_t *ptr;
-    size_t res;
+    int64_t const off = ftello(c->f);
+    if (off_) *off_ = off;
+    if (fread(data, 4, 1, c->f) != 1) return -1; // EOF
+    *sz = rl32(data);
+    if (!c->broken) {
+        if (fread(data, 8, 1, c->f) != 1) return -1;
+        *ts = rl64(data);
+    } else {
+        if (fseeko(c->f, 8, SEEK_CUR)) return -1;
+        *ts = off > 32 ? c->last_ts + c->step : 0;
+    }
+    return 0;
+}
 
-    const int64_t off = ftello(c->f);
-    if ((res = fread(data, 4, 1, c->f)) != 1)
-        return -1; // EOF
-    const ptrdiff_t sz = rl32(data);
-    if ((res = fread(data, 8, 1, c->f)) != 1)
-        return -1; // EOF
-    ptr = dav1d_data_create(buf, sz);
-    if (!ptr) return -1;
-    buf->m.offset = off;
-    buf->m.timestamp = rl64(data);
-    if ((res = fread(ptr, sz, 1, c->f)) != 1) {
+static int ivf_read(IvfInputContext *const c, Dav1dData *const buf) {
+    uint8_t *ptr;
+    ptrdiff_t sz;
+    int64_t off;
+    uint64_t ts;
+    if (ivf_read_header(c, &sz, &off, &ts)) return -1;
+    if (!(ptr = dav1d_data_create(buf, sz))) return -1;
+    if (fread(ptr, sz, 1, c->f) != 1) {
         fprintf(stderr, "Failed to read frame data: %s\n", strerror(errno));
         dav1d_data_unref(buf);
         return -1;
     }
+    buf->m.offset = off;
+    buf->m.timestamp = ts;
+    c->last_ts = ts;
+    return 0;
+}
 
+static int ivf_seek(IvfInputContext *const c, const uint64_t pts) {
+    uint64_t cur;
+    const uint64_t ts = llround((pts * c->timebase) / 1000000000.0);
+    if (ts <= c->last_ts)
+        if (fseeko(c->f, 32, SEEK_SET)) goto error;
+    while (1) {
+        ptrdiff_t sz;
+        if (ivf_read_header(c, &sz, NULL, &cur)) goto error;
+        if (cur >= ts) break;
+        if (fseeko(c->f, sz, SEEK_CUR)) goto error;
+        c->last_ts = cur;
+    }
+    if (fseeko(c->f, -12, SEEK_CUR)) goto error;
     return 0;
+error:
+    fprintf(stderr, "Failed to seek: %s\n", strerror(errno));
+    return -1;
 }
 
 static void ivf_close(IvfInputContext *const c) {
@@ -154,5 +199,6 @@ const Demuxer ivf_demuxer = {
     .probe_sz = sizeof(probe_data),
     .open = ivf_open,
     .read = ivf_read,
+    .seek = ivf_seek,
     .close = ivf_close,
 };
diff --git a/tools/input/section5.c b/tools/input/section5.c
index 0c2ce28..b7cc203 100644
--- a/tools/input/section5.c
+++ b/tools/input/section5.c
@@ -181,5 +181,6 @@ const Demuxer section5_demuxer = {
     .probe_sz = PROBE_SIZE,
     .open = section5_open,
     .read = section5_read,
+    .seek = NULL,
     .close = section5_close,
 };