summaryrefslogtreecommitdiffstats
path: root/third_party/jpeg-xl/lib/jxl/coeff_order.cc
blob: 43adafd82a35e13c16e7996837313e8ed3fc6e97 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "lib/jxl/coeff_order.h"

#include <stdint.h>

#include <algorithm>
#include <vector>

#include "lib/jxl/ans_params.h"
#include "lib/jxl/base/padded_bytes.h"
#include "lib/jxl/base/profiler.h"
#include "lib/jxl/base/span.h"
#include "lib/jxl/coeff_order_fwd.h"
#include "lib/jxl/dec_ans.h"
#include "lib/jxl/dec_bit_reader.h"
#include "lib/jxl/entropy_coder.h"
#include "lib/jxl/lehmer_code.h"
#include "lib/jxl/modular/encoding/encoding.h"
#include "lib/jxl/modular/modular_image.h"

namespace jxl {

uint32_t CoeffOrderContext(uint32_t val) {
  uint32_t token, nbits, bits;
  HybridUintConfig(0, 0, 0).Encode(val, &token, &nbits, &bits);
  return std::min(token, kPermutationContexts - 1);
}

namespace {
Status ReadPermutation(size_t skip, size_t size, coeff_order_t* order,
                       BitReader* br, ANSSymbolReader* reader,
                       const std::vector<uint8_t>& context_map) {
  std::vector<LehmerT> lehmer(size);
  // temp space needs to be as large as the next power of 2, so doubling the
  // allocated size is enough.
  std::vector<uint32_t> temp(size * 2);
  uint32_t end =
      reader->ReadHybridUint(CoeffOrderContext(size), br, context_map) + skip;
  if (end > size) {
    return JXL_FAILURE("Invalid permutation size");
  }
  uint32_t last = 0;
  for (size_t i = skip; i < end; ++i) {
    lehmer[i] =
        reader->ReadHybridUint(CoeffOrderContext(last), br, context_map);
    last = lehmer[i];
    if (lehmer[i] + i >= size) {
      return JXL_FAILURE("Invalid lehmer code");
    }
  }
  if (order == nullptr) return true;
  DecodeLehmerCode(lehmer.data(), temp.data(), size, order);
  return true;
}

}  // namespace

Status DecodePermutation(size_t skip, size_t size, coeff_order_t* order,
                         BitReader* br) {
  std::vector<uint8_t> context_map;
  ANSCode code;
  JXL_RETURN_IF_ERROR(
      DecodeHistograms(br, kPermutationContexts, &code, &context_map));
  ANSSymbolReader reader(&code, br);
  JXL_RETURN_IF_ERROR(
      ReadPermutation(skip, size, order, br, &reader, context_map));
  if (!reader.CheckANSFinalState()) {
    return JXL_FAILURE("Invalid ANS stream");
  }
  return true;
}

namespace {

Status DecodeCoeffOrder(AcStrategy acs, coeff_order_t* order, BitReader* br,
                        ANSSymbolReader* reader,
                        std::vector<coeff_order_t>& natural_order,
                        const std::vector<uint8_t>& context_map) {
  PROFILER_FUNC;
  const size_t llf = acs.covered_blocks_x() * acs.covered_blocks_y();
  const size_t size = kDCTBlockSize * llf;

  JXL_RETURN_IF_ERROR(
      ReadPermutation(llf, size, order, br, reader, context_map));
  if (order == nullptr) return true;
  for (size_t k = 0; k < size; ++k) {
    order[k] = natural_order[order[k]];
  }
  return true;
}

}  // namespace

Status DecodeCoeffOrders(uint16_t used_orders, uint32_t used_acs,
                         coeff_order_t* order, BitReader* br) {
  uint16_t computed = 0;
  std::vector<uint8_t> context_map;
  ANSCode code;
  std::unique_ptr<ANSSymbolReader> reader;
  std::vector<coeff_order_t> natural_order;
  // Bitstream does not have histograms if no coefficient order is used.
  if (used_orders != 0) {
    JXL_RETURN_IF_ERROR(
        DecodeHistograms(br, kPermutationContexts, &code, &context_map));
    reader = make_unique<ANSSymbolReader>(&code, br);
  }
  uint32_t acs_mask = 0;
  for (uint8_t o = 0; o < AcStrategy::kNumValidStrategies; ++o) {
    if ((used_acs & (1 << o)) == 0) continue;
    acs_mask |= 1 << kStrategyOrder[o];
  }
  for (uint8_t o = 0; o < AcStrategy::kNumValidStrategies; ++o) {
    uint8_t ord = kStrategyOrder[o];
    if (computed & (1 << ord)) continue;
    computed |= 1 << ord;
    AcStrategy acs = AcStrategy::FromRawStrategy(o);
    bool used = (acs_mask & (1 << ord)) != 0;

    const size_t llf = acs.covered_blocks_x() * acs.covered_blocks_y();
    const size_t size = kDCTBlockSize * llf;

    if (used || (used_orders & (1 << ord))) {
      if (natural_order.size() < size) natural_order.resize(size);
      acs.ComputeNaturalCoeffOrder(natural_order.data());
    }

    if ((used_orders & (1 << ord)) == 0) {
      // No need to set the default order if no ACS uses this order.
      if (used) {
        for (size_t c = 0; c < 3; c++) {
          memcpy(&order[CoeffOrderOffset(ord, c)], natural_order.data(),
                 size * sizeof(*order));
        }
      }
    } else {
      for (size_t c = 0; c < 3; c++) {
        coeff_order_t* dest = used ? &order[CoeffOrderOffset(ord, c)] : nullptr;
        JXL_RETURN_IF_ERROR(DecodeCoeffOrder(acs, dest, br, reader.get(),
                                             natural_order, context_map));
      }
    }
  }
  if (used_orders && !reader->CheckANSFinalState()) {
    return JXL_FAILURE("Invalid ANS stream");
  }
  return true;
}

}  // namespace jxl