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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
//! The arithmetic decoder, described in Annex C.
//!
//! The arithmetic decoder keeps track of some state and continuously receives
//! context labels as input, each time yielding a new bit from the original data
//! as output.
use super::mq::QE_TABLE;
pub(crate) struct ArithmeticDecoder<'a> {
/// The underlying encoded data.
data: &'a [u8],
/// The C-register (see Table C.1).
c: u32,
/// The A-register (see Table C.1).
a: u32,
/// The pointer to the current byte.
base_pointer: u32,
/// The bit shift counter.
shift_count: u32,
}
impl<'a> ArithmeticDecoder<'a> {
pub(crate) fn new(data: &'a [u8]) -> Self {
let mut decoder = ArithmeticDecoder {
data,
c: 0,
a: 0,
base_pointer: 0,
shift_count: 0,
};
decoder.initialize();
decoder
}
/// Read the next bit using the given context label.
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
pub(crate) fn read_bit(&mut self, context: &mut ArithmeticDecoderContext) -> u32 {
self.decode(context)
}
/// The INITDEC procedure from C.3.5.
///
/// We use the version from Annex G in <https://www.itu.int/rec/T-REC-T.88-201808-I>.
pub(crate) fn initialize(&mut self) {
self.c = (u32::from(self.current_byte()) ^ 0xff) << 16;
self.read_byte();
self.c <<= 7;
self.shift_count -= 7;
self.a = 0x8000;
}
/// The BYTEIN procedure from C.3.4.
///
/// We use the version from Annex G from <https://www.itu.int/rec/T-REC-T.88-201808-I>.
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
fn read_byte(&mut self) {
if self.current_byte() == 0xff {
let b1 = self.next_byte();
if b1 > 0x8f {
self.shift_count = 8;
} else {
self.base_pointer += 1;
self.c = self
.c
.wrapping_add(0xfe00)
.wrapping_sub(u32::from(self.current_byte()) << 9);
self.shift_count = 7;
}
} else {
self.base_pointer += 1;
self.c = self
.c
.wrapping_add(0xff00)
.wrapping_sub(u32::from(self.current_byte()) << 8);
self.shift_count = 8;
}
}
/// The RENORMD procedure from C.3.3.
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
fn renormalize(&mut self) {
// Original code:
// loop {
// if self.shift_count == 0 {
// self.read_byte();
// }
//
// self.a <<= 1;
// self.c <<= 1;
// self.shift_count -= 1;
//
// if self.a & 0x8000 != 0 {
// break;
// }
// }
// Optimization: Batch shifts.
while self.a & 0x8000 == 0 {
if self.shift_count == 0 {
self.read_byte();
}
let shifts_needed = self.a.leading_zeros() - 16;
let batch = shifts_needed.min(self.shift_count);
self.a <<= batch;
self.c <<= batch;
self.shift_count -= batch;
}
}
/// The DECODE procedure from C.3.2.
///
/// We use the version from Annex G from <https://www.itu.int/rec/T-REC-T.88-201808-I>.
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
pub(crate) fn decode(&mut self, context: &mut ArithmeticDecoderContext) -> u32 {
let qe_entry = &QE_TABLE[context.index() as usize];
self.a -= qe_entry.qe;
// This is a faster version that reduces branching, which has shown
// itself to be the main limiting factor for better performance.
// We short-circuit the case where just the most probably symbol is
// returned, and otherwise use a code path that works for both,
// MPS_EXCHANGE and LPS_EXCHANGE.
if (self.c >> 16) < self.a && self.a & 0x8000 != 0 {
return context.mps();
}
// Unified branchless MPS_EXCHANGE / LPS_EXCHANGE. In the Annex C.3.2
// procedures, the only difference is that LPS flips the role of cond:
// exchange_mps: d = mps ^ cond, flip when cond, index = cond*nlps + inv*nmps
// exchange_lps: d = mps ^ inv_cond, flip when inv_cond, index = cond*nmps + inv*nlps
//
// This is equivalent to XOR-ing cond with is_lps, so we can handle
// both paths with a single branchless computation.
//
// As can be seen above, renormalization is always performed.
let is_lps = u32::from((self.c >> 16) >= self.a);
// LPS: C -= A << 16 (no-op when MPS).
let lps_mask = is_lps.wrapping_neg(); // 0xFFFF_FFFF if LPS, 0 if MPS
self.c -= (self.a << 16) & lps_mask;
// Same condition as in exchange_mps / exchange_lps.
let cond = u32::from(self.a < qe_entry.qe);
// LPS: a = qe (no-op when MPS, a stays as a - qe).
self.a = (self.a & !lps_mask) | (qe_entry.qe & lps_mask);
// exchange_mps: d = mps ^ cond → cond ^ 0
// exchange_lps: d = mps ^ inv_cond → cond ^ 1
// unified: d = mps ^ (cond ^ is_lps)
let d = context.mps() ^ cond ^ is_lps;
// exchange_mps: flip mps when cond & switch → (cond ^ 0) & switch
// exchange_lps: flip mps when inv_cond & switch → (cond ^ 1) & switch
// unified: flip mps when (cond ^ is_lps) & switch
context.xor_mps((cond ^ is_lps) & u32::from(qe_entry.switch));
// exchange_mps: index = cond * nlps + inv_cond * nmps
// exchange_lps: index = cond * nmps + inv_cond * nlps (swapped)
// unified: the result is always exactly nmps or nlps —
// pick nlps when (cond ^ is_lps) == 1, nmps otherwise.
let pick_nlps = u8::from(cond ^ is_lps != 0).wrapping_neg(); // 0xFF or 0x00
context.set_index(qe_entry.nmps ^ ((qe_entry.nmps ^ qe_entry.nlps) & pick_nlps));
self.renormalize();
d
}
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
fn current_byte(&self) -> u8 {
self.data
.get(self.base_pointer as usize)
.copied()
// "The number of bytes corresponding to the coding passes is
// specified in the packet header. Often at that point there are
// more symbols to be decoded. Therefore, the decoder shall extend
// the input bit stream to the arithmetic coder with 0xFF bytes,
// as necessary, until all symbols have been decoded."
.unwrap_or(0xFF)
}
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
fn next_byte(&self) -> u8 {
self.data
.get((self.base_pointer + 1) as usize)
.copied()
.unwrap_or(0xFF)
}
}
// Previously, we stored the context as 2 u32's, but doing it with a bit-packed
// u8 seems to be slightly better (though it doesn't make that huge of a
// difference).
/// Bits 0-6 = index (0-46).
/// Bit 7 = mps (0 or 1).
#[derive(Copy, Clone, Debug, Default)]
pub(crate) struct ArithmeticDecoderContext(u8);
impl ArithmeticDecoderContext {
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
pub(crate) fn index(self) -> u32 {
u32::from(self.0 & 0x7F)
}
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
pub(crate) fn mps(self) -> u32 {
u32::from(self.0 >> 7)
}
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
fn set_index(&mut self, index: u8) {
self.0 = (self.0 & 0x80) | index;
}
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
fn xor_mps(&mut self, val: u32) {
self.0 ^= u8::from(val & 1 != 0) << 7;
}
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
pub(crate) fn reset(&mut self) {
self.0 = 0;
}
#[expect(
clippy::inline_always,
reason = "MQ state transitions are measured per-symbol hot paths"
)]
#[inline(always)]
pub(crate) fn reset_with_index(&mut self, index: u8) {
self.0 = index;
}
}