* all: Added matroska writing. Patch by Mike Matsnev.
[x264/x264-sandbox.git] / matroska.c
1 /*****************************************************************************
2  * matroska.c:
3  *****************************************************************************
4  * Copyright (C) 2005 x264 project
5  * $Id: $
6  *
7  * Authors: Mike Matsnev
8  *
9  * This program is free software; you can redistribute it and/or modify
10  * it under the terms of the GNU General Public License as published by
11  * the Free Software Foundation; either version 2 of the License, or
12  * (at your option) any later version.
13  *
14  * This program is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
17  * GNU General Public License for more details.
18  *
19  * You should have received a copy of the GNU General Public License
20  * along with this program; if not, write to the Free Software
21  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111, USA.
22  *****************************************************************************/
23
24 #include <stdio.h>
25 #include <stdlib.h>
26 #include <string.h>
27
28 #ifdef HAVE_STDINT_H
29 #include <stdint.h>
30 #else
31 #include <inttypes.h>
32 #endif
33
34 #include "matroska.h"
35
36 #define CLSIZE    1048576
37
38 struct mk_Context {
39   struct mk_Context *next, **prev, *parent;
40   struct mk_Writer  *owner;
41   unsigned          id;
42
43   void              *data;
44   unsigned          d_cur, d_max;
45 };
46
47 typedef struct mk_Context mk_Context;
48
49 struct mk_Writer {
50   FILE                *fp;
51
52   unsigned            duration_ptr;
53
54   mk_Context          *root, *cluster, *frame;
55   mk_Context          *freelist;
56   mk_Context          *actlist;
57
58   int64_t             def_duration;
59   int64_t             timescale;
60   int64_t             cluster_tc_scaled;
61   int64_t             frame_tc, prev_frame_tc_scaled, max_frame_tc;
62
63   char                wrote_header, in_frame, keyframe;
64 };
65
66 static mk_Context *mk_createContext(mk_Writer *w, mk_Context *parent, unsigned id) {
67   mk_Context  *c;
68
69   if (w->freelist) {
70     c = w->freelist;
71     w->freelist = w->freelist->next;
72   } else {
73     c = malloc(sizeof(*c));
74     memset(c, 0, sizeof(*c));
75   }
76
77   if (c == NULL)
78     return NULL;
79
80   c->parent = parent;
81   c->owner = w;
82   c->id = id;
83
84   if (c->owner->actlist)
85     c->owner->actlist->prev = &c->next;
86   c->next = c->owner->actlist;
87   c->prev = &c->owner->actlist;
88
89   return c;
90 }
91
92 static int        mk_appendContextData(mk_Context *c, const void *data, unsigned size) {
93   unsigned  ns = c->d_cur + size;
94
95   if (ns > c->d_max) {
96     void      *dp;
97     unsigned  dn = c->d_max ? c->d_max << 1 : 16;
98     while (ns > dn)
99       dn <<= 1;
100
101     dp = realloc(c->data, dn);
102     if (dp == NULL)
103       return -1;
104
105     c->data = dp;
106     c->d_max = dn;
107   }
108
109   memcpy((char*)c->data + c->d_cur, data, size);
110
111   c->d_cur = ns;
112
113   return 0;
114 }
115
116 static void       mk_writeID(mk_Context *c, unsigned id) {
117   unsigned char   c_id[4] = { id >> 24, id >> 16, id >> 8, id };
118
119   if (c_id[0])
120     mk_appendContextData(c, c_id, 4);
121   else if (c_id[1])
122     mk_appendContextData(c, c_id+1, 3);
123   else if (c_id[2])
124     mk_appendContextData(c, c_id+2, 2);
125   else
126     mk_appendContextData(c, c_id+3, 1);
127 }
128
129 static void       mk_writeSize(mk_Context *c, unsigned size) {
130   unsigned char   c_size[5] = { 0x08, size >> 24, size >> 16, size >> 8, size };
131
132   if (size < 0x7f) {
133     c_size[4] |= 0x80;
134     mk_appendContextData(c, c_size+4, 1);
135   } else if (size < 0x3fff) {
136     c_size[3] |= 0x40;
137     mk_appendContextData(c, c_size+3, 2);
138   } else if (size < 0x1fffff) {
139     c_size[2] |= 0x20;
140     mk_appendContextData(c, c_size+2, 3);
141   } else if (size < 0x0fffffff) {
142     c_size[1] |= 0x10;
143     mk_appendContextData(c, c_size+1, 4);
144   } else
145     mk_appendContextData(c, c_size, 5);
146 }
147
148 static void       mk_flushContextID(mk_Context *c) {
149   unsigned char ff = 0xff;
150
151   if (c->id == 0)
152     return;
153
154   mk_writeID(c->parent, c->id);
155   mk_appendContextData(c->parent, &ff, 1);
156
157   c->id = 0;
158 }
159
160 static void       mk_flushContextData(mk_Context *c) {
161   if (c->d_cur == 0)
162     return;
163
164   if (c->parent)
165     mk_appendContextData(c->parent, c->data, c->d_cur);
166   else
167     fwrite(c->data, c->d_cur, 1, c->owner->fp);
168
169   c->d_cur = 0;
170 }
171
172 static unsigned   mk_closeContext(mk_Context *c, unsigned off) {
173   if (c->id) {
174     mk_writeID(c->parent, c->id);
175     mk_writeSize(c->parent, c->d_cur);
176   }
177
178   if (c->parent)
179     off += c->parent->d_cur;
180
181   mk_flushContextData(c);
182
183   if (c->next)
184     c->next->prev = c->prev;
185   *(c->prev) = c->next;
186   c->next = c->owner->freelist;
187   c->owner->freelist = c;
188
189   return off;
190 }
191
192 static void       mk_destroyContexts(mk_Writer *w) {
193   mk_Context  *cur, *next;
194
195   for (cur = w->freelist; cur; cur = next) {
196     next = cur->next;
197     free(cur->data);
198     free(cur);
199   }
200
201   for (cur = w->actlist; cur; cur = next) {
202     next = cur->next;
203     free(cur->data);
204     free(cur);
205   }
206
207   w->freelist = w->actlist = w->root = NULL;
208 }
209
210 static void       mk_writeStr(mk_Context *c, unsigned id, const char *str) {
211   size_t  len = strlen(str);
212
213   mk_writeID(c, id);
214   mk_writeSize(c, len);
215   mk_appendContextData(c, str, len);
216 }
217
218 static void       mk_writeBin(mk_Context *c, unsigned id, const void *data, unsigned size) {
219   mk_writeID(c, id);
220   mk_writeSize(c, size);
221   mk_appendContextData(c, data, size);
222 }
223
224 static void       mk_writeUInt(mk_Context *c, unsigned id, uint64_t ui) {
225   unsigned char   c_ui[8] = { ui >> 56, ui >> 48, ui >> 40, ui >> 32, ui >> 24, ui >> 16, ui >> 8, ui };
226   unsigned        i = 0;
227
228   mk_writeID(c, id);
229   while (i < 7 && c_ui[i] == 0)
230     ++i;
231   mk_writeSize(c, 8 - i);
232   mk_appendContextData(c, c_ui+i, 8 - i);
233 }
234
235 static void       mk_writeSInt(mk_Context *c, unsigned id, int64_t si) {
236   unsigned char   c_si[8] = { si >> 56, si >> 48, si >> 40, si >> 32, si >> 24, si >> 16, si >> 8, si };
237   unsigned        i = 0;
238
239   mk_writeID(c, id);
240   if (si < 0)
241     while (i < 7 && c_si[i] == 0xff && c_si[i+1] & 0x80)
242       ++i;
243   else
244     while (i < 7 && c_si[i] == 0 && !(c_si[i+1] & 0x80))
245       ++i;
246   mk_writeSize(c, 8 - i);
247   mk_appendContextData(c, c_si+i, 8 - i);
248 }
249
250 static void       mk_writeFloatRaw(mk_Context *c, float f) {
251   union {
252     float f;
253     unsigned u;
254   } u;
255   unsigned char c_f[4];
256
257   u.f = f;
258   c_f[0] = u.u >> 24;
259   c_f[1] = u.u >> 16;
260   c_f[2] = u.u >> 8;
261   c_f[3] = u.u;
262
263   mk_appendContextData(c, c_f, 4);
264 }
265
266 static void       mk_writeFloat(mk_Context *c, unsigned id, float f) {
267   mk_writeID(c, id);
268   mk_writeSize(c, 4);
269   mk_writeFloatRaw(c, f);
270 }
271
272 static unsigned   mk_ebmlSizeSize(unsigned s) {
273   if (s < 0x7f)
274     return 1;
275   if (s < 0x3fff)
276     return 2;
277   if (s < 0x1fffff)
278     return 3;
279   if (s < 0x0fffffff)
280     return 4;
281   return 5;
282 }
283
284 static unsigned   mk_ebmlSIntSize(int64_t si) {
285   unsigned char   c_si[8] = { si >> 56, si >> 48, si >> 40, si >> 32, si >> 24, si >> 16, si >> 8, si };
286   unsigned        i = 0;
287
288   if (si < 0)
289     while (i < 7 && c_si[i] == 0xff && c_si[i+1] & 0x80)
290       ++i;
291   else
292     while (i < 7 && c_si[i] == 0 && !(c_si[i+1] & 0x80))
293       ++i;
294
295   return 8 - i;
296 }
297
298 mk_Writer *mk_createWriter(const char *filename) {
299   mk_Writer *w = malloc(sizeof(*w));
300   if (w == NULL)
301     return NULL;
302
303   memset(w, 0, sizeof(*w));
304
305   w->root = mk_createContext(w, NULL, 0);
306   if (w->root == NULL) {
307     free(w);
308     return NULL;
309   }
310
311   w->fp = fopen(filename, "wb");
312   if (w->fp == NULL) {
313     mk_destroyContexts(w);
314     free(w);
315     return NULL;
316   }
317
318   w->timescale = 1000000;
319
320   return w;
321 }
322
323 int       mk_writeHeader(mk_Writer *w, const char *writingApp,
324                          const char *codecID,
325                          const void *codecPrivate, unsigned codecPrivateSize,
326                          int64_t default_frame_duration,
327                          int64_t timescale,
328                          unsigned width, unsigned height,
329                          unsigned d_width, unsigned d_height)
330 {
331   mk_Context  *c, *ti, *v;
332
333   if (w->wrote_header)
334     return -1;
335
336   w->timescale = timescale;
337   w->def_duration = default_frame_duration;
338
339   if ((c = mk_createContext(w, w->root, 0x1a45dfa3)) == NULL) // EBML
340     return -1;
341   mk_writeUInt(c, 0x4286, 1); // EBMLVersion
342   mk_writeUInt(c, 0x42f7, 1); // EBMLReadVersion
343   mk_writeUInt(c, 0x42f2, 4); // EBMLMaxIDLength
344   mk_writeUInt(c, 0x42f3, 8); // EBMLMaxSizeLength
345   mk_writeStr(c, 0x4282, "matroska"); // DocType
346   mk_writeUInt(c, 0x4287, 1); // DocTypeVersion
347   mk_writeUInt(c, 0x4285, 1); // DocTypeReadversion
348   mk_closeContext(c, 0);
349
350   if ((c = mk_createContext(w, w->root, 0x18538067)) == NULL) // Segment
351     return -1;
352   mk_flushContextID(c);
353   mk_closeContext(c, 0);
354
355   if ((c = mk_createContext(w, w->root, 0x1549a966)) == NULL) // SegmentInfo
356     return -1;
357   mk_writeStr(c, 0x4d80, "Haali Matroska Writer b0");
358   mk_writeStr(c, 0x5741, writingApp);
359   mk_writeUInt(c, 0x2ad7b1, w->timescale);
360   mk_writeFloat(c, 0x4489, 0);
361   w->duration_ptr = mk_closeContext(c, c->d_cur - 4);
362
363   if ((c = mk_createContext(w, w->root, 0x1654ae6b)) == NULL) // tracks
364     return -1;
365   if ((ti = mk_createContext(w, c, 0xae)) == NULL) // TrackEntry
366     return -1;
367   mk_writeUInt(ti, 0xd7, 1); // TrackNumber
368   mk_writeUInt(ti, 0x73c5, 1); // TrackUID
369   mk_writeUInt(ti, 0x83, 1); // TrackType
370   mk_writeUInt(ti, 0x9c, 0); // FlagLacing
371   mk_writeStr(ti, 0x86, codecID); // CodecID
372   if (codecPrivateSize)
373     mk_writeBin(ti, 0x63a2, codecPrivate, codecPrivateSize); // CodecPrivate
374   if (default_frame_duration)
375     mk_writeUInt(ti, 0x23e383, default_frame_duration); // DefaultDuration
376
377   if ((v = mk_createContext(w, ti, 0xe0)) == NULL) // Video
378     return -1;
379   mk_writeUInt(v, 0xb0, width);
380   mk_writeUInt(v, 0xba, height);
381   mk_writeUInt(v, 0x54b0, d_width);
382   mk_writeUInt(v, 0x54ba, d_height);
383   mk_closeContext(v, 0);
384
385   mk_closeContext(ti, 0);
386
387   mk_closeContext(c, 0);
388
389   mk_flushContextData(w->root);
390
391   w->wrote_header = 1;
392
393   return 0;
394 }
395
396 static void mk_closeCluster(mk_Writer *w) {
397   if (w->cluster == NULL)
398     return;
399   mk_closeContext(w->cluster, 0);
400   w->cluster = NULL;
401   mk_flushContextData(w->root);
402 }
403
404 int       mk_flushFrame(mk_Writer *w)
405 {
406   int64_t       delta, ref = 0;
407   unsigned      fsize, bgsize;
408   unsigned char c_delta_flags[3];
409
410   if (!w->in_frame)
411     return 0;
412
413   delta = w->frame_tc/w->timescale - w->cluster_tc_scaled;
414   if (delta > 32767ll || delta < -32768ll)
415     mk_closeCluster(w);
416
417   if (w->cluster == NULL) {
418     w->cluster_tc_scaled = w->frame_tc / w->timescale;
419     w->cluster = mk_createContext(w, w->root, 0x1f43b675); // Cluster
420     if (w->cluster == NULL)
421       return -1;
422
423     mk_writeUInt(w->cluster, 0xe7, w->cluster_tc_scaled); // Timecode
424
425     delta = 0;
426   }
427
428   fsize = w->frame ? w->frame->d_cur : 0;
429   bgsize = fsize + 4 + mk_ebmlSizeSize(fsize + 4) + 1;
430   if (!w->keyframe) {
431     ref = w->prev_frame_tc_scaled - w->cluster_tc_scaled - delta;
432     bgsize += 1 + 1 + mk_ebmlSIntSize(ref);
433   }
434
435   mk_writeID(w->cluster, 0xa0); // BlockGroup
436   mk_writeSize(w->cluster, bgsize);
437   mk_writeID(w->cluster, 0xa1); // Block
438   mk_writeSize(w->cluster, fsize + 4);
439   mk_writeSize(w->cluster, 1); // track number
440
441   c_delta_flags[0] = delta >> 8;
442   c_delta_flags[1] = delta;
443   c_delta_flags[2] = 0;
444   mk_appendContextData(w->cluster, c_delta_flags, 3);
445   if (w->frame) {
446     mk_appendContextData(w->cluster, w->frame->data, w->frame->d_cur);
447     w->frame->d_cur = 0;
448   }
449   if (!w->keyframe)
450     mk_writeSInt(w->cluster, 0xfb, ref); // ReferenceBlock
451
452   w->in_frame = 0;
453   w->prev_frame_tc_scaled = w->cluster_tc_scaled + delta;
454
455   if (w->cluster->d_cur > CLSIZE)
456     mk_closeCluster(w);
457
458   return 0;
459 }
460
461 int       mk_startFrame(mk_Writer *w) {
462   if (mk_flushFrame(w) < 0)
463     return -1;
464
465   w->in_frame = 1;
466   w->keyframe = 0;
467
468   return 0;
469 }
470
471 int       mk_setFrameFlags(mk_Writer *w,int64_t timestamp, int keyframe) {
472   if (!w->in_frame)
473     return -1;
474
475   w->frame_tc = timestamp;
476   w->keyframe = keyframe != 0;
477
478   if (w->max_frame_tc < timestamp)
479     w->max_frame_tc = timestamp;
480
481   return 0;
482 }
483
484 int       mk_addFrameData(mk_Writer *w, const void *data, unsigned size) {
485   if (!w->in_frame)
486     return -1;
487
488   if (w->frame == NULL)
489     if ((w->frame = mk_createContext(w, NULL, 0)) == NULL)
490       return -1;
491
492   return mk_appendContextData(w->frame, data, size);
493 }
494
495 void      mk_close(mk_Writer *w) {
496   mk_flushFrame(w);
497   mk_closeCluster(w);
498   if (w->wrote_header) {
499     fseek(w->fp, w->duration_ptr, SEEK_SET);
500     mk_writeFloatRaw(w->root, (float)((double)(w->max_frame_tc+w->def_duration) / w->timescale));
501     mk_flushContextData(w->root);
502   }
503   mk_destroyContexts(w);
504   fclose(w->fp);
505   free(w);
506 }
507
508 #ifdef TESTING
509 int   main(int argc, char **argv) {
510   char    f[4] = { 1, 2, 3, 4 };
511   mk_Writer *w = mk_createWriter("foo.mkv");
512   mk_writeHeader(w, "mktest", "V_MPEG/ISO/AVC", NULL, 0, 40000000, 352, 288, 3, 4);
513   mk_startFrame(w);
514   mk_setFrameFlags(w, 0, 1);
515   mk_addFrameData(w, f, 1);
516   mk_startFrame(w);
517   mk_setFrameFlags(w, 40000000, 0);
518   mk_addFrameData(w, f, 2);
519   mk_startFrame(w);
520   mk_setFrameFlags(w, 80000000, 0);
521   mk_addFrameData(w, f, 4);
522   mk_close(w);
523   return 0;
524 }
525 #endif