git.sesse.net Git - ffmpeg/blob - libavcodec/cbs_av1.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 #include "libavutil/pixfmt.h"
  21
  22 #include "cbs.h"
  23 #include "cbs_internal.h"
  24 #include "cbs_av1.h"
  25 #include "internal.h"
  26
  27
  28 static int cbs_av1_read_uvlc(CodedBitstreamContext *ctx, GetBitContext *gbc,
  29                              const char *name, uint32_t *write_to,
  30                              uint32_t range_min, uint32_t range_max)
  31 {
  32     uint32_t zeroes, bits_value, value;
  33     int position;
  34
  35     if (ctx->trace_enable)
  36         position = get_bits_count(gbc);
  37
  38     zeroes = 0;
  39     while (1) {
  40         if (get_bits_left(gbc) < 1) {
  41             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid uvlc code at "
  42                    "%s: bitstream ended.\n", name);
  43             return AVERROR_INVALIDDATA;
  44         }
  45
  46         if (get_bits1(gbc))
  47             break;
  48         ++zeroes;
  49     }
  50
  51     if (zeroes >= 32) {
  52         value = MAX_UINT_BITS(32);
  53     } else {
  54         if (get_bits_left(gbc) < zeroes) {
  55             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid uvlc code at "
  56                    "%s: bitstream ended.\n", name);
  57             return AVERROR_INVALIDDATA;
  58         }
  59
  60         bits_value = get_bits_long(gbc, zeroes);
  61         value = bits_value + (UINT32_C(1) << zeroes) - 1;
  62     }
  63
  64     if (ctx->trace_enable) {
  65         char bits[65];
  66         int i, j, k;
  67
  68         if (zeroes >= 32) {
  69             while (zeroes > 32) {
  70                 k = FFMIN(zeroes - 32, 32);
  71                 for (i = 0; i < k; i++)
  72                     bits[i] = '0';
  73                 bits[i] = 0;
  74                 ff_cbs_trace_syntax_element(ctx, position, name,
  75                                             NULL, bits, 0);
  76                 zeroes -= k;
  77                 position += k;
  78             }
  79         }
  80
  81         for (i = 0; i < zeroes; i++)
  82             bits[i] = '0';
  83         bits[i++] = '1';
  84
  85         if (zeroes < 32) {
  86             for (j = 0; j < zeroes; j++)
  87                 bits[i++] = (bits_value >> (zeroes - j - 1) & 1) ? '1' : '0';
  88         }
  89
  90         bits[i] = 0;
  91         ff_cbs_trace_syntax_element(ctx, position, name,
  92                                     NULL, bits, value);
  93     }
  94
  95     if (value < range_min || value > range_max) {
  96         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
  97                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
  98                name, value, range_min, range_max);
  99         return AVERROR_INVALIDDATA;
 100     }
 101
 102     *write_to = value;
 103     return 0;
 104 }
 105
 106 static int cbs_av1_write_uvlc(CodedBitstreamContext *ctx, PutBitContext *pbc,
 107                               const char *name, uint32_t value,
 108                               uint32_t range_min, uint32_t range_max)
 109 {
 110     uint32_t v;
 111     int position, zeroes;
 112
 113     if (value < range_min || value > range_max) {
 114         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 115                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
 116                name, value, range_min, range_max);
 117         return AVERROR_INVALIDDATA;
 118     }
 119
 120     if (ctx->trace_enable)
 121         position = put_bits_count(pbc);
 122
 123     if (value == 0) {
 124         zeroes = 0;
 125         put_bits(pbc, 1, 1);
 126     } else {
 127         zeroes = av_log2(value + 1);
 128         v = value - (1 << zeroes) + 1;
 129         put_bits(pbc, zeroes + 1, 1);
 130         put_bits(pbc, zeroes, v);
 131     }
 132
 133     if (ctx->trace_enable) {
 134         char bits[65];
 135         int i, j;
 136         i = 0;
 137         for (j = 0; j < zeroes; j++)
 138             bits[i++] = '0';
 139         bits[i++] = '1';
 140         for (j = 0; j < zeroes; j++)
 141             bits[i++] = (v >> (zeroes - j - 1) & 1) ? '1' : '0';
 142         bits[i++] = 0;
 143         ff_cbs_trace_syntax_element(ctx, position, name, NULL,
 144                                     bits, value);
 145     }
 146
 147     return 0;
 148 }
 149
 150 static int cbs_av1_read_leb128(CodedBitstreamContext *ctx, GetBitContext *gbc,
 151                                const char *name, uint64_t *write_to)
 152 {
 153     uint64_t value;
 154     int position, err, i;
 155
 156     if (ctx->trace_enable)
 157         position = get_bits_count(gbc);
 158
 159     value = 0;
 160     for (i = 0; i < 8; i++) {
 161         int subscript[2] = { 1, i };
 162         uint32_t byte;
 163         err = ff_cbs_read_unsigned(ctx, gbc, 8, "leb128_byte[i]", subscript,
 164                                    &byte, 0x00, 0xff);
 165         if (err < 0)
 166             return err;
 167
 168         value |= (uint64_t)(byte & 0x7f) << (i * 7);
 169         if (!(byte & 0x80))
 170             break;
 171     }
 172
 173     if (value > UINT32_MAX)
 174         return AVERROR_INVALIDDATA;
 175
 176     if (ctx->trace_enable)
 177         ff_cbs_trace_syntax_element(ctx, position, name, NULL, "", value);
 178
 179     *write_to = value;
 180     return 0;
 181 }
 182
 183 static int cbs_av1_write_leb128(CodedBitstreamContext *ctx, PutBitContext *pbc,
 184                                 const char *name, uint64_t value)
 185 {
 186     int position, err, len, i;
 187     uint8_t byte;
 188
 189     len = (av_log2(value) + 7) / 7;
 190
 191     if (ctx->trace_enable)
 192         position = put_bits_count(pbc);
 193
 194     for (i = 0; i < len; i++) {
 195         int subscript[2] = { 1, i };
 196
 197         byte = value >> (7 * i) & 0x7f;
 198         if (i < len - 1)
 199             byte |= 0x80;
 200
 201         err = ff_cbs_write_unsigned(ctx, pbc, 8, "leb128_byte[i]", subscript,
 202                                     byte, 0x00, 0xff);
 203         if (err < 0)
 204             return err;
 205     }
 206
 207     if (ctx->trace_enable)
 208         ff_cbs_trace_syntax_element(ctx, position, name, NULL, "", value);
 209
 210     return 0;
 211 }
 212
 213 static int cbs_av1_read_ns(CodedBitstreamContext *ctx, GetBitContext *gbc,
 214                            uint32_t n, const char *name,
 215                            const int *subscripts, uint32_t *write_to)
 216 {
 217     uint32_t m, v, extra_bit, value;
 218     int position, w;
 219
 220     av_assert0(n > 0);
 221
 222     if (ctx->trace_enable)
 223         position = get_bits_count(gbc);
 224
 225     w = av_log2(n) + 1;
 226     m = (1 << w) - n;
 227
 228     if (get_bits_left(gbc) < w) {
 229         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid non-symmetric value at "
 230                "%s: bitstream ended.\n", name);
 231         return AVERROR_INVALIDDATA;
 232     }
 233
 234     if (w - 1 > 0)
 235         v = get_bits(gbc, w - 1);
 236     else
 237         v = 0;
 238
 239     if (v < m) {
 240         value = v;
 241     } else {
 242         extra_bit = get_bits1(gbc);
 243         value = (v << 1) - m + extra_bit;
 244     }
 245
 246     if (ctx->trace_enable) {
 247         char bits[33];
 248         int i;
 249         for (i = 0; i < w - 1; i++)
 250             bits[i] = (v >> i & 1) ? '1' : '0';
 251         if (v >= m)
 252             bits[i++] = extra_bit ? '1' : '0';
 253         bits[i] = 0;
 254
 255         ff_cbs_trace_syntax_element(ctx, position,
 256                                     name, subscripts, bits, value);
 257     }
 258
 259     *write_to = value;
 260     return 0;
 261 }
 262
 263 static int cbs_av1_write_ns(CodedBitstreamContext *ctx, PutBitContext *pbc,
 264                             uint32_t n, const char *name,
 265                             const int *subscripts, uint32_t value)
 266 {
 267     uint32_t w, m, v, extra_bit;
 268     int position;
 269
 270     if (value > n) {
 271         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 272                "%"PRIu32", but must be in [0,%"PRIu32"].\n",
 273                name, value, n);
 274         return AVERROR_INVALIDDATA;
 275     }
 276
 277     if (ctx->trace_enable)
 278         position = put_bits_count(pbc);
 279
 280     w = av_log2(n) + 1;
 281     m = (1 << w) - n;
 282
 283     if (put_bits_left(pbc) < w)
 284         return AVERROR(ENOSPC);
 285
 286     if (value < m) {
 287         v = value;
 288         put_bits(pbc, w - 1, v);
 289     } else {
 290         v = m + ((value - m) >> 1);
 291         extra_bit = (value - m) & 1;
 292         put_bits(pbc, w - 1, v);
 293         put_bits(pbc, 1, extra_bit);
 294     }
 295
 296     if (ctx->trace_enable) {
 297         char bits[33];
 298         int i;
 299         for (i = 0; i < w - 1; i++)
 300             bits[i] = (v >> i & 1) ? '1' : '0';
 301         if (value >= m)
 302             bits[i++] = extra_bit ? '1' : '0';
 303         bits[i] = 0;
 304
 305         ff_cbs_trace_syntax_element(ctx, position,
 306                                     name, subscripts, bits, value);
 307     }
 308
 309     return 0;
 310 }
 311
 312 static int cbs_av1_read_increment(CodedBitstreamContext *ctx, GetBitContext *gbc,
 313                                   uint32_t range_min, uint32_t range_max,
 314                                   const char *name, uint32_t *write_to)
 315 {
 316     uint32_t value;
 317     int position, i;
 318     char bits[33];
 319
 320     av_assert0(range_min <= range_max && range_max - range_min < sizeof(bits) - 1);
 321     if (ctx->trace_enable)
 322         position = get_bits_count(gbc);
 323
 324     for (i = 0, value = range_min; value < range_max;) {
 325         if (get_bits_left(gbc) < 1) {
 326             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid increment value at "
 327                    "%s: bitstream ended.\n", name);
 328             return AVERROR_INVALIDDATA;
 329         }
 330         if (get_bits1(gbc)) {
 331             bits[i++] = '1';
 332             ++value;
 333         } else {
 334             bits[i++] = '0';
 335             break;
 336         }
 337     }
 338
 339     if (ctx->trace_enable) {
 340         bits[i] = 0;
 341         ff_cbs_trace_syntax_element(ctx, position,
 342                                     name, NULL, bits, value);
 343     }
 344
 345     *write_to = value;
 346     return 0;
 347 }
 348
 349 static int cbs_av1_write_increment(CodedBitstreamContext *ctx, PutBitContext *pbc,
 350                                    uint32_t range_min, uint32_t range_max,
 351                                    const char *name, uint32_t value)
 352 {
 353     int len;
 354
 355     av_assert0(range_min <= range_max && range_max - range_min < 32);
 356     if (value < range_min || value > range_max) {
 357         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 358                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
 359                name, value, range_min, range_max);
 360         return AVERROR_INVALIDDATA;
 361     }
 362
 363     if (value == range_max)
 364         len = range_max - range_min;
 365     else
 366         len = value - range_min + 1;
 367     if (put_bits_left(pbc) < len)
 368         return AVERROR(ENOSPC);
 369
 370     if (ctx->trace_enable) {
 371         char bits[33];
 372         int i;
 373         for (i = 0; i < len; i++) {
 374             if (range_min + i == value)
 375                 bits[i] = '0';
 376             else
 377                 bits[i] = '1';
 378         }
 379         bits[i] = 0;
 380         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc),
 381                                     name, NULL, bits, value);
 382     }
 383
 384     if (len > 0)
 385         put_bits(pbc, len, (1 << len) - 1 - (value != range_max));
 386
 387     return 0;
 388 }
 389
 390 static int cbs_av1_read_subexp(CodedBitstreamContext *ctx, GetBitContext *gbc,
 391                                uint32_t range_max, const char *name,
 392                                const int *subscripts, uint32_t *write_to)
 393 {
 394     uint32_t value;
 395     int position, err;
 396     uint32_t max_len, len, range_offset, range_bits;
 397
 398     if (ctx->trace_enable)
 399         position = get_bits_count(gbc);
 400
 401     av_assert0(range_max > 0);
 402     max_len = av_log2(range_max - 1) - 3;
 403
 404     err = cbs_av1_read_increment(ctx, gbc, 0, max_len,
 405                                  "subexp_more_bits", &len);
 406     if (err < 0)
 407         return err;
 408
 409     if (len) {
 410         range_bits   = 2 + len;
 411         range_offset = 1 << range_bits;
 412     } else {
 413         range_bits   = 3;
 414         range_offset = 0;
 415     }
 416
 417     if (len < max_len) {
 418         err = ff_cbs_read_unsigned(ctx, gbc, range_bits,
 419                                    "subexp_bits", NULL, &value,
 420                                    0, MAX_UINT_BITS(range_bits));
 421         if (err < 0)
 422             return err;
 423
 424     } else {
 425         err = cbs_av1_read_ns(ctx, gbc, range_max - range_offset,
 426                               "subexp_final_bits", NULL, &value);
 427         if (err < 0)
 428             return err;
 429     }
 430     value += range_offset;
 431
 432     if (ctx->trace_enable)
 433         ff_cbs_trace_syntax_element(ctx, position,
 434                                     name, subscripts, "", value);
 435
 436     *write_to = value;
 437     return err;
 438 }
 439
 440 static int cbs_av1_write_subexp(CodedBitstreamContext *ctx, PutBitContext *pbc,
 441                                 uint32_t range_max, const char *name,
 442                                 const int *subscripts, uint32_t value)
 443 {
 444     int position, err;
 445     uint32_t max_len, len, range_offset, range_bits;
 446
 447     if (value > range_max) {
 448         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 449                "%"PRIu32", but must be in [0,%"PRIu32"].\n",
 450                name, value, range_max);
 451         return AVERROR_INVALIDDATA;
 452     }
 453
 454     if (ctx->trace_enable)
 455         position = put_bits_count(pbc);
 456
 457     av_assert0(range_max > 0);
 458     max_len = av_log2(range_max - 1) - 3;
 459
 460     if (value < 8) {
 461         range_bits   = 3;
 462         range_offset = 0;
 463         len = 0;
 464     } else {
 465         range_bits = av_log2(value);
 466         len = range_bits - 2;
 467         if (len > max_len) {
 468             // The top bin is combined with the one below it.
 469             av_assert0(len == max_len + 1);
 470             --range_bits;
 471             len = max_len;
 472         }
 473         range_offset = 1 << range_bits;
 474     }
 475
 476     err = cbs_av1_write_increment(ctx, pbc, 0, max_len,
 477                                   "subexp_more_bits", len);
 478     if (err < 0)
 479         return err;
 480
 481     if (len < max_len) {
 482         err = ff_cbs_write_unsigned(ctx, pbc, range_bits,
 483                                     "subexp_bits", NULL,
 484                                     value - range_offset,
 485                                     0, MAX_UINT_BITS(range_bits));
 486         if (err < 0)
 487             return err;
 488
 489     } else {
 490         err = cbs_av1_write_ns(ctx, pbc, range_max - range_offset,
 491                                "subexp_final_bits", NULL,
 492                                value - range_offset);
 493         if (err < 0)
 494             return err;
 495     }
 496
 497     if (ctx->trace_enable)
 498         ff_cbs_trace_syntax_element(ctx, position,
 499                                     name, subscripts, "", value);
 500
 501     return err;
 502 }
 503
 504
 505 static int cbs_av1_tile_log2(int blksize, int target)
 506 {
 507     int k;
 508     for (k = 0; (blksize << k) < target; k++);
 509     return k;
 510 }
 511
 512 static int cbs_av1_get_relative_dist(const AV1RawSequenceHeader *seq,
 513                                      unsigned int a, unsigned int b)
 514 {
 515     unsigned int diff, m;
 516     if (!seq->enable_order_hint)
 517         return 0;
 518     diff = a - b;
 519     m = 1 << seq->order_hint_bits_minus_1;
 520     diff = (diff & (m - 1)) - (diff & m);
 521     return diff;
 522 }
 523
 524 static size_t cbs_av1_get_payload_bytes_left(GetBitContext *gbc)
 525 {
 526     GetBitContext tmp = *gbc;
 527     size_t size = 0;
 528     for (int i = 0; get_bits_left(&tmp) >= 8; i++) {
 529         if (get_bits(&tmp, 8))
 530             size = i;
 531     }
 532     return size;
 533 }
 534
 535
 536 #define HEADER(name) do { \
 537         ff_cbs_trace_header(ctx, name); \
 538     } while (0)
 539
 540 #define CHECK(call) do { \
 541         err = (call); \
 542         if (err < 0) \
 543             return err; \
 544     } while (0)
 545
 546 #define FUNC_NAME(rw, codec, name) cbs_ ## codec ## _ ## rw ## _ ## name
 547 #define FUNC_AV1(rw, name) FUNC_NAME(rw, av1, name)
 548 #define FUNC(name) FUNC_AV1(READWRITE, name)
 549
 550 #define SUBSCRIPTS(subs, ...) (subs > 0 ? ((int[subs + 1]){ subs, __VA_ARGS__ }) : NULL)
 551
 552 #define fb(width, name) \
 553         xf(width, name, current->name, 0, MAX_UINT_BITS(width), 0, )
 554 #define fc(width, name, range_min, range_max) \
 555         xf(width, name, current->name, range_min, range_max, 0, )
 556 #define flag(name) fb(1, name)
 557 #define su(width, name) \
 558         xsu(width, name, current->name, 0, )
 559
 560 #define fbs(width, name, subs, ...) \
 561         xf(width, name, current->name, 0, MAX_UINT_BITS(width), subs, __VA_ARGS__)
 562 #define fcs(width, name, range_min, range_max, subs, ...) \
 563         xf(width, name, current->name, range_min, range_max, subs, __VA_ARGS__)
 564 #define flags(name, subs, ...) \
 565         xf(1, name, current->name, 0, 1, subs, __VA_ARGS__)
 566 #define sus(width, name, subs, ...) \
 567         xsu(width, name, current->name, subs, __VA_ARGS__)
 568
 569 #define fixed(width, name, value) do { \
 570         av_unused uint32_t fixed_value = value; \
 571         xf(width, name, fixed_value, value, value, 0, ); \
 572     } while (0)
 573
 574
 575 #define READ
 576 #define READWRITE read
 577 #define RWContext GetBitContext
 578
 579 #define xf(width, name, var, range_min, range_max, subs, ...) do { \
 580         uint32_t value; \
 581         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
 582                                    SUBSCRIPTS(subs, __VA_ARGS__), \
 583                                    &value, range_min, range_max)); \
 584         var = value; \
 585     } while (0)
 586
 587 #define xsu(width, name, var, subs, ...) do { \
 588         int32_t value; \
 589         CHECK(ff_cbs_read_signed(ctx, rw, width, #name, \
 590                                  SUBSCRIPTS(subs, __VA_ARGS__), &value, \
 591                                  MIN_INT_BITS(width), \
 592                                  MAX_INT_BITS(width))); \
 593         var = value; \
 594     } while (0)
 595
 596 #define uvlc(name, range_min, range_max) do { \
 597         uint32_t value; \
 598         CHECK(cbs_av1_read_uvlc(ctx, rw, #name, \
 599                                 &value, range_min, range_max)); \
 600         current->name = value; \
 601     } while (0)
 602
 603 #define ns(max_value, name, subs, ...) do { \
 604         uint32_t value; \
 605         CHECK(cbs_av1_read_ns(ctx, rw, max_value, #name, \
 606                               SUBSCRIPTS(subs, __VA_ARGS__), &value)); \
 607         current->name = value; \
 608     } while (0)
 609
 610 #define increment(name, min, max) do { \
 611         uint32_t value; \
 612         CHECK(cbs_av1_read_increment(ctx, rw, min, max, #name, &value)); \
 613         current->name = value; \
 614     } while (0)
 615
 616 #define subexp(name, max, subs, ...) do { \
 617         uint32_t value; \
 618         CHECK(cbs_av1_read_subexp(ctx, rw, max, #name, \
 619                                   SUBSCRIPTS(subs, __VA_ARGS__), &value)); \
 620         current->name = value; \
 621     } while (0)
 622
 623 #define delta_q(name) do { \
 624         uint8_t delta_coded; \
 625         int8_t delta_q; \
 626         xf(1, name.delta_coded, delta_coded, 0, 1, 0, ); \
 627         if (delta_coded) \
 628             xsu(1 + 6, name.delta_q, delta_q, 0, ); \
 629         else \
 630             delta_q = 0; \
 631         current->name = delta_q; \
 632     } while (0)
 633
 634 #define leb128(name) do { \
 635         uint64_t value; \
 636         CHECK(cbs_av1_read_leb128(ctx, rw, #name, &value)); \
 637         current->name = value; \
 638     } while (0)
 639
 640 #define infer(name, value) do { \
 641         current->name = value; \
 642     } while (0)
 643
 644 #define byte_alignment(rw) (get_bits_count(rw) % 8)
 645
 646 #include "cbs_av1_syntax_template.c"
 647
 648 #undef READ
 649 #undef READWRITE
 650 #undef RWContext
 651 #undef xf
 652 #undef xsu
 653 #undef uvlc
 654 #undef ns
 655 #undef increment
 656 #undef subexp
 657 #undef delta_q
 658 #undef leb128
 659 #undef infer
 660 #undef byte_alignment
 661
 662
 663 #define WRITE
 664 #define READWRITE write
 665 #define RWContext PutBitContext
 666
 667 #define xf(width, name, var, range_min, range_max, subs, ...) do { \
 668         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
 669                                     SUBSCRIPTS(subs, __VA_ARGS__), \
 670                                     var, range_min, range_max)); \
 671     } while (0)
 672
 673 #define xsu(width, name, var, subs, ...) do { \
 674         CHECK(ff_cbs_write_signed(ctx, rw, width, #name, \
 675                                   SUBSCRIPTS(subs, __VA_ARGS__), var, \
 676                                   MIN_INT_BITS(width), \
 677                                   MAX_INT_BITS(width))); \
 678     } while (0)
 679
 680 #define uvlc(name, range_min, range_max) do { \
 681         CHECK(cbs_av1_write_uvlc(ctx, rw, #name, current->name, \
 682                                  range_min, range_max)); \
 683     } while (0)
 684
 685 #define ns(max_value, name, subs, ...) do { \
 686         CHECK(cbs_av1_write_ns(ctx, rw, max_value, #name, \
 687                                SUBSCRIPTS(subs, __VA_ARGS__), \
 688                                current->name)); \
 689     } while (0)
 690
 691 #define increment(name, min, max) do { \
 692         CHECK(cbs_av1_write_increment(ctx, rw, min, max, #name, \
 693                                       current->name)); \
 694     } while (0)
 695
 696 #define subexp(name, max, subs, ...) do { \
 697         CHECK(cbs_av1_write_subexp(ctx, rw, max, #name, \
 698                                    SUBSCRIPTS(subs, __VA_ARGS__), \
 699                                    current->name)); \
 700     } while (0)
 701
 702 #define delta_q(name) do { \
 703         xf(1, name.delta_coded, current->name != 0, 0, 1, 0, ); \
 704         if (current->name) \
 705             xsu(1 + 6, name.delta_q, current->name, 0, ); \
 706     } while (0)
 707
 708 #define leb128(name) do { \
 709         CHECK(cbs_av1_write_leb128(ctx, rw, #name, current->name)); \
 710     } while (0)
 711
 712 #define infer(name, value) do { \
 713         if (current->name != (value)) { \
 714             av_log(ctx->log_ctx, AV_LOG_WARNING, "Warning: " \
 715                    "%s does not match inferred value: " \
 716                    "%"PRId64", but should be %"PRId64".\n", \
 717                    #name, (int64_t)current->name, (int64_t)(value)); \
 718         } \
 719     } while (0)
 720
 721 #define byte_alignment(rw) (put_bits_count(rw) % 8)
 722
 723 #include "cbs_av1_syntax_template.c"
 724
 725 #undef WRITE
 726 #undef READWRITE
 727 #undef RWContext
 728 #undef xf
 729 #undef xsu
 730 #undef uvlc
 731 #undef ns
 732 #undef increment
 733 #undef subexp
 734 #undef delta_q
 735 #undef leb128
 736 #undef infer
 737 #undef byte_alignment
 738
 739
 740 static int cbs_av1_split_fragment(CodedBitstreamContext *ctx,
 741                                   CodedBitstreamFragment *frag,
 742                                   int header)
 743 {
 744     GetBitContext gbc;
 745     uint8_t *data;
 746     size_t size;
 747     uint64_t obu_length;
 748     int pos, err, trace;
 749
 750     // Don't include this parsing in trace output.
 751     trace = ctx->trace_enable;
 752     ctx->trace_enable = 0;
 753
 754     data = frag->data;
 755     size = frag->data_size;
 756
 757     if (INT_MAX / 8 < size) {
 758         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid fragment: "
 759                "too large (%"SIZE_SPECIFIER" bytes).\n", size);
 760         err = AVERROR_INVALIDDATA;
 761         goto fail;
 762     }
 763
 764     while (size > 0) {
 765         AV1RawOBUHeader header;
 766         uint64_t obu_size;
 767
 768         init_get_bits(&gbc, data, 8 * size);
 769
 770         err = cbs_av1_read_obu_header(ctx, &gbc, &header);
 771         if (err < 0)
 772             goto fail;
 773
 774         if (header.obu_has_size_field) {
 775             if (get_bits_left(&gbc) < 8) {
 776                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid OBU: fragment "
 777                        "too short (%"SIZE_SPECIFIER" bytes).\n", size);
 778                 err = AVERROR_INVALIDDATA;
 779                 goto fail;
 780             }
 781             err = cbs_av1_read_leb128(ctx, &gbc, "obu_size", &obu_size);
 782             if (err < 0)
 783                 goto fail;
 784         } else
 785             obu_size = size - 1 - header.obu_extension_flag;
 786
 787         pos = get_bits_count(&gbc);
 788         av_assert0(pos % 8 == 0 && pos / 8 <= size);
 789
 790         obu_length = pos / 8 + obu_size;
 791
 792         if (size < obu_length) {
 793             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid OBU length: "
 794                    "%"PRIu64", but only %"SIZE_SPECIFIER" bytes remaining in fragment.\n",
 795                    obu_length, size);
 796             err = AVERROR_INVALIDDATA;
 797             goto fail;
 798         }
 799
 800         err = ff_cbs_insert_unit_data(ctx, frag, -1, header.obu_type,
 801                                       data, obu_length, frag->data_ref);
 802         if (err < 0)
 803             goto fail;
 804
 805         data += obu_length;
 806         size -= obu_length;
 807     }
 808
 809     err = 0;
 810 fail:
 811     ctx->trace_enable = trace;
 812     return err;
 813 }
 814
 815 static void cbs_av1_free_tile_data(AV1RawTileData *td)
 816 {
 817     av_buffer_unref(&td->data_ref);
 818 }
 819
 820 static void cbs_av1_free_padding(AV1RawPadding *pd)
 821 {
 822     av_buffer_unref(&pd->payload_ref);
 823 }
 824
 825 static void cbs_av1_free_metadata(AV1RawMetadata *md)
 826 {
 827     switch (md->metadata_type) {
 828     case AV1_METADATA_TYPE_ITUT_T35:
 829         av_buffer_unref(&md->metadata.itut_t35.payload_ref);
 830         break;
 831     }
 832 }
 833
 834 static void cbs_av1_free_obu(void *opaque, uint8_t *content)
 835 {
 836     AV1RawOBU *obu = (AV1RawOBU*)content;
 837
 838     switch (obu->header.obu_type) {
 839     case AV1_OBU_TILE_GROUP:
 840         cbs_av1_free_tile_data(&obu->obu.tile_group.tile_data);
 841         break;
 842     case AV1_OBU_FRAME:
 843         cbs_av1_free_tile_data(&obu->obu.frame.tile_group.tile_data);
 844         break;
 845     case AV1_OBU_TILE_LIST:
 846         cbs_av1_free_tile_data(&obu->obu.tile_list.tile_data);
 847         break;
 848     case AV1_OBU_METADATA:
 849         cbs_av1_free_metadata(&obu->obu.metadata);
 850         break;
 851     case AV1_OBU_PADDING:
 852         cbs_av1_free_padding(&obu->obu.padding);
 853         break;
 854     }
 855
 856     av_freep(&obu);
 857 }
 858
 859 static int cbs_av1_ref_tile_data(CodedBitstreamContext *ctx,
 860                                  CodedBitstreamUnit *unit,
 861                                  GetBitContext *gbc,
 862                                  AV1RawTileData *td)
 863 {
 864     int pos;
 865
 866     pos = get_bits_count(gbc);
 867     if (pos >= 8 * unit->data_size) {
 868         av_log(ctx->log_ctx, AV_LOG_ERROR, "Bitstream ended before "
 869                "any data in tile group (%d bits read).\n", pos);
 870         return AVERROR_INVALIDDATA;
 871     }
 872     // Must be byte-aligned at this point.
 873     av_assert0(pos % 8 == 0);
 874
 875     td->data_ref = av_buffer_ref(unit->data_ref);
 876     if (!td->data_ref)
 877         return AVERROR(ENOMEM);
 878
 879     td->data      = unit->data      + pos / 8;
 880     td->data_size = unit->data_size - pos / 8;
 881
 882     return 0;
 883 }
 884
 885 static int cbs_av1_read_unit(CodedBitstreamContext *ctx,
 886                              CodedBitstreamUnit *unit)
 887 {
 888     CodedBitstreamAV1Context *priv = ctx->priv_data;
 889     AV1RawOBU *obu;
 890     GetBitContext gbc;
 891     int err, start_pos, end_pos;
 892
 893     err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*obu),
 894                                     &cbs_av1_free_obu);
 895     if (err < 0)
 896         return err;
 897     obu = unit->content;
 898
 899     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 900     if (err < 0)
 901         return err;
 902
 903     err = cbs_av1_read_obu_header(ctx, &gbc, &obu->header);
 904     if (err < 0)
 905         return err;
 906     av_assert0(obu->header.obu_type == unit->type);
 907
 908     if (obu->header.obu_has_size_field) {
 909         uint64_t obu_size;
 910         err = cbs_av1_read_leb128(ctx, &gbc, "obu_size", &obu_size);
 911         if (err < 0)
 912             return err;
 913         obu->obu_size = obu_size;
 914     } else {
 915         if (unit->data_size < 1 + obu->header.obu_extension_flag) {
 916             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid OBU length: "
 917                    "unit too short (%"SIZE_SPECIFIER").\n", unit->data_size);
 918             return AVERROR_INVALIDDATA;
 919         }
 920         obu->obu_size = unit->data_size - 1 - obu->header.obu_extension_flag;
 921     }
 922
 923     start_pos = get_bits_count(&gbc);
 924
 925     if (obu->header.obu_extension_flag) {
 926         priv->temporal_id = obu->header.temporal_id;
 927         priv->spatial_id  = obu->header.spatial_id;
 928
 929         if (obu->header.obu_type != AV1_OBU_SEQUENCE_HEADER &&
 930             obu->header.obu_type != AV1_OBU_TEMPORAL_DELIMITER &&
 931             priv->operating_point_idc) {
 932             int in_temporal_layer =
 933                 (priv->operating_point_idc >>  priv->temporal_id    ) & 1;
 934             int in_spatial_layer  =
 935                 (priv->operating_point_idc >> (priv->spatial_id + 8)) & 1;
 936             if (!in_temporal_layer || !in_spatial_layer) {
 937                 // Decoding will drop this OBU at this operating point.
 938             }
 939         }
 940     } else {
 941         priv->temporal_id = 0;
 942         priv->spatial_id  = 0;
 943     }
 944
 945     priv->ref = (AV1ReferenceFrameState *)&priv->read_ref;
 946
 947     switch (obu->header.obu_type) {
 948     case AV1_OBU_SEQUENCE_HEADER:
 949         {
 950             err = cbs_av1_read_sequence_header_obu(ctx, &gbc,
 951                                                    &obu->obu.sequence_header);
 952             if (err < 0)
 953                 return err;
 954
 955             av_buffer_unref(&priv->sequence_header_ref);
 956             priv->sequence_header = NULL;
 957
 958             priv->sequence_header_ref = av_buffer_ref(unit->content_ref);
 959             if (!priv->sequence_header_ref)
 960                 return AVERROR(ENOMEM);
 961             priv->sequence_header = &obu->obu.sequence_header;
 962         }
 963         break;
 964     case AV1_OBU_TEMPORAL_DELIMITER:
 965         {
 966             err = cbs_av1_read_temporal_delimiter_obu(ctx, &gbc);
 967             if (err < 0)
 968                 return err;
 969         }
 970         break;
 971     case AV1_OBU_FRAME_HEADER:
 972     case AV1_OBU_REDUNDANT_FRAME_HEADER:
 973         {
 974             err = cbs_av1_read_frame_header_obu(ctx, &gbc,
 975                                                 &obu->obu.frame_header,
 976                                                 obu->header.obu_type ==
 977                                                 AV1_OBU_REDUNDANT_FRAME_HEADER,
 978                                                 unit->data_ref);
 979             if (err < 0)
 980                 return err;
 981         }
 982         break;
 983     case AV1_OBU_TILE_GROUP:
 984         {
 985             err = cbs_av1_read_tile_group_obu(ctx, &gbc,
 986                                               &obu->obu.tile_group);
 987             if (err < 0)
 988                 return err;
 989
 990             err = cbs_av1_ref_tile_data(ctx, unit, &gbc,
 991                                         &obu->obu.tile_group.tile_data);
 992             if (err < 0)
 993                 return err;
 994         }
 995         break;
 996     case AV1_OBU_FRAME:
 997         {
 998             err = cbs_av1_read_frame_obu(ctx, &gbc, &obu->obu.frame,
 999                                          unit->data_ref);
1000             if (err < 0)
1001                 return err;
1002
1003             err = cbs_av1_ref_tile_data(ctx, unit, &gbc,
1004                                         &obu->obu.frame.tile_group.tile_data);
1005             if (err < 0)
1006                 return err;
1007         }
1008         break;
1009     case AV1_OBU_TILE_LIST:
1010         {
1011             err = cbs_av1_read_tile_list_obu(ctx, &gbc,
1012                                              &obu->obu.tile_list);
1013             if (err < 0)
1014                 return err;
1015
1016             err = cbs_av1_ref_tile_data(ctx, unit, &gbc,
1017                                         &obu->obu.tile_list.tile_data);
1018             if (err < 0)
1019                 return err;
1020         }
1021         break;
1022     case AV1_OBU_METADATA:
1023         {
1024             err = cbs_av1_read_metadata_obu(ctx, &gbc, &obu->obu.metadata);
1025             if (err < 0)
1026                 return err;
1027         }
1028         break;
1029     case AV1_OBU_PADDING:
1030         {
1031             err = cbs_av1_read_padding_obu(ctx, &gbc, &obu->obu.padding);
1032             if (err < 0)
1033                 return err;
1034         }
1035         break;
1036     default:
1037         return AVERROR(ENOSYS);
1038     }
1039
1040     end_pos = get_bits_count(&gbc);
1041     av_assert0(end_pos <= unit->data_size * 8);
1042
1043     if (obu->obu_size > 0 &&
1044         obu->header.obu_type != AV1_OBU_TILE_GROUP &&
1045         obu->header.obu_type != AV1_OBU_TILE_LIST &&
1046         obu->header.obu_type != AV1_OBU_FRAME) {
1047         int nb_bits = obu->obu_size * 8 + start_pos - end_pos;
1048
1049         if (nb_bits <= 0)
1050             return AVERROR_INVALIDDATA;
1051
1052         err = cbs_av1_read_trailing_bits(ctx, &gbc, nb_bits);
1053         if (err < 0)
1054             return err;
1055     }
1056
1057     return 0;
1058 }
1059
1060 static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
1061                              CodedBitstreamUnit *unit,
1062                              PutBitContext *pbc)
1063 {
1064     CodedBitstreamAV1Context *priv = ctx->priv_data;
1065     AV1RawOBU *obu = unit->content;
1066     PutBitContext pbc_tmp;
1067     AV1RawTileData *td;
1068     size_t header_size;
1069     int err, start_pos, end_pos, data_pos;
1070
1071     // OBUs in the normal bitstream format must contain a size field
1072     // in every OBU (in annex B it is optional, but we don't support
1073     // writing that).
1074     obu->header.obu_has_size_field = 1;
1075
1076     err = cbs_av1_write_obu_header(ctx, pbc, &obu->header);
1077     if (err < 0)
1078         return err;
1079
1080     if (obu->header.obu_has_size_field) {
1081         pbc_tmp = *pbc;
1082         // Add space for the size field to fill later.
1083         put_bits32(pbc, 0);
1084         put_bits32(pbc, 0);
1085     }
1086
1087     td = NULL;
1088     start_pos = put_bits_count(pbc);
1089
1090     priv->ref = (AV1ReferenceFrameState *)&priv->write_ref;
1091
1092     switch (obu->header.obu_type) {
1093     case AV1_OBU_SEQUENCE_HEADER:
1094         {
1095             err = cbs_av1_write_sequence_header_obu(ctx, pbc,
1096                                                     &obu->obu.sequence_header);
1097             if (err < 0)
1098                 return err;
1099
1100             av_buffer_unref(&priv->sequence_header_ref);
1101             priv->sequence_header = NULL;
1102
1103             priv->sequence_header_ref = av_buffer_ref(unit->content_ref);
1104             if (!priv->sequence_header_ref)
1105                 return AVERROR(ENOMEM);
1106             priv->sequence_header = &obu->obu.sequence_header;
1107         }
1108         break;
1109     case AV1_OBU_TEMPORAL_DELIMITER:
1110         {
1111             err = cbs_av1_write_temporal_delimiter_obu(ctx, pbc);
1112             if (err < 0)
1113                 return err;
1114         }
1115         break;
1116     case AV1_OBU_FRAME_HEADER:
1117     case AV1_OBU_REDUNDANT_FRAME_HEADER:
1118         {
1119             err = cbs_av1_write_frame_header_obu(ctx, pbc,
1120                                                  &obu->obu.frame_header,
1121                                                  obu->header.obu_type ==
1122                                                  AV1_OBU_REDUNDANT_FRAME_HEADER,
1123                                                  NULL);
1124             if (err < 0)
1125                 return err;
1126         }
1127         break;
1128     case AV1_OBU_TILE_GROUP:
1129         {
1130             err = cbs_av1_write_tile_group_obu(ctx, pbc,
1131                                                &obu->obu.tile_group);
1132             if (err < 0)
1133                 return err;
1134
1135             td = &obu->obu.tile_group.tile_data;
1136         }
1137         break;
1138     case AV1_OBU_FRAME:
1139         {
1140             err = cbs_av1_write_frame_obu(ctx, pbc, &obu->obu.frame, NULL);
1141             if (err < 0)
1142                 return err;
1143
1144             td = &obu->obu.frame.tile_group.tile_data;
1145         }
1146         break;
1147     case AV1_OBU_TILE_LIST:
1148         {
1149             err = cbs_av1_write_tile_list_obu(ctx, pbc, &obu->obu.tile_list);
1150             if (err < 0)
1151                 return err;
1152
1153             td = &obu->obu.tile_list.tile_data;
1154         }
1155         break;
1156     case AV1_OBU_METADATA:
1157         {
1158             err = cbs_av1_write_metadata_obu(ctx, pbc, &obu->obu.metadata);
1159             if (err < 0)
1160                 return err;
1161         }
1162         break;
1163     case AV1_OBU_PADDING:
1164         {
1165             err = cbs_av1_write_padding_obu(ctx, pbc, &obu->obu.padding);
1166             if (err < 0)
1167                 return err;
1168         }
1169         break;
1170     default:
1171         return AVERROR(ENOSYS);
1172     }
1173
1174     end_pos = put_bits_count(pbc);
1175     header_size = (end_pos - start_pos + 7) / 8;
1176     if (td) {
1177         obu->obu_size = header_size + td->data_size;
1178     } else if (header_size > 0) {
1179         // Add trailing bits and recalculate.
1180         err = cbs_av1_write_trailing_bits(ctx, pbc, 8 - end_pos % 8);
1181         if (err < 0)
1182             return err;
1183         end_pos = put_bits_count(pbc);
1184         obu->obu_size = header_size = (end_pos - start_pos + 7) / 8;
1185     } else {
1186         // Empty OBU.
1187         obu->obu_size = 0;
1188     }
1189
1190     end_pos = put_bits_count(pbc);
1191     // Must now be byte-aligned.
1192     av_assert0(end_pos % 8 == 0);
1193     flush_put_bits(pbc);
1194     start_pos /= 8;
1195     end_pos   /= 8;
1196
1197     *pbc = pbc_tmp;
1198     err = cbs_av1_write_leb128(ctx, pbc, "obu_size", obu->obu_size);
1199     if (err < 0)
1200         return err;
1201
1202     data_pos = put_bits_count(pbc) / 8;
1203     flush_put_bits(pbc);
1204     av_assert0(data_pos <= start_pos);
1205
1206     if (8 * obu->obu_size > put_bits_left(pbc))
1207         return AVERROR(ENOSPC);
1208
1209     if (obu->obu_size > 0) {
1210         memmove(pbc->buf + data_pos,
1211                 pbc->buf + start_pos, header_size);
1212         skip_put_bytes(pbc, header_size);
1213
1214         if (td) {
1215             memcpy(pbc->buf + data_pos + header_size,
1216                    td->data, td->data_size);
1217             skip_put_bytes(pbc, td->data_size);
1218         }
1219     }
1220
1221     // OBU data must be byte-aligned.
1222     av_assert0(put_bits_count(pbc) % 8 == 0);
1223
1224     return 0;
1225 }
1226
1227 static int cbs_av1_assemble_fragment(CodedBitstreamContext *ctx,
1228                                      CodedBitstreamFragment *frag)
1229 {
1230     size_t size, pos;
1231     int i;
1232
1233     size = 0;
1234     for (i = 0; i < frag->nb_units; i++)
1235         size += frag->units[i].data_size;
1236
1237     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
1238     if (!frag->data_ref)
1239         return AVERROR(ENOMEM);
1240     frag->data = frag->data_ref->data;
1241     memset(frag->data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
1242
1243     pos = 0;
1244     for (i = 0; i < frag->nb_units; i++) {
1245         memcpy(frag->data + pos, frag->units[i].data,
1246                frag->units[i].data_size);
1247         pos += frag->units[i].data_size;
1248     }
1249     av_assert0(pos == size);
1250     frag->data_size = size;
1251
1252     return 0;
1253 }
1254
1255 static void cbs_av1_close(CodedBitstreamContext *ctx)
1256 {
1257     CodedBitstreamAV1Context *priv = ctx->priv_data;
1258
1259     av_buffer_unref(&priv->sequence_header_ref);
1260     av_buffer_unref(&priv->frame_header_ref);
1261 }
1262
1263 const CodedBitstreamType ff_cbs_type_av1 = {
1264     .codec_id          = AV_CODEC_ID_AV1,
1265
1266     .priv_data_size    = sizeof(CodedBitstreamAV1Context),
1267
1268     .split_fragment    = &cbs_av1_split_fragment,
1269     .read_unit         = &cbs_av1_read_unit,
1270     .write_unit        = &cbs_av1_write_obu,
1271     .assemble_fragment = &cbs_av1_assemble_fragment,
1272
1273     .close             = &cbs_av1_close,
1274 };