avformat/mxfdec: fix sorting of index segments
[ffmpeg.git] / libavcodec / cbs_mpeg2.c
1 /*
2  * This file is part of FFmpeg.
3  *
4  * FFmpeg is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU Lesser General Public
6  * License as published by the Free Software Foundation; either
7  * version 2.1 of the License, or (at your option) any later version.
8  *
9  * FFmpeg is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
12  * Lesser General Public License for more details.
13  *
14  * You should have received a copy of the GNU Lesser General Public
15  * License along with FFmpeg; if not, write to the Free Software
16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
17  */
18
19 #include "libavutil/avassert.h"
20
21 #include "cbs.h"
22 #include "cbs_internal.h"
23 #include "cbs_mpeg2.h"
24 #include "internal.h"
25
26
27 #define HEADER(name) do { \
28         ff_cbs_trace_header(ctx, name); \
29     } while (0)
30
31 #define CHECK(call) do { \
32         err = (call); \
33         if (err < 0) \
34             return err; \
35     } while (0)
36
37 #define FUNC_NAME(rw, codec, name) cbs_ ## codec ## _ ## rw ## _ ## name
38 #define FUNC_MPEG2(rw, name) FUNC_NAME(rw, mpeg2, name)
39 #define FUNC(name) FUNC_MPEG2(READWRITE, name)
40
41
42 #define READ
43 #define READWRITE read
44 #define RWContext GetBitContext
45
46 #define xui(width, name, var) do { \
47         uint32_t value = 0; \
48         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
49                                    &value, 0, (1 << width) - 1)); \
50         var = value; \
51     } while (0)
52
53 #define ui(width, name) \
54         xui(width, name, current->name)
55
56 #define marker_bit() do { \
57         av_unused uint32_t one; \
58         CHECK(ff_cbs_read_unsigned(ctx, rw, 1, "marker_bit", &one, 1, 1)); \
59     } while (0)
60
61 #define nextbits(width, compare, var) \
62     (get_bits_left(rw) >= width && \
63      (var = show_bits(rw, width)) == (compare))
64
65 #include "cbs_mpeg2_syntax_template.c"
66
67 #undef READ
68 #undef READWRITE
69 #undef RWContext
70 #undef xui
71 #undef ui
72 #undef marker_bit
73 #undef nextbits
74
75
76 #define WRITE
77 #define READWRITE write
78 #define RWContext PutBitContext
79
80 #define xui(width, name, var) do { \
81         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
82                                     var, 0, (1 << width) - 1)); \
83     } while (0)
84
85 #define ui(width, name) \
86         xui(width, name, current->name)
87
88 #define marker_bit() do { \
89         CHECK(ff_cbs_write_unsigned(ctx, rw, 1, "marker_bit", 1, 1, 1)); \
90     } while (0)
91
92 #define nextbits(width, compare, var) (var)
93
94 #include "cbs_mpeg2_syntax_template.c"
95
96 #undef READ
97 #undef READWRITE
98 #undef RWContext
99 #undef xui
100 #undef ui
101 #undef marker_bit
102 #undef nextbits
103
104
105 static void cbs_mpeg2_free_user_data(void *unit, uint8_t *content)
106 {
107     MPEG2RawUserData *user = (MPEG2RawUserData*)content;
108     av_buffer_unref(&user->user_data_ref);
109     av_freep(&content);
110 }
111
112 static void cbs_mpeg2_free_slice(void *unit, uint8_t *content)
113 {
114     MPEG2RawSlice *slice = (MPEG2RawSlice*)content;
115     av_buffer_unref(&slice->header.extra_information_ref);
116     av_buffer_unref(&slice->data_ref);
117     av_freep(&content);
118 }
119
120 static int cbs_mpeg2_split_fragment(CodedBitstreamContext *ctx,
121                                     CodedBitstreamFragment *frag,
122                                     int header)
123 {
124     const uint8_t *start, *end;
125     uint8_t *unit_data;
126     uint32_t start_code = -1, next_start_code = -1;
127     size_t unit_size;
128     int err, i, unit_type;
129
130     start = avpriv_find_start_code(frag->data, frag->data + frag->data_size,
131                                    &start_code);
132     for (i = 0;; i++) {
133         end = avpriv_find_start_code(start, frag->data + frag->data_size,
134                                      &next_start_code);
135
136         unit_type = start_code & 0xff;
137
138         // The start and end pointers point at to the byte following the
139         // start_code_identifier in the start code that they found.
140         if (end == frag->data + frag->data_size) {
141             // We didn't find a start code, so this is the final unit.
142             unit_size = end - (start - 1);
143         } else {
144             // Unit runs from start to the beginning of the start code
145             // pointed to by end (including any padding zeroes).
146             unit_size = (end - 4) - (start - 1);
147         }
148
149         unit_data = av_malloc(unit_size + AV_INPUT_BUFFER_PADDING_SIZE);
150         if (!unit_data)
151             return AVERROR(ENOMEM);
152         memcpy(unit_data, start - 1, unit_size);
153         memset(unit_data + unit_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
154
155         err = ff_cbs_insert_unit_data(ctx, frag, i, unit_type,
156                                       unit_data, unit_size, NULL);
157         if (err < 0) {
158             av_freep(&unit_data);
159             return err;
160         }
161
162         if (end == frag->data + frag->data_size)
163             break;
164
165         start_code = next_start_code;
166         start = end;
167     }
168
169     return 0;
170 }
171
172 static int cbs_mpeg2_read_unit(CodedBitstreamContext *ctx,
173                                CodedBitstreamUnit *unit)
174 {
175     GetBitContext gbc;
176     int err;
177
178     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
179     if (err < 0)
180         return err;
181
182     if (MPEG2_START_IS_SLICE(unit->type)) {
183         MPEG2RawSlice *slice;
184         int pos, len;
185
186         err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*slice),
187                                         &cbs_mpeg2_free_slice);
188         if (err < 0)
189             return err;
190         slice = unit->content;
191
192         err = cbs_mpeg2_read_slice_header(ctx, &gbc, &slice->header);
193         if (err < 0)
194             return err;
195
196         pos = get_bits_count(&gbc);
197         len = unit->data_size;
198
199         slice->data_size = len - pos / 8;
200         slice->data_ref  = av_buffer_alloc(slice->data_size +
201                                            AV_INPUT_BUFFER_PADDING_SIZE);
202         if (!slice->data_ref)
203             return AVERROR(ENOMEM);
204         slice->data = slice->data_ref->data;
205
206         memcpy(slice->data,
207                unit->data + pos / 8, slice->data_size);
208         memset(slice->data + slice->data_size, 0,
209                AV_INPUT_BUFFER_PADDING_SIZE);
210         slice->data_bit_start = pos % 8;
211
212     } else {
213         switch (unit->type) {
214 #define START(start_code, type, read_func, free_func) \
215         case start_code: \
216             { \
217                 type *header; \
218                 err = ff_cbs_alloc_unit_content(ctx, unit, \
219                                                 sizeof(*header), free_func); \
220                 if (err < 0) \
221                     return err; \
222                 header = unit->content; \
223                 err = cbs_mpeg2_read_ ## read_func(ctx, &gbc, header); \
224                 if (err < 0) \
225                     return err; \
226             } \
227             break;
228             START(0x00, MPEG2RawPictureHeader,  picture_header,  NULL);
229             START(0xb2, MPEG2RawUserData,       user_data,
230                                             &cbs_mpeg2_free_user_data);
231             START(0xb3, MPEG2RawSequenceHeader, sequence_header, NULL);
232             START(0xb5, MPEG2RawExtensionData,  extension_data,  NULL);
233             START(0xb8, MPEG2RawGroupOfPicturesHeader,
234                                        group_of_pictures_header, NULL);
235 #undef START
236         default:
237             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unknown start code %02"PRIx32".\n",
238                    unit->type);
239             return AVERROR_INVALIDDATA;
240         }
241     }
242
243     return 0;
244 }
245
246 static int cbs_mpeg2_write_header(CodedBitstreamContext *ctx,
247                                   CodedBitstreamUnit *unit,
248                                   PutBitContext *pbc)
249 {
250     int err;
251
252     switch (unit->type) {
253 #define START(start_code, type, func) \
254     case start_code: \
255         err = cbs_mpeg2_write_ ## func(ctx, pbc, unit->content); \
256         break;
257         START(0x00, MPEG2RawPictureHeader,  picture_header);
258         START(0xb2, MPEG2RawUserData,       user_data);
259         START(0xb3, MPEG2RawSequenceHeader, sequence_header);
260         START(0xb5, MPEG2RawExtensionData,  extension_data);
261         START(0xb8, MPEG2RawGroupOfPicturesHeader, group_of_pictures_header);
262 #undef START
263     default:
264         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for start "
265                "code %02"PRIx32".\n", unit->type);
266         return AVERROR_PATCHWELCOME;
267     }
268
269     return err;
270 }
271
272 static int cbs_mpeg2_write_slice(CodedBitstreamContext *ctx,
273                                  CodedBitstreamUnit *unit,
274                                  PutBitContext *pbc)
275 {
276     MPEG2RawSlice *slice = unit->content;
277     GetBitContext gbc;
278     size_t bits_left;
279     int err;
280
281     err = cbs_mpeg2_write_slice_header(ctx, pbc, &slice->header);
282     if (err < 0)
283         return err;
284
285     if (slice->data) {
286         if (slice->data_size * 8 + 8 > put_bits_left(pbc))
287             return AVERROR(ENOSPC);
288
289         init_get_bits(&gbc, slice->data, slice->data_size * 8);
290         skip_bits_long(&gbc, slice->data_bit_start);
291
292         while (get_bits_left(&gbc) > 15)
293             put_bits(pbc, 16, get_bits(&gbc, 16));
294
295         bits_left = get_bits_left(&gbc);
296         put_bits(pbc, bits_left, get_bits(&gbc, bits_left));
297
298         // Align with zeroes.
299         while (put_bits_count(pbc) % 8 != 0)
300             put_bits(pbc, 1, 0);
301     }
302
303     return 0;
304 }
305
306 static int cbs_mpeg2_write_unit(CodedBitstreamContext *ctx,
307                                 CodedBitstreamUnit *unit)
308 {
309     CodedBitstreamMPEG2Context *priv = ctx->priv_data;
310     PutBitContext pbc;
311     int err;
312
313     if (!priv->write_buffer) {
314         // Initial write buffer size is 1MB.
315         priv->write_buffer_size = 1024 * 1024;
316
317     reallocate_and_try_again:
318         err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
319         if (err < 0) {
320             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
321                    "sufficiently large write buffer (last attempt "
322                    "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
323             return err;
324         }
325     }
326
327     init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
328
329     if (unit->type >= 0x01 && unit->type <= 0xaf)
330         err = cbs_mpeg2_write_slice(ctx, unit, &pbc);
331     else
332         err = cbs_mpeg2_write_header(ctx, unit, &pbc);
333
334     if (err == AVERROR(ENOSPC)) {
335         // Overflow.
336         priv->write_buffer_size *= 2;
337         goto reallocate_and_try_again;
338     }
339     if (err < 0) {
340         // Write failed for some other reason.
341         return err;
342     }
343
344     if (put_bits_count(&pbc) % 8)
345         unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
346     else
347         unit->data_bit_padding = 0;
348
349     unit->data_size = (put_bits_count(&pbc) + 7) / 8;
350     flush_put_bits(&pbc);
351
352     err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
353     if (err < 0)
354         return err;
355
356     memcpy(unit->data, priv->write_buffer, unit->data_size);
357
358     return 0;
359 }
360
361 static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
362                                        CodedBitstreamFragment *frag)
363 {
364     uint8_t *data;
365     size_t size, dp, sp;
366     int i;
367
368     size = 0;
369     for (i = 0; i < frag->nb_units; i++)
370         size += 3 + frag->units[i].data_size;
371
372     frag->data_ref = av_buffer_alloc(size);
373     if (!frag->data_ref)
374         return AVERROR(ENOMEM);
375     data = frag->data_ref->data;
376
377     dp = 0;
378     for (i = 0; i < frag->nb_units; i++) {
379         CodedBitstreamUnit *unit = &frag->units[i];
380
381         data[dp++] = 0;
382         data[dp++] = 0;
383         data[dp++] = 1;
384
385         for (sp = 0; sp < unit->data_size; sp++)
386             data[dp++] = unit->data[sp];
387     }
388
389     av_assert0(dp == size);
390
391     frag->data      = data;
392     frag->data_size = size;
393
394     return 0;
395 }
396
397 static void cbs_mpeg2_close(CodedBitstreamContext *ctx)
398 {
399     CodedBitstreamMPEG2Context *priv = ctx->priv_data;
400
401     av_freep(&priv->write_buffer);
402 }
403
404 const CodedBitstreamType ff_cbs_type_mpeg2 = {
405     .codec_id          = AV_CODEC_ID_MPEG2VIDEO,
406
407     .priv_data_size    = sizeof(CodedBitstreamMPEG2Context),
408
409     .split_fragment    = &cbs_mpeg2_split_fragment,
410     .read_unit         = &cbs_mpeg2_read_unit,
411     .write_unit        = &cbs_mpeg2_write_unit,
412     .assemble_fragment = &cbs_mpeg2_assemble_fragment,
413
414     .close             = &cbs_mpeg2_close,
415 };