lavfi/framepool: add audio support

1 files changed, 109 insertions, 0 deletions

diff --git a/libavfilter/framepool.c b/libavfilter/framepool.c
index 36c6e8fd04..588dcb6da2 100644
--- a/libavfilter/framepool.c
+++ b/libavfilter/framepool.c
@@ -20,6 +20,7 @@
 
 #include "framepool.h"
 #include "libavutil/avassert.h"
+#include "libavutil/avutil.h"
 #include "libavutil/buffer.h"
 #include "libavutil/frame.h"
 #include "libavutil/imgutils.h"
@@ -28,8 +29,18 @@
 
 struct FFFramePool {
 
+    enum AVMediaType type;
+
+    /* video */
     int width;
     int height;
+
+    /* audio */
+    int planes;
+    int channels;
+    int nb_samples;
+
+    /* common */
     int format;
     int align;
     int linesize[4];
@@ -54,6 +65,7 @@ FFFramePool *ff_frame_pool_video_init(AVBufferRef* (*alloc)(int size),
     if (!pool)
         return NULL;
 
+    pool->type = AVMEDIA_TYPE_VIDEO;
     pool->width = width;
     pool->height = height;
     pool->format = format;
@@ -104,6 +116,44 @@ fail:
     return NULL;
 }
 
+FFFramePool *ff_frame_pool_audio_init(AVBufferRef* (*alloc)(int size),
+                                      int channels,
+                                      int nb_samples,
+                                      enum AVSampleFormat format,
+                                      int align)
+{
+    int ret, planar;
+    FFFramePool *pool;
+
+    pool = av_mallocz(sizeof(FFFramePool));
+    if (!pool)
+        return NULL;
+
+    planar = av_sample_fmt_is_planar(format);
+
+    pool->type = AVMEDIA_TYPE_AUDIO;
+    pool->planes = planar ? channels : 1;
+    pool->channels = channels;
+    pool->nb_samples = nb_samples;
+    pool->format = format;
+    pool->align = align;
+
+    ret = av_samples_get_buffer_size(&pool->linesize[0], channels,
+                                     nb_samples, format, 0);
+    if (ret < 0)
+        goto fail;
+
+    pool->pools[0] = av_buffer_pool_init(pool->linesize[0], NULL);
+    if (!pool->pools[0])
+        goto fail;
+
+    return pool;
+
+fail:
+    ff_frame_pool_uninit(&pool);
+    return NULL;
+}
+
 int ff_frame_pool_get_video_config(FFFramePool *pool,
                                    int *width,
                                    int *height,
@@ -113,6 +163,8 @@ int ff_frame_pool_get_video_config(FFFramePool *pool,
     if (!pool)
         return AVERROR(EINVAL);
 
+    av_assert0(pool->type == AVMEDIA_TYPE_VIDEO);
+
     *width = pool->width;
     *height = pool->height;
     *format = pool->format;
@@ -121,6 +173,24 @@ int ff_frame_pool_get_video_config(FFFramePool *pool,
     return 0;
 }
 
+int ff_frame_pool_get_audio_config(FFFramePool *pool,
+                                   int *channels,
+                                   int *nb_samples,
+                                   enum AVSampleFormat *format,
+                                   int *align)
+{
+    if (!pool)
+        return AVERROR(EINVAL);
+
+    av_assert0(pool->type == AVMEDIA_TYPE_AUDIO);
+
+    *channels = pool->channels;
+    *nb_samples = pool->nb_samples;
+    *format = pool->format;
+    *align = pool->align;
+
+    return 0;
+}
 
 AVFrame *ff_frame_pool_get(FFFramePool *pool)
 {
@@ -133,6 +203,8 @@ AVFrame *ff_frame_pool_get(FFFramePool *pool)
         return NULL;
     }
 
+    switch(pool->type) {
+    case AVMEDIA_TYPE_VIDEO:
     desc = av_pix_fmt_desc_get(pool->format);
     if (!desc) {
         goto fail;
@@ -167,6 +239,43 @@ AVFrame *ff_frame_pool_get(FFFramePool *pool)
     }
 
     frame->extended_data = frame->data;
+    break;
+    case AVMEDIA_TYPE_AUDIO:
+        frame->nb_samples = pool->nb_samples;
+        av_frame_set_channels(frame, pool->channels);
+        frame->format = pool->format;
+        frame->linesize[0] = pool->linesize[0];
+
+        if (pool->planes > AV_NUM_DATA_POINTERS) {
+            frame->extended_data = av_mallocz_array(pool->planes,
+                                                    sizeof(*frame->extended_data));
+            frame->nb_extended_buf = pool->planes - AV_NUM_DATA_POINTERS;
+            frame->extended_buf = av_mallocz_array(frame->nb_extended_buf,
+                                                   sizeof(*frame->extended_buf));
+            if (!frame->extended_data || !frame->extended_buf)
+                goto fail;
+        } else {
+            frame->extended_data = frame->data;
+            av_assert0(frame->nb_extended_buf == 0);
+        }
+
+        for (i = 0; i < FFMIN(pool->planes, AV_NUM_DATA_POINTERS); i++) {
+            frame->buf[i] = av_buffer_pool_get(pool->pools[0]);
+            if (!frame->buf[i])
+                goto fail;
+            frame->extended_data[i] = frame->data[i] = frame->buf[i]->data;
+        }
+        for (i = 0; i < frame->nb_extended_buf; i++) {
+            frame->extended_buf[i] = av_buffer_pool_get(pool->pools[0]);
+            if (!frame->extended_buf[i])
+                goto fail;
+            frame->extended_data[i + AV_NUM_DATA_POINTERS] = frame->extended_buf[i]->data;
+        }
+
+        break;
+    default:
+        av_assert0(0);
+    }
 
     return frame;
 fail: