avcodec/cbs: Factor out common code for writing units

All cbs-functions to write units share a common pattern:
1. They check whether they have a write buffer (that is used to store
the unit's data until the needed size becomes known after writing the
unit when a dedicated buffer will be allocated).
2. They use this buffer for a PutBitContext.
3. The (codec-specific) writing takes place through the PutBitContext.
4. The return value is checked. AVERROR(ENOSPC) here always indicates
that the buffer was too small and leads to a reallocation of said
buffer.
5. The final buffer will be allocated and the data copied.

This commit factors this common code out in a single function in cbs.c.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@gmail.com>
This commit is contained in:
Andreas Rheinhardt 2019-11-17 08:34:35 +01:00 committed by Mark Thompson
parent 5d8d9e032c
commit 7c92eaace2
13 changed files with 94 additions and 322 deletions

View File

@ -95,10 +95,12 @@ int ff_cbs_init(CodedBitstreamContext **ctx_ptr,
ctx->log_ctx = log_ctx;
ctx->codec = type;
ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
if (!ctx->priv_data) {
av_freep(&ctx);
return AVERROR(ENOMEM);
if (type->priv_data_size) {
ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
if (!ctx->priv_data) {
av_freep(&ctx);
return AVERROR(ENOMEM);
}
}
ctx->decompose_unit_types = NULL;
@ -120,6 +122,7 @@ void ff_cbs_close(CodedBitstreamContext **ctx_ptr)
if (ctx->codec && ctx->codec->close)
ctx->codec->close(ctx);
av_freep(&ctx->write_buffer);
av_freep(&ctx->priv_data);
av_freep(ctx_ptr);
}
@ -280,6 +283,57 @@ int ff_cbs_read(CodedBitstreamContext *ctx,
return cbs_read_fragment_content(ctx, frag);
}
static int cbs_write_unit_data(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit)
{
PutBitContext pbc;
int ret;
if (!ctx->write_buffer) {
// Initial write buffer size is 1MB.
ctx->write_buffer_size = 1024 * 1024;
reallocate_and_try_again:
ret = av_reallocp(&ctx->write_buffer, ctx->write_buffer_size);
if (ret < 0) {
av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
"sufficiently large write buffer (last attempt "
"%"SIZE_SPECIFIER" bytes).\n", ctx->write_buffer_size);
return ret;
}
}
init_put_bits(&pbc, ctx->write_buffer, ctx->write_buffer_size);
ret = ctx->codec->write_unit(ctx, unit, &pbc);
if (ret < 0) {
if (ret == AVERROR(ENOSPC)) {
// Overflow.
ctx->write_buffer_size *= 2;
goto reallocate_and_try_again;
}
// Write failed for some other reason.
return ret;
}
// Overflow but we didn't notice.
av_assert0(put_bits_count(&pbc) <= 8 * ctx->write_buffer_size);
if (put_bits_count(&pbc) % 8)
unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
else
unit->data_bit_padding = 0;
flush_put_bits(&pbc);
ret = ff_cbs_alloc_unit_data(ctx, unit, put_bits_count(&pbc) / 8);
if (ret < 0)
return ret;
memcpy(unit->data, ctx->write_buffer, unit->data_size);
return 0;
}
int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
CodedBitstreamFragment *frag)
@ -295,7 +349,7 @@ int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
av_buffer_unref(&unit->data_ref);
unit->data = NULL;
err = ctx->codec->write_unit(ctx, unit);
err = cbs_write_unit_data(ctx, unit);
if (err < 0) {
av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to write unit %d "
"(type %"PRIu32").\n", i, unit->type);

View File

@ -210,6 +210,13 @@ typedef struct CodedBitstreamContext {
* From AV_LOG_*; defaults to AV_LOG_TRACE.
*/
int trace_level;
/**
* Write buffer. Used as intermediate buffer when writing units.
* For internal use of cbs only.
*/
uint8_t *write_buffer;
size_t write_buffer_size;
} CodedBitstreamContext;

View File

@ -1199,66 +1199,19 @@ static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
return AVERROR(ENOSPC);
if (obu->obu_size > 0) {
memmove(priv->write_buffer + data_pos,
priv->write_buffer + start_pos, header_size);
memmove(pbc->buf + data_pos,
pbc->buf + start_pos, header_size);
skip_put_bytes(pbc, header_size);
if (td) {
memcpy(priv->write_buffer + data_pos + header_size,
memcpy(pbc->buf + data_pos + header_size,
td->data, td->data_size);
skip_put_bytes(pbc, td->data_size);
}
}
return 0;
}
static int cbs_av1_write_unit(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit)
{
CodedBitstreamAV1Context *priv = ctx->priv_data;
PutBitContext pbc;
int err;
if (!priv->write_buffer) {
// Initial write buffer size is 1MB.
priv->write_buffer_size = 1024 * 1024;
reallocate_and_try_again:
err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
if (err < 0) {
av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
"sufficiently large write buffer (last attempt "
"%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
return err;
}
}
init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
err = cbs_av1_write_obu(ctx, unit, &pbc);
if (err == AVERROR(ENOSPC)) {
// Overflow.
priv->write_buffer_size *= 2;
goto reallocate_and_try_again;
}
if (err < 0)
return err;
// Overflow but we didn't notice.
av_assert0(put_bits_count(&pbc) <= 8 * priv->write_buffer_size);
// OBU data must be byte-aligned.
av_assert0(put_bits_count(&pbc) % 8 == 0);
unit->data_size = put_bits_count(&pbc) / 8;
flush_put_bits(&pbc);
err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
if (err < 0)
return err;
memcpy(unit->data, priv->write_buffer, unit->data_size);
av_assert0(put_bits_count(pbc) % 8 == 0);
return 0;
}
@ -1297,8 +1250,6 @@ static void cbs_av1_close(CodedBitstreamContext *ctx)
av_buffer_unref(&priv->sequence_header_ref);
av_buffer_unref(&priv->frame_header_ref);
av_freep(&priv->write_buffer);
}
const CodedBitstreamType ff_cbs_type_av1 = {
@ -1308,7 +1259,7 @@ const CodedBitstreamType ff_cbs_type_av1 = {
.split_fragment = &cbs_av1_split_fragment,
.read_unit = &cbs_av1_read_unit,
.write_unit = &cbs_av1_write_unit,
.write_unit = &cbs_av1_write_obu,
.assemble_fragment = &cbs_av1_assemble_fragment,
.close = &cbs_av1_close,

View File

@ -442,10 +442,6 @@ typedef struct CodedBitstreamAV1Context {
int tile_rows;
AV1ReferenceFrameState ref[AV1_NUM_REF_FRAMES];
// Write buffer.
uint8_t *write_buffer;
size_t write_buffer_size;
} CodedBitstreamAV1Context;

View File

@ -1380,65 +1380,6 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
return 0;
}
static int cbs_h2645_write_nal_unit(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit)
{
CodedBitstreamH2645Context *priv = ctx->priv_data;
enum AVCodecID codec_id = ctx->codec->codec_id;
PutBitContext pbc;
int err;
if (!priv->write_buffer) {
// Initial write buffer size is 1MB.
priv->write_buffer_size = 1024 * 1024;
reallocate_and_try_again:
err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
if (err < 0) {
av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
"sufficiently large write buffer (last attempt "
"%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
return err;
}
}
init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
if (codec_id == AV_CODEC_ID_H264)
err = cbs_h264_write_nal_unit(ctx, unit, &pbc);
else
err = cbs_h265_write_nal_unit(ctx, unit, &pbc);
if (err == AVERROR(ENOSPC)) {
// Overflow.
priv->write_buffer_size *= 2;
goto reallocate_and_try_again;
}
// Overflow but we didn't notice.
av_assert0(put_bits_count(&pbc) <= 8 * priv->write_buffer_size);
if (err < 0) {
// Write failed for some other reason.
return err;
}
if (put_bits_count(&pbc) % 8)
unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
else
unit->data_bit_padding = 0;
unit->data_size = (put_bits_count(&pbc) + 7) / 8;
flush_put_bits(&pbc);
err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
if (err < 0)
return err;
memcpy(unit->data, priv->write_buffer, unit->data_size);
return 0;
}
static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
CodedBitstreamFragment *frag)
{
@ -1533,8 +1474,6 @@ static void cbs_h264_close(CodedBitstreamContext *ctx)
ff_h2645_packet_uninit(&h264->common.read_packet);
av_freep(&h264->common.write_buffer);
for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
av_buffer_unref(&h264->sps_ref[i]);
for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
@ -1548,8 +1487,6 @@ static void cbs_h265_close(CodedBitstreamContext *ctx)
ff_h2645_packet_uninit(&h265->common.read_packet);
av_freep(&h265->common.write_buffer);
for (i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++)
av_buffer_unref(&h265->vps_ref[i]);
for (i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++)
@ -1565,7 +1502,7 @@ const CodedBitstreamType ff_cbs_type_h264 = {
.split_fragment = &cbs_h2645_split_fragment,
.read_unit = &cbs_h264_read_nal_unit,
.write_unit = &cbs_h2645_write_nal_unit,
.write_unit = &cbs_h264_write_nal_unit,
.assemble_fragment = &cbs_h2645_assemble_fragment,
.close = &cbs_h264_close,
@ -1578,7 +1515,7 @@ const CodedBitstreamType ff_cbs_type_h265 = {
.split_fragment = &cbs_h2645_split_fragment,
.read_unit = &cbs_h265_read_nal_unit,
.write_unit = &cbs_h2645_write_nal_unit,
.write_unit = &cbs_h265_write_nal_unit,
.assemble_fragment = &cbs_h2645_assemble_fragment,
.close = &cbs_h265_close,

View File

@ -19,9 +19,6 @@
#ifndef AVCODEC_CBS_H2645_H
#define AVCODEC_CBS_H2645_H
#include <stddef.h>
#include <stdint.h>
#include "h2645_parse.h"
@ -33,10 +30,6 @@ typedef struct CodedBitstreamH2645Context {
int nal_length_size;
// Packet reader.
H2645Packet read_packet;
// Write buffer
uint8_t *write_buffer;
size_t write_buffer_size;
} CodedBitstreamH2645Context;

View File

@ -44,9 +44,11 @@ typedef struct CodedBitstreamType {
int (*read_unit)(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit);
// Write the unit->data bitstream from unit->content.
// Write the data bitstream from unit->content into pbc.
// Return value AVERROR(ENOSPC) indicates that pbc was too small.
int (*write_unit)(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit);
CodedBitstreamUnit *unit,
PutBitContext *pbc);
// Read the data from all of frag->units and assemble it into
// a bitstream for the whole fragment.

View File

@ -377,58 +377,13 @@ static int cbs_jpeg_write_segment(CodedBitstreamContext *ctx,
}
static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit)
CodedBitstreamUnit *unit,
PutBitContext *pbc)
{
CodedBitstreamJPEGContext *priv = ctx->priv_data;
PutBitContext pbc;
int err;
if (!priv->write_buffer) {
// Initial write buffer size is 1MB.
priv->write_buffer_size = 1024 * 1024;
reallocate_and_try_again:
err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
if (err < 0) {
av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
"sufficiently large write buffer (last attempt "
"%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
return err;
}
}
init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
if (unit->type == JPEG_MARKER_SOS)
err = cbs_jpeg_write_scan(ctx, unit, &pbc);
return cbs_jpeg_write_scan (ctx, unit, pbc);
else
err = cbs_jpeg_write_segment(ctx, unit, &pbc);
if (err == AVERROR(ENOSPC)) {
// Overflow.
priv->write_buffer_size *= 2;
goto reallocate_and_try_again;
}
if (err < 0) {
// Write failed for some other reason.
return err;
}
if (put_bits_count(&pbc) % 8)
unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
else
unit->data_bit_padding = 0;
unit->data_size = (put_bits_count(&pbc) + 7) / 8;
flush_put_bits(&pbc);
err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
if (err < 0)
return err;
memcpy(unit->data, priv->write_buffer, unit->data_size);
return 0;
return cbs_jpeg_write_segment(ctx, unit, pbc);
}
static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
@ -499,22 +454,11 @@ static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
return 0;
}
static void cbs_jpeg_close(CodedBitstreamContext *ctx)
{
CodedBitstreamJPEGContext *priv = ctx->priv_data;
av_freep(&priv->write_buffer);
}
const CodedBitstreamType ff_cbs_type_jpeg = {
.codec_id = AV_CODEC_ID_MJPEG,
.priv_data_size = sizeof(CodedBitstreamJPEGContext),
.split_fragment = &cbs_jpeg_split_fragment,
.read_unit = &cbs_jpeg_read_unit,
.write_unit = &cbs_jpeg_write_unit,
.assemble_fragment = &cbs_jpeg_assemble_fragment,
.close = &cbs_jpeg_close,
};

View File

@ -120,11 +120,4 @@ typedef struct JPEGRawComment {
} JPEGRawComment;
typedef struct CodedBitstreamJPEGContext {
// Write buffer.
uint8_t *write_buffer;
size_t write_buffer_size;
} CodedBitstreamJPEGContext;
#endif /* AVCODEC_CBS_JPEG_H */

View File

@ -371,58 +371,13 @@ static int cbs_mpeg2_write_slice(CodedBitstreamContext *ctx,
}
static int cbs_mpeg2_write_unit(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit)
CodedBitstreamUnit *unit,
PutBitContext *pbc)
{
CodedBitstreamMPEG2Context *priv = ctx->priv_data;
PutBitContext pbc;
int err;
if (!priv->write_buffer) {
// Initial write buffer size is 1MB.
priv->write_buffer_size = 1024 * 1024;
reallocate_and_try_again:
err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
if (err < 0) {
av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
"sufficiently large write buffer (last attempt "
"%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
return err;
}
}
init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
if (MPEG2_START_IS_SLICE(unit->type))
err = cbs_mpeg2_write_slice(ctx, unit, &pbc);
return cbs_mpeg2_write_slice (ctx, unit, pbc);
else
err = cbs_mpeg2_write_header(ctx, unit, &pbc);
if (err == AVERROR(ENOSPC)) {
// Overflow.
priv->write_buffer_size *= 2;
goto reallocate_and_try_again;
}
if (err < 0) {
// Write failed for some other reason.
return err;
}
if (put_bits_count(&pbc) % 8)
unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
else
unit->data_bit_padding = 0;
unit->data_size = (put_bits_count(&pbc) + 7) / 8;
flush_put_bits(&pbc);
err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
if (err < 0)
return err;
memcpy(unit->data, priv->write_buffer, unit->data_size);
return 0;
return cbs_mpeg2_write_header(ctx, unit, pbc);
}
static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
@ -462,13 +417,6 @@ static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
return 0;
}
static void cbs_mpeg2_close(CodedBitstreamContext *ctx)
{
CodedBitstreamMPEG2Context *priv = ctx->priv_data;
av_freep(&priv->write_buffer);
}
const CodedBitstreamType ff_cbs_type_mpeg2 = {
.codec_id = AV_CODEC_ID_MPEG2VIDEO,
@ -478,6 +426,4 @@ const CodedBitstreamType ff_cbs_type_mpeg2 = {
.read_unit = &cbs_mpeg2_read_unit,
.write_unit = &cbs_mpeg2_write_unit,
.assemble_fragment = &cbs_mpeg2_assemble_fragment,
.close = &cbs_mpeg2_close,
};

View File

@ -225,10 +225,6 @@ typedef struct CodedBitstreamMPEG2Context {
uint8_t scalable_mode;
uint8_t progressive_sequence;
uint8_t number_of_frame_centre_offsets;
// Write buffer.
uint8_t *write_buffer;
size_t write_buffer_size;
} CodedBitstreamMPEG2Context;

View File

@ -522,62 +522,28 @@ static int cbs_vp9_read_unit(CodedBitstreamContext *ctx,
}
static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
CodedBitstreamUnit *unit)
CodedBitstreamUnit *unit,
PutBitContext *pbc)
{
CodedBitstreamVP9Context *priv = ctx->priv_data;
VP9RawFrame *frame = unit->content;
PutBitContext pbc;
int err;
if (!priv->write_buffer) {
// Initial write buffer size is 1MB.
priv->write_buffer_size = 1024 * 1024;
reallocate_and_try_again:
err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
if (err < 0) {
av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
"sufficiently large write buffer (last attempt "
"%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
return err;
}
}
init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
err = cbs_vp9_write_frame(ctx, &pbc, frame);
if (err == AVERROR(ENOSPC)) {
priv->write_buffer_size *= 2;
goto reallocate_and_try_again;
}
err = cbs_vp9_write_frame(ctx, pbc, frame);
if (err < 0)
return err;
// Frame must be byte-aligned.
av_assert0(put_bits_count(&pbc) % 8 == 0);
unit->data_size = put_bits_count(&pbc) / 8;
unit->data_bit_padding = 0;
flush_put_bits(&pbc);
av_assert0(put_bits_count(pbc) % 8 == 0);
if (frame->data) {
if (unit->data_size + frame->data_size >
priv->write_buffer_size) {
priv->write_buffer_size *= 2;
goto reallocate_and_try_again;
}
if (frame->data_size > put_bits_left(pbc) / 8)
return AVERROR(ENOSPC);
memcpy(priv->write_buffer + unit->data_size,
frame->data, frame->data_size);
unit->data_size += frame->data_size;
flush_put_bits(pbc);
memcpy(put_bits_ptr(pbc), frame->data, frame->data_size);
skip_put_bytes(pbc, frame->data_size);
}
err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
if (err < 0)
return err;
memcpy(unit->data, priv->write_buffer, unit->data_size);
return 0;
}
@ -671,13 +637,6 @@ static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
return 0;
}
static void cbs_vp9_close(CodedBitstreamContext *ctx)
{
CodedBitstreamVP9Context *priv = ctx->priv_data;
av_freep(&priv->write_buffer);
}
const CodedBitstreamType ff_cbs_type_vp9 = {
.codec_id = AV_CODEC_ID_VP9,
@ -687,6 +646,4 @@ const CodedBitstreamType ff_cbs_type_vp9 = {
.read_unit = &cbs_vp9_read_unit,
.write_unit = &cbs_vp9_write_unit,
.assemble_fragment = &cbs_vp9_assemble_fragment,
.close = &cbs_vp9_close,
};

View File

@ -207,10 +207,6 @@ typedef struct CodedBitstreamVP9Context {
int bit_depth;
VP9ReferenceFrameState ref[VP9_NUM_REF_FRAMES];
// Write buffer.
uint8_t *write_buffer;
size_t write_buffer_size;
} CodedBitstreamVP9Context;