1 /*
2  * This file is part of mpv.
3  *
4  * mpv 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  * mpv 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
12  * GNU 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 mpv.  If not, see <http://www.gnu.org/licenses/>.
16  */
17 
18 #include <libavutil/frame.h>
19 #include <libavutil/mem.h>
20 
21 #include "common/common.h"
22 
23 #include "chmap.h"
24 #include "fmt-conversion.h"
25 #include "format.h"
26 #include "aframe.h"
27 
28 struct mp_aframe {
29     AVFrame *av_frame;
30     // We support channel layouts different from AVFrame channel masks
31     struct mp_chmap chmap;
32     // We support spdif formats, which are allocated as AV_SAMPLE_FMT_S16.
33     int format;
34     double pts;
35     double speed;
36 };
37 
38 struct avframe_opaque {
39     double speed;
40 };
41 
free_frame(void * ptr)42 static void free_frame(void *ptr)
43 {
44     struct mp_aframe *frame = ptr;
45     av_frame_free(&frame->av_frame);
46 }
47 
mp_aframe_create(void)48 struct mp_aframe *mp_aframe_create(void)
49 {
50     struct mp_aframe *frame = talloc_zero(NULL, struct mp_aframe);
51     frame->av_frame = av_frame_alloc();
52     if (!frame->av_frame)
53         abort();
54     talloc_set_destructor(frame, free_frame);
55     mp_aframe_reset(frame);
56     return frame;
57 }
58 
mp_aframe_new_ref(struct mp_aframe * frame)59 struct mp_aframe *mp_aframe_new_ref(struct mp_aframe *frame)
60 {
61     if (!frame)
62         return NULL;
63 
64     struct mp_aframe *dst = mp_aframe_create();
65 
66     dst->chmap = frame->chmap;
67     dst->format = frame->format;
68     dst->pts = frame->pts;
69     dst->speed = frame->speed;
70 
71     if (mp_aframe_is_allocated(frame)) {
72         if (av_frame_ref(dst->av_frame, frame->av_frame) < 0)
73             abort();
74     } else {
75         // av_frame_ref() would fail.
76         mp_aframe_config_copy(dst, frame);
77     }
78 
79     return dst;
80 }
81 
82 // Revert to state after mp_aframe_create().
mp_aframe_reset(struct mp_aframe * frame)83 void mp_aframe_reset(struct mp_aframe *frame)
84 {
85     av_frame_unref(frame->av_frame);
86     frame->chmap.num = 0;
87     frame->format = 0;
88     frame->pts = MP_NOPTS_VALUE;
89     frame->speed = 1.0;
90 }
91 
92 // Remove all actual audio data and leave only the metadata.
mp_aframe_unref_data(struct mp_aframe * frame)93 void mp_aframe_unref_data(struct mp_aframe *frame)
94 {
95     // In a fucked up way, this is less complex than just unreffing the data.
96     struct mp_aframe *tmp = mp_aframe_create();
97     MPSWAP(struct mp_aframe, *tmp, *frame);
98     mp_aframe_reset(frame);
99     mp_aframe_config_copy(frame, tmp);
100     talloc_free(tmp);
101 }
102 
103 // Allocate this much data. Returns false for failure (data already allocated,
104 // invalid sample count or format, allocation failures).
105 // Normally you're supposed to use a frame pool and mp_aframe_pool_allocate().
mp_aframe_alloc_data(struct mp_aframe * frame,int samples)106 bool mp_aframe_alloc_data(struct mp_aframe *frame, int samples)
107 {
108     if (mp_aframe_is_allocated(frame))
109         return false;
110     struct mp_aframe_pool *p = mp_aframe_pool_create(NULL);
111     int r = mp_aframe_pool_allocate(p, frame, samples);
112     talloc_free(p);
113     return r >= 0;
114 }
115 
116 // Return a new reference to the data in av_frame. av_frame itself is not
117 // touched. Returns NULL if not representable, or if input is NULL.
118 // Does not copy the timestamps.
mp_aframe_from_avframe(struct AVFrame * av_frame)119 struct mp_aframe *mp_aframe_from_avframe(struct AVFrame *av_frame)
120 {
121     if (!av_frame || av_frame->width > 0 || av_frame->height > 0)
122         return NULL;
123 
124     int format = af_from_avformat(av_frame->format);
125     if (!format && av_frame->format != AV_SAMPLE_FMT_NONE)
126         return NULL;
127 
128     struct mp_aframe *frame = mp_aframe_create();
129 
130     // This also takes care of forcing refcounting.
131     if (av_frame_ref(frame->av_frame, av_frame) < 0)
132         abort();
133 
134     frame->format = format;
135     mp_chmap_from_lavc(&frame->chmap, frame->av_frame->channel_layout);
136 
137     // FFmpeg being a stupid POS again
138     if (frame->chmap.num != frame->av_frame->channels)
139         mp_chmap_from_channels(&frame->chmap, av_frame->channels);
140 
141     if (av_frame->opaque_ref) {
142         struct avframe_opaque *op = (void *)av_frame->opaque_ref->data;
143         frame->speed = op->speed;
144     }
145 
146     return frame;
147 }
148 
149 // Return a new reference to the data in frame. Returns NULL is not
150 // representable (), or if input is NULL.
151 // Does not copy the timestamps.
mp_aframe_to_avframe(struct mp_aframe * frame)152 struct AVFrame *mp_aframe_to_avframe(struct mp_aframe *frame)
153 {
154     if (!frame)
155         return NULL;
156 
157     if (af_to_avformat(frame->format) != frame->av_frame->format)
158         return NULL;
159 
160     if (!mp_chmap_is_lavc(&frame->chmap))
161         return NULL;
162 
163     if (!frame->av_frame->opaque_ref && frame->speed != 1.0) {
164         frame->av_frame->opaque_ref =
165             av_buffer_alloc(sizeof(struct avframe_opaque));
166         if (!frame->av_frame->opaque_ref)
167             return NULL;
168 
169         struct avframe_opaque *op = (void *)frame->av_frame->opaque_ref->data;
170         op->speed = frame->speed;
171     }
172 
173     return av_frame_clone(frame->av_frame);
174 }
175 
mp_aframe_to_avframe_and_unref(struct mp_aframe * frame)176 struct AVFrame *mp_aframe_to_avframe_and_unref(struct mp_aframe *frame)
177 {
178     AVFrame *av = mp_aframe_to_avframe(frame);
179     talloc_free(frame);
180     return av;
181 }
182 
183 // You must not use this.
mp_aframe_get_raw_avframe(struct mp_aframe * frame)184 struct AVFrame *mp_aframe_get_raw_avframe(struct mp_aframe *frame)
185 {
186     return frame->av_frame;
187 }
188 
189 // Return whether it has associated audio data. (If not, metadata only.)
mp_aframe_is_allocated(struct mp_aframe * frame)190 bool mp_aframe_is_allocated(struct mp_aframe *frame)
191 {
192     return frame->av_frame->buf[0] || frame->av_frame->extended_data[0];
193 }
194 
195 // Clear dst, and then copy the configuration to it.
mp_aframe_config_copy(struct mp_aframe * dst,struct mp_aframe * src)196 void mp_aframe_config_copy(struct mp_aframe *dst, struct mp_aframe *src)
197 {
198     mp_aframe_reset(dst);
199 
200     dst->chmap = src->chmap;
201     dst->format = src->format;
202 
203     mp_aframe_copy_attributes(dst, src);
204 
205     dst->av_frame->sample_rate = src->av_frame->sample_rate;
206     dst->av_frame->format = src->av_frame->format;
207     dst->av_frame->channel_layout = src->av_frame->channel_layout;
208     // FFmpeg being a stupid POS again
209     dst->av_frame->channels = src->av_frame->channels;
210 }
211 
212 // Copy "soft" attributes from src to dst, excluding things which affect
213 // frame allocation and organization.
mp_aframe_copy_attributes(struct mp_aframe * dst,struct mp_aframe * src)214 void mp_aframe_copy_attributes(struct mp_aframe *dst, struct mp_aframe *src)
215 {
216     dst->pts = src->pts;
217     dst->speed = src->speed;
218 
219     int rate = dst->av_frame->sample_rate;
220 
221     if (av_frame_copy_props(dst->av_frame, src->av_frame) < 0)
222         abort();
223 
224     dst->av_frame->sample_rate = rate;
225 }
226 
227 // Return whether a and b use the same physical audio format. Extra metadata
228 // such as PTS, per-frame signalling, and AVFrame side data is not compared.
mp_aframe_config_equals(struct mp_aframe * a,struct mp_aframe * b)229 bool mp_aframe_config_equals(struct mp_aframe *a, struct mp_aframe *b)
230 {
231     struct mp_chmap ca = {0}, cb = {0};
232     mp_aframe_get_chmap(a, &ca);
233     mp_aframe_get_chmap(b, &cb);
234     return mp_chmap_equals(&ca, &cb) &&
235            mp_aframe_get_rate(a) == mp_aframe_get_rate(b) &&
236            mp_aframe_get_format(a) == mp_aframe_get_format(b);
237 }
238 
239 // Return whether all required format fields have been set.
mp_aframe_config_is_valid(struct mp_aframe * frame)240 bool mp_aframe_config_is_valid(struct mp_aframe *frame)
241 {
242     return frame->format && frame->chmap.num && frame->av_frame->sample_rate;
243 }
244 
245 // Return the pointer to the first sample for each plane. The pointers stay
246 // valid until the next call that mutates frame somehow. You must not write to
247 // the audio data. Returns NULL if no frame allocated.
mp_aframe_get_data_ro(struct mp_aframe * frame)248 uint8_t **mp_aframe_get_data_ro(struct mp_aframe *frame)
249 {
250     return mp_aframe_is_allocated(frame) ? frame->av_frame->extended_data : NULL;
251 }
252 
253 // Like mp_aframe_get_data_ro(), but you can write to the audio data.
254 // Additionally, it will return NULL if copy-on-write fails.
mp_aframe_get_data_rw(struct mp_aframe * frame)255 uint8_t **mp_aframe_get_data_rw(struct mp_aframe *frame)
256 {
257     if (!mp_aframe_is_allocated(frame))
258         return NULL;
259     if (av_frame_make_writable(frame->av_frame) < 0)
260         return NULL;
261     return frame->av_frame->extended_data;
262 }
263 
mp_aframe_get_format(struct mp_aframe * frame)264 int mp_aframe_get_format(struct mp_aframe *frame)
265 {
266     return frame->format;
267 }
268 
mp_aframe_get_chmap(struct mp_aframe * frame,struct mp_chmap * out)269 bool mp_aframe_get_chmap(struct mp_aframe *frame, struct mp_chmap *out)
270 {
271     if (!mp_chmap_is_valid(&frame->chmap))
272         return false;
273     *out = frame->chmap;
274     return true;
275 }
276 
mp_aframe_get_channels(struct mp_aframe * frame)277 int mp_aframe_get_channels(struct mp_aframe *frame)
278 {
279     return frame->chmap.num;
280 }
281 
mp_aframe_get_rate(struct mp_aframe * frame)282 int mp_aframe_get_rate(struct mp_aframe *frame)
283 {
284     return frame->av_frame->sample_rate;
285 }
286 
mp_aframe_get_size(struct mp_aframe * frame)287 int mp_aframe_get_size(struct mp_aframe *frame)
288 {
289     return frame->av_frame->nb_samples;
290 }
291 
mp_aframe_get_pts(struct mp_aframe * frame)292 double mp_aframe_get_pts(struct mp_aframe *frame)
293 {
294     return frame->pts;
295 }
296 
mp_aframe_set_format(struct mp_aframe * frame,int format)297 bool mp_aframe_set_format(struct mp_aframe *frame, int format)
298 {
299     if (mp_aframe_is_allocated(frame))
300         return false;
301     enum AVSampleFormat av_format = af_to_avformat(format);
302     if (av_format == AV_SAMPLE_FMT_NONE && format) {
303         if (!af_fmt_is_spdif(format))
304             return false;
305         av_format = AV_SAMPLE_FMT_S16;
306     }
307     frame->format = format;
308     frame->av_frame->format = av_format;
309     return true;
310 }
311 
mp_aframe_set_chmap(struct mp_aframe * frame,struct mp_chmap * in)312 bool mp_aframe_set_chmap(struct mp_aframe *frame, struct mp_chmap *in)
313 {
314     if (!mp_chmap_is_valid(in) && !mp_chmap_is_empty(in))
315         return false;
316     if (mp_aframe_is_allocated(frame) && in->num != frame->chmap.num)
317         return false;
318     uint64_t lavc_layout = mp_chmap_to_lavc_unchecked(in);
319     if (!lavc_layout && in->num)
320         return false;
321     frame->chmap = *in;
322     frame->av_frame->channel_layout = lavc_layout;
323     // FFmpeg being a stupid POS again
324     frame->av_frame->channels = frame->chmap.num;
325     return true;
326 }
327 
mp_aframe_set_rate(struct mp_aframe * frame,int rate)328 bool mp_aframe_set_rate(struct mp_aframe *frame, int rate)
329 {
330     if (rate < 1 || rate > 10000000)
331         return false;
332     frame->av_frame->sample_rate = rate;
333     return true;
334 }
335 
mp_aframe_set_size(struct mp_aframe * frame,int samples)336 bool mp_aframe_set_size(struct mp_aframe *frame, int samples)
337 {
338     if (!mp_aframe_is_allocated(frame) || mp_aframe_get_size(frame) < samples)
339         return false;
340     frame->av_frame->nb_samples = MPMAX(samples, 0);
341     return true;
342 }
343 
mp_aframe_set_pts(struct mp_aframe * frame,double pts)344 void mp_aframe_set_pts(struct mp_aframe *frame, double pts)
345 {
346     frame->pts = pts;
347 }
348 
349 // Set a speed factor. This is multiplied with the sample rate to get the
350 // "effective" samplerate (mp_aframe_get_effective_rate()), which will be used
351 // to do PTS calculations. If speed!=1.0, the PTS values always refer to the
352 // original PTS (before changing speed), and if you want reasonably continuous
353 // PTS between frames, you need to use the effective samplerate.
mp_aframe_set_speed(struct mp_aframe * frame,double factor)354 void mp_aframe_set_speed(struct mp_aframe *frame, double factor)
355 {
356     frame->speed = factor;
357 }
358 
359 // Adjust current speed factor.
mp_aframe_mul_speed(struct mp_aframe * frame,double factor)360 void mp_aframe_mul_speed(struct mp_aframe *frame, double factor)
361 {
362     frame->speed *= factor;
363 }
364 
mp_aframe_get_speed(struct mp_aframe * frame)365 double mp_aframe_get_speed(struct mp_aframe *frame)
366 {
367     return frame->speed;
368 }
369 
370 // Matters for speed changed frames (such as a frame which has been resampled
371 // to play at a different speed).
372 // Return the sample rate at which the frame would have to be played to result
373 // in the same duration as the original frame before the speed change.
374 // This is used for A/V sync.
mp_aframe_get_effective_rate(struct mp_aframe * frame)375 double mp_aframe_get_effective_rate(struct mp_aframe *frame)
376 {
377     return mp_aframe_get_rate(frame) / frame->speed;
378 }
379 
380 // Return number of data pointers.
mp_aframe_get_planes(struct mp_aframe * frame)381 int mp_aframe_get_planes(struct mp_aframe *frame)
382 {
383     return af_fmt_is_planar(mp_aframe_get_format(frame))
384            ? mp_aframe_get_channels(frame) : 1;
385 }
386 
387 // Return number of bytes between 2 consecutive samples on the same plane.
mp_aframe_get_sstride(struct mp_aframe * frame)388 size_t mp_aframe_get_sstride(struct mp_aframe *frame)
389 {
390     int format = mp_aframe_get_format(frame);
391     return af_fmt_to_bytes(format) *
392            (af_fmt_is_planar(format) ? 1 : mp_aframe_get_channels(frame));
393 }
394 
395 // Return total number of samples on each plane.
mp_aframe_get_total_plane_samples(struct mp_aframe * frame)396 int mp_aframe_get_total_plane_samples(struct mp_aframe *frame)
397 {
398     return frame->av_frame->nb_samples *
399            (af_fmt_is_planar(mp_aframe_get_format(frame))
400             ? 1 : mp_aframe_get_channels(frame));
401 }
402 
mp_aframe_format_str_buf(char * buf,size_t buf_size,struct mp_aframe * fmt)403 char *mp_aframe_format_str_buf(char *buf, size_t buf_size, struct mp_aframe *fmt)
404 {
405     char ch[128];
406     mp_chmap_to_str_buf(ch, sizeof(ch), &fmt->chmap);
407     char *hr_ch = mp_chmap_to_str_hr(&fmt->chmap);
408     if (strcmp(hr_ch, ch) != 0)
409         mp_snprintf_cat(ch, sizeof(ch), " (%s)", hr_ch);
410     snprintf(buf, buf_size, "%dHz %s %dch %s", fmt->av_frame->sample_rate,
411              ch, fmt->chmap.num, af_fmt_to_str(fmt->format));
412     return buf;
413 }
414 
415 // Set data to the audio after the given number of samples (i.e. slice it).
mp_aframe_skip_samples(struct mp_aframe * f,int samples)416 void mp_aframe_skip_samples(struct mp_aframe *f, int samples)
417 {
418     assert(samples >= 0 && samples <= mp_aframe_get_size(f));
419 
420     if (av_frame_make_writable(f->av_frame) < 0)
421         return; // go complain to ffmpeg
422 
423     int num_planes = mp_aframe_get_planes(f);
424     size_t sstride = mp_aframe_get_sstride(f);
425     for (int n = 0; n < num_planes; n++) {
426         memmove(f->av_frame->extended_data[n],
427                 f->av_frame->extended_data[n] + samples * sstride,
428                 (f->av_frame->nb_samples - samples) * sstride);
429     }
430 
431     f->av_frame->nb_samples -= samples;
432 
433     if (f->pts != MP_NOPTS_VALUE)
434         f->pts += samples / mp_aframe_get_effective_rate(f);
435 }
436 
437 // Return the timestamp of the sample just after the end of this frame.
mp_aframe_end_pts(struct mp_aframe * f)438 double mp_aframe_end_pts(struct mp_aframe *f)
439 {
440     double rate = mp_aframe_get_effective_rate(f);
441     if (f->pts == MP_NOPTS_VALUE || rate <= 0)
442         return MP_NOPTS_VALUE;
443     return f->pts + f->av_frame->nb_samples / rate;
444 }
445 
446 // Return the duration in seconds of the frame (0 if invalid).
mp_aframe_duration(struct mp_aframe * f)447 double mp_aframe_duration(struct mp_aframe *f)
448 {
449     double rate = mp_aframe_get_effective_rate(f);
450     if (rate <= 0)
451         return 0;
452     return f->av_frame->nb_samples / rate;
453 }
454 
455 // Clip the given frame to the given timestamp range. Adjusts the frame size
456 // and timestamp.
457 // Refuses to change spdif frames.
mp_aframe_clip_timestamps(struct mp_aframe * f,double start,double end)458 void mp_aframe_clip_timestamps(struct mp_aframe *f, double start, double end)
459 {
460     double f_end = mp_aframe_end_pts(f);
461     double rate = mp_aframe_get_effective_rate(f);
462     if (f_end == MP_NOPTS_VALUE)
463         return;
464     if (end != MP_NOPTS_VALUE) {
465         if (f_end >= end) {
466             if (f->pts >= end) {
467                 f->av_frame->nb_samples = 0;
468             } else {
469                 if (af_fmt_is_spdif(mp_aframe_get_format(f)))
470                     return;
471                 int new = (end - f->pts) * rate;
472                 f->av_frame->nb_samples = MPCLAMP(new, 0, f->av_frame->nb_samples);
473             }
474         }
475     }
476     if (start != MP_NOPTS_VALUE) {
477         if (f->pts < start) {
478             if (f_end <= start) {
479                 f->av_frame->nb_samples = 0;
480                 f->pts = f_end;
481             } else {
482                 if (af_fmt_is_spdif(mp_aframe_get_format(f)))
483                     return;
484                 int skip = (start - f->pts) * rate;
485                 skip = MPCLAMP(skip, 0, f->av_frame->nb_samples);
486                 mp_aframe_skip_samples(f, skip);
487             }
488         }
489     }
490 }
491 
mp_aframe_copy_samples(struct mp_aframe * dst,int dst_offset,struct mp_aframe * src,int src_offset,int samples)492 bool mp_aframe_copy_samples(struct mp_aframe *dst, int dst_offset,
493                             struct mp_aframe *src, int src_offset,
494                             int samples)
495 {
496     if (!mp_aframe_config_equals(dst, src))
497         return false;
498 
499     if (mp_aframe_get_size(dst) < dst_offset + samples ||
500         mp_aframe_get_size(src) < src_offset + samples)
501         return false;
502 
503     uint8_t **s = mp_aframe_get_data_ro(src);
504     uint8_t **d = mp_aframe_get_data_rw(dst);
505     if (!s || !d)
506         return false;
507 
508     int planes = mp_aframe_get_planes(dst);
509     size_t sstride = mp_aframe_get_sstride(dst);
510 
511     for (int n = 0; n < planes; n++) {
512         memcpy(d[n] + dst_offset * sstride, s[n] + src_offset * sstride,
513                samples * sstride);
514     }
515 
516     return true;
517 }
518 
mp_aframe_set_silence(struct mp_aframe * f,int offset,int samples)519 bool mp_aframe_set_silence(struct mp_aframe *f, int offset, int samples)
520 {
521     if (mp_aframe_get_size(f) < offset + samples)
522         return false;
523 
524     int format = mp_aframe_get_format(f);
525     uint8_t **d = mp_aframe_get_data_rw(f);
526     if (!d)
527         return false;
528 
529     int planes = mp_aframe_get_planes(f);
530     size_t sstride = mp_aframe_get_sstride(f);
531 
532     for (int n = 0; n < planes; n++)
533         af_fill_silence(d[n] + offset * sstride, samples * sstride, format);
534 
535     return true;
536 }
537 
mp_aframe_reverse(struct mp_aframe * f)538 bool mp_aframe_reverse(struct mp_aframe *f)
539 {
540     int format = mp_aframe_get_format(f);
541     size_t bps = af_fmt_to_bytes(format);
542     if (!af_fmt_is_pcm(format) || bps > 16)
543         return false;
544 
545     uint8_t **d = mp_aframe_get_data_rw(f);
546     if (!d)
547         return false;
548 
549     int planes = mp_aframe_get_planes(f);
550     int samples = mp_aframe_get_size(f);
551     int channels = mp_aframe_get_channels(f);
552     size_t sstride = mp_aframe_get_sstride(f);
553 
554     int plane_samples = channels;
555     if (af_fmt_is_planar(format))
556         plane_samples = 1;
557 
558     for (int p = 0; p < planes; p++) {
559         for (int n = 0; n < samples / 2; n++) {
560             int s1_offset = n * sstride;
561             int s2_offset = (samples - 1 - n) * sstride;
562             for (int c = 0; c < plane_samples; c++) {
563                 // Nobody said it'd be fast.
564                 char tmp[16];
565                 uint8_t *s1 = d[p] + s1_offset + c * bps;
566                 uint8_t *s2 = d[p] + s2_offset + c * bps;
567                 memcpy(tmp, s2, bps);
568                 memcpy(s2, s1, bps);
569                 memcpy(s1, tmp, bps);
570             }
571         }
572     }
573 
574     return true;
575 }
576 
mp_aframe_approx_byte_size(struct mp_aframe * frame)577 int mp_aframe_approx_byte_size(struct mp_aframe *frame)
578 {
579     // God damn, AVFrame is too fucking annoying. Just go with the size that
580     // allocating a new frame would use.
581     int planes = mp_aframe_get_planes(frame);
582     size_t sstride = mp_aframe_get_sstride(frame);
583     int samples = frame->av_frame->nb_samples;
584     int plane_size = MP_ALIGN_UP(sstride * MPMAX(samples, 1), 32);
585     return plane_size * planes + sizeof(*frame);
586 }
587 
588 struct mp_aframe_pool {
589     AVBufferPool *avpool;
590     int element_size;
591 };
592 
mp_aframe_pool_create(void * ta_parent)593 struct mp_aframe_pool *mp_aframe_pool_create(void *ta_parent)
594 {
595     return talloc_zero(ta_parent, struct mp_aframe_pool);
596 }
597 
mp_aframe_pool_destructor(void * p)598 static void mp_aframe_pool_destructor(void *p)
599 {
600     struct mp_aframe_pool *pool = p;
601     av_buffer_pool_uninit(&pool->avpool);
602 }
603 
604 // Like mp_aframe_allocate(), but use the pool to allocate data.
mp_aframe_pool_allocate(struct mp_aframe_pool * pool,struct mp_aframe * frame,int samples)605 int mp_aframe_pool_allocate(struct mp_aframe_pool *pool, struct mp_aframe *frame,
606                             int samples)
607 {
608     int planes = mp_aframe_get_planes(frame);
609     size_t sstride = mp_aframe_get_sstride(frame);
610     // FFmpeg hardcodes similar hidden possibly-requirements in a number of
611     // places: av_frame_get_buffer(), libavcodec's get_buffer(), mem.c,
612     // probably more.
613     int align_samples = MP_ALIGN_UP(MPMAX(samples, 1), 32);
614     int plane_size = MP_ALIGN_UP(sstride * align_samples, 64);
615     int size = plane_size * planes;
616 
617     if (size <= 0 || mp_aframe_is_allocated(frame))
618         return -1;
619 
620     if (!pool->avpool || size > pool->element_size) {
621         size_t alloc = ta_calc_prealloc_elems(size);
622         if (alloc >= INT_MAX)
623             return -1;
624         av_buffer_pool_uninit(&pool->avpool);
625         pool->element_size = alloc;
626         pool->avpool = av_buffer_pool_init(pool->element_size, NULL);
627         if (!pool->avpool)
628             return -1;
629         talloc_set_destructor(pool, mp_aframe_pool_destructor);
630     }
631 
632     // Yes, you have to do all this shit manually.
633     // At least it's less stupid than av_frame_get_buffer(), which just wipes
634     // the entire frame struct on error for no reason.
635     AVFrame *av_frame = frame->av_frame;
636     if (av_frame->extended_data != av_frame->data)
637         av_freep(&av_frame->extended_data); // sigh
638     if (planes > AV_NUM_DATA_POINTERS) {
639         av_frame->extended_data =
640             av_mallocz_array(planes, sizeof(av_frame->extended_data[0]));
641         if (!av_frame->extended_data)
642             abort();
643     } else {
644         av_frame->extended_data = av_frame->data;
645     }
646     av_frame->buf[0] = av_buffer_pool_get(pool->avpool);
647     if (!av_frame->buf[0])
648         return -1;
649     av_frame->linesize[0] = samples * sstride;
650     for (int n = 0; n < planes; n++)
651         av_frame->extended_data[n] = av_frame->buf[0]->data + n * plane_size;
652     if (planes > AV_NUM_DATA_POINTERS) {
653         for (int n = 0; n < AV_NUM_DATA_POINTERS; n++)
654             av_frame->data[n] = av_frame->extended_data[n];
655     }
656     av_frame->nb_samples = samples;
657 
658     return 0;
659 }
660