avcodec: add a get_encode_buffer() callback to AVCodecContext

This callback is functionally the same as get_buffer2() is for decoders, and
implements for the new encode API the functionality of the old encode API had
where the user could provide their own buffers.

Reviewed-by: Lynne <dev@lynne.ee>
Reviewed-by: Michael Niedermayer <michael@niedermayer.cc>
Reviewed-by: Mark Thompson <sw@jkqxz.net>
Signed-off-by: James Almer <jamrial@gmail.com>

6 files changed, 124 insertions, 4 deletions

diff --git a/libavcodec/avcodec.h b/libavcodec/avcodec.h
index 3ecb3a4a60..fbd4804160 100644
--- a/libavcodec/avcodec.h
+++ b/libavcodec/avcodec.h
@@ -513,6 +513,11 @@ typedef struct AVProducerReferenceTime {
  */
 #define AV_GET_BUFFER_FLAG_REF (1 << 0)
 
+/**
+ * The encoder will keep a reference to the packet and may reuse it later.
+ */
+#define AV_GET_ENCODE_BUFFER_FLAG_REF (1 << 0)
+
 struct AVCodecInternal;
 
 /**
@@ -2339,6 +2344,44 @@ typedef struct AVCodecContext {
      * - encoding: set by user
      */
     int export_side_data;
+
+    /**
+     * This callback is called at the beginning of each packet to get a data
+     * buffer for it.
+     *
+     * The following field will be set in the packet before this callback is
+     * called:
+     * - size
+     * This callback must use the above value to calculate the required buffer size,
+     * which must padded by at least AV_INPUT_BUFFER_PADDING_SIZE bytes.
+     *
+     * This callback must fill the following fields in the packet:
+     * - data: alignment requirements for AVPacket apply, if any. Some architectures and
+     *   encoders may benefit from having aligned data.
+     * - buf: must contain a pointer to an AVBufferRef structure. The packet's
+     *   data pointer must be contained in it. See: av_buffer_create(), av_buffer_alloc(),
+     *   and av_buffer_ref().
+     *
+     * If AV_CODEC_CAP_DR1 is not set then get_encode_buffer() must call
+     * avcodec_default_get_encode_buffer() instead of providing a buffer allocated by
+     * some other means.
+     *
+     * The flags field may contain a combination of AV_GET_ENCODE_BUFFER_FLAG_ flags.
+     * They may be used for example to hint what use the buffer may get after being
+     * created.
+     * Implementations of this callback may ignore flags they don't understand.
+     * If AV_GET_ENCODE_BUFFER_FLAG_REF is set in flags then the packet may be reused
+     * (read and/or written to if it is writable) later by libavcodec.
+     *
+     * This callback must be thread-safe, as when frame threading is used, it may
+     * be called from multiple threads simultaneously.
+     *
+     * @see avcodec_default_get_encode_buffer()
+     *
+     * - encoding: Set by libavcodec, user can override.
+     * - decoding: unused
+     */
+    int (*get_encode_buffer)(struct AVCodecContext *s, AVPacket *pkt, int flags);
 } AVCodecContext;
 
 #if FF_API_CODEC_GET_SET
@@ -2899,6 +2942,13 @@ void avsubtitle_free(AVSubtitle *sub);
 int avcodec_default_get_buffer2(AVCodecContext *s, AVFrame *frame, int flags);
 
 /**
+ * The default callback for AVCodecContext.get_encode_buffer(). It is made public so
+ * it can be called by custom get_encode_buffer() implementations for encoders without
+ * AV_CODEC_CAP_DR1 set.
+ */
+int avcodec_default_get_encode_buffer(AVCodecContext *s, AVPacket *pkt, int flags);
+
+/**
  * Modify width and height values so that they will result in a memory
  * buffer that is acceptable for the codec if you do not use any horizontal
  * padding.
diff --git a/libavcodec/codec.h b/libavcodec/codec.h
index 729df0e304..1bb70260ac 100644
--- a/libavcodec/codec.h
+++ b/libavcodec/codec.h
@@ -43,9 +43,11 @@
  */
 #define AV_CODEC_CAP_DRAW_HORIZ_BAND     (1 <<  0)
 /**
- * Codec uses get_buffer() for allocating buffers and supports custom allocators.
- * If not set, it might not use get_buffer() at all or use operations that
- * assume the buffer was allocated by avcodec_default_get_buffer.
+ * Codec uses get_buffer() or get_encode_buffer() for allocating buffers and
+ * supports custom allocators.
+ * If not set, it might not use get_buffer() or get_encode_buffer() at all, or
+ * use operations that assume the buffer was allocated by
+ * avcodec_default_get_buffer2 or avcodec_default_get_encode_buffer.
  */
 #define AV_CODEC_CAP_DR1                 (1 <<  1)
 #define AV_CODEC_CAP_TRUNCATED           (1 <<  3)
diff --git a/libavcodec/encode.c b/libavcodec/encode.c
index 282337e453..abec818eb9 100644
--- a/libavcodec/encode.c
+++ b/libavcodec/encode.c
@@ -56,6 +56,59 @@ int ff_alloc_packet2(AVCodecContext *avctx, AVPacket *avpkt, int64_t size, int64
     return 0;
 }
 
+int avcodec_default_get_encode_buffer(AVCodecContext *avctx, AVPacket *avpkt, int flags)
+{
+    int ret;
+
+    if (avpkt->size < 0 || avpkt->size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE)
+        return AVERROR(EINVAL);
+
+    if (avpkt->data || avpkt->buf) {
+        av_log(avctx, AV_LOG_ERROR, "avpkt->{data,buf} != NULL in avcodec_default_get_encode_buffer()\n");
+        return AVERROR(EINVAL);
+    }
+
+    ret = av_buffer_realloc(&avpkt->buf, avpkt->size + AV_INPUT_BUFFER_PADDING_SIZE);
+    if (ret < 0) {
+        av_log(avctx, AV_LOG_ERROR, "Failed to allocate packet of size %d\n", avpkt->size);
+        return ret;
+    }
+    avpkt->data = avpkt->buf->data;
+    memset(avpkt->data + avpkt->size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
+
+    return 0;
+}
+
+int ff_get_encode_buffer(AVCodecContext *avctx, AVPacket *avpkt, int64_t size, int flags)
+{
+    int ret;
+
+    if (size < 0 || size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE)
+        return AVERROR(EINVAL);
+
+    av_assert0(!avpkt->data && !avpkt->buf);
+
+    avpkt->size = size;
+    ret = avctx->get_encode_buffer(avctx, avpkt, flags);
+    if (ret < 0)
+        goto fail;
+
+    if (!avpkt->data || !avpkt->buf) {
+        av_log(avctx, AV_LOG_ERROR, "No buffer returned by get_encode_buffer()\n");
+        ret = AVERROR(EINVAL);
+        goto fail;
+    }
+
+    ret = 0;
+fail:
+    if (ret < 0) {
+        av_log(avctx, AV_LOG_ERROR, "get_encode_buffer() failed\n");
+        av_packet_unref(avpkt);
+    }
+
+    return ret;
+}
+
 /**
  * Pad last frame with silence.
  */
@@ -377,6 +430,12 @@ static int compat_encode(AVCodecContext *avctx, AVPacket *avpkt,
             av_log(avctx, AV_LOG_WARNING, "AVFrame.width or height is not set\n");
     }
 
+    if (avctx->codec->capabilities & AV_CODEC_CAP_DR1) {
+        av_log(avctx, AV_LOG_WARNING, "The deprecated avcodec_encode_* API does not support "
+                                      "AV_CODEC_CAP_DR1 encoders\n");
+        return AVERROR(ENOSYS);
+    }
+
     ret = avcodec_send_frame(avctx, frame);
     if (ret == AVERROR_EOF)
         ret = 0;
diff --git a/libavcodec/encode.h b/libavcodec/encode.h
index dfa9cb2d97..089c31091c 100644
--- a/libavcodec/encode.h
+++ b/libavcodec/encode.h
@@ -24,6 +24,7 @@
 #include "libavutil/frame.h"
 
 #include "avcodec.h"
+#include "packet.h"
 
 /**
  * Called by encoders to get the next frame for encoding.
@@ -36,4 +37,11 @@
  */
 int ff_encode_get_frame(AVCodecContext *avctx, AVFrame *frame);
 
+/**
+ * Get a buffer for a packet. This is a wrapper around
+ * AVCodecContext.get_encode_buffer() and should be used instead calling get_encode_buffer()
+ * directly.
+ */
+int ff_get_encode_buffer(AVCodecContext *avctx, AVPacket *avpkt, int64_t size, int flags);
+
 #endif /* AVCODEC_ENCODE_H */
diff --git a/libavcodec/options.c b/libavcodec/options.c
index 6e904b61d8..61689b48d9 100644
--- a/libavcodec/options.c
+++ b/libavcodec/options.c
@@ -130,6 +130,7 @@ static int init_context_defaults(AVCodecContext *s, const AVCodec *codec)
     s->pkt_timebase        = (AVRational){ 0, 1 };
     s->get_buffer2         = avcodec_default_get_buffer2;
     s->get_format          = avcodec_default_get_format;
+    s->get_encode_buffer   = avcodec_default_get_encode_buffer;
     s->execute             = avcodec_default_execute;
     s->execute2            = avcodec_default_execute2;
     s->sample_aspect_ratio = (AVRational){0,1};
diff --git a/libavcodec/version.h b/libavcodec/version.h
index 3124ec8061..4299ad4239 100644
--- a/libavcodec/version.h
+++ b/libavcodec/version.h
@@ -28,7 +28,7 @@
 #include "libavutil/version.h"
 
 #define LIBAVCODEC_VERSION_MAJOR  58
-#define LIBAVCODEC_VERSION_MINOR 130
+#define LIBAVCODEC_VERSION_MINOR 131
 #define LIBAVCODEC_VERSION_MICRO 100
 
 #define LIBAVCODEC_VERSION_INT  AV_VERSION_INT(LIBAVCODEC_VERSION_MAJOR, \