1"""A picture that stands in for an H.264 frame a peer is not sent: every macroblock skipped, so it
2repeats the picture before it, under the left-out frame's own frame_num, picture order, reference
3list modification, weights, and reference marking. A decoder that refuses a frame_num gap
4(OpenH264, which Firefox decodes WebRTC H.264 with where it has no system FFmpeg) then takes the
5stream whole, its decoded picture buffer the encoder's in every structure but the pixels of the
6left-out frames, which the encoder predicts past (`RTCRtpSender._repair`).
7
8Only the slice header's entropy context index, quantizer, and deblocking change: `cabac_init_idc`
90 and `SKIP_QP` let one arithmetic-coded run of skips serve every slice of a size, and a picture of
10skips has no edge to filter. A slice this cannot rewrite (an I slice, slice groups, separate colour
11planes, a redundant picture) gets no stand-in.
12"""
13from typing import Dict, List, Optional, Tuple
14
15# Table 9-44: rangeTabLPS, by pStateIdx and qCodIRangeIdx.
16RANGE_LPS = (
17 (128, 176, 208, 240), (128, 167, 197, 227), (128, 158, 187, 216), (123, 150, 178, 205),
18 (116, 142, 169, 195), (111, 135, 160, 185), (105, 128, 152, 175), (100, 122, 144, 166),
19 (95, 116, 137, 158), (90, 110, 130, 150), (85, 104, 123, 142), (81, 99, 117, 135),
20 (77, 94, 111, 128), (73, 89, 105, 122), (69, 85, 100, 116), (66, 80, 95, 110),
21 (62, 76, 90, 104), (59, 72, 86, 99), (56, 69, 81, 94), (53, 65, 77, 89),
22 (51, 62, 73, 85), (48, 59, 69, 80), (46, 56, 66, 76), (43, 53, 63, 72),
23 (41, 50, 59, 69), (39, 48, 56, 65), (37, 45, 54, 62), (35, 43, 51, 59),
24 (33, 41, 48, 56), (32, 39, 46, 53), (30, 37, 43, 50), (29, 35, 41, 48),
25 (27, 33, 39, 45), (26, 31, 37, 43), (24, 30, 35, 41), (23, 28, 33, 39),
26 (22, 27, 32, 37), (21, 26, 30, 35), (20, 24, 29, 33), (19, 23, 27, 31),
27 (18, 22, 26, 30), (17, 21, 25, 28), (16, 20, 23, 27), (15, 19, 22, 25),
28 (14, 18, 21, 24), (14, 17, 20, 23), (13, 16, 19, 22), (12, 15, 18, 21),
29 (12, 14, 17, 20), (11, 14, 16, 19), (11, 13, 15, 18), (10, 12, 15, 17),
30 (10, 12, 14, 16), (9, 11, 13, 15), (9, 11, 12, 14), (8, 10, 12, 14),
31 (8, 9, 11, 13), (7, 9, 11, 12), (7, 9, 10, 12), (7, 8, 10, 11),
32 (6, 8, 9, 11), (6, 7, 9, 10), (6, 7, 8, 9), (2, 2, 2, 2),
33)
34# Table 9-45: transIdxLPS, by pStateIdx.
35TRANS_LPS = (
36 0, 0, 1, 2, 2, 4, 4, 5, 6, 7, 8, 9, 9, 11, 11, 12, 13, 13, 15, 15, 16, 16, 18, 18, 19, 19, 21, 21, 22, 22,
37 23, 24, 24, 25, 26, 26, 27, 27, 28, 29, 29, 30, 30, 30, 31, 32, 32, 33, 33, 33, 34, 34, 35, 35, 35, 36, 36,
38 36, 37, 37, 37, 38, 38, 63,
39)
40# Table 9-13: (m, n) of mb_skip_flag's context in a P slice whose neighbours are skipped or absent
41# (ctxIdx 11), at cabac_init_idc 0.
42SKIP_CONTEXT = (23, 33)
43# The quantizer a stand-in declares. It codes no residual, so only the arithmetic coder's start reads it.
44SKIP_QP = 26
45P_SLICE, I_SLICE = 0, 2
46NAL_SLICE, NAL_IDR, NAL_SPS, NAL_PPS = 1, 5, 7, 8
47# The bytes of a slice a stand-in is read from: its header, which encoders keep to a few dozen bytes.
48# A header running past them gets no stand-in.
49HEAD_BYTES = 256
50
51
52class _Reader:
53 """Exp-Golomb and fixed-width reads over a NAL unit's payload, emulation prevention removed."""
54
55 def __init__(self, nal: bytes) -> None:
56 rbsp, zeros = bytearray(), 0
57 for byte in nal[1:]:
58 if zeros >= 2 and byte == 3:
59 zeros = 0
60 continue
61 zeros = zeros + 1 if byte == 0 else 0
62 rbsp.append(byte)
63 self.value = int.from_bytes(rbsp, "big") if rbsp else 0
64 self.size = len(rbsp) * 8
65 self.pos = 0
66
67 def u(self, n: int) -> int:
68 if n == 0:
69 return 0
70 if self.pos + n > self.size:
71 raise ValueError("read past the end of the NAL unit")
72 self.pos += n
73 return (self.value >> (self.size - self.pos)) & ((1 << n) - 1)
74
75 def ue(self) -> int:
76 zeros = 0
77 while self.u(1) == 0:
78 zeros += 1
79 if zeros > 31:
80 raise ValueError("exp-Golomb code too long")
81 return (1 << zeros) - 1 + self.u(zeros)
82
83 def se(self) -> int:
84 k = self.ue()
85 return (k + 1) // 2 if k % 2 else -(k // 2)
86
87
88class _Writer:
89 """Bits gathered into an RBSP, written out as a NAL unit with emulation prevention."""
90
91 def __init__(self) -> None:
92 self.bits: List[str] = []
93
94 def u(self, n: int, v: int) -> None:
95 if n:
96 self.bits.append(format(v, f"0{n}b"))
97
98 def ue(self, v: int) -> None:
99 code = v + 1
100 self.bits.append("0" * (code.bit_length() - 1) + format(code, "b"))
101
102 def se(self, v: int) -> None:
103 self.ue(2 * v - 1 if v > 0 else -2 * v)
104
105 def raw(self, bits: str) -> None:
106 self.bits.append(bits)
107
108 def length(self) -> int:
109 return sum(len(b) for b in self.bits)
110
111 def nal(self, header: int) -> bytes:
112 bits = "".join(self.bits)
113 bits += "0" * (-len(bits) % 8)
114 rbsp = int(bits, 2).to_bytes(len(bits) // 8, "big") if bits else b""
115 out, zeros = bytearray([header]), 0
116 for byte in rbsp:
117 if zeros >= 2 and byte <= 3:
118 out.append(3)
119 zeros = 0
120 out.append(byte)
121 zeros = zeros + 1 if byte == 0 else 0
122 return bytes(out)
123
124
125def _skip_scaling_list(r: _Reader, size: int) -> None:
126 last = following = 8
127 for _ in range(size):
128 if following:
129 following = (last + r.se() + 256) % 256
130 last = following or last
131
132
133def parse_sps(nal: bytes) -> Tuple[int, dict]:
134 """The fields of a sequence parameter set a slice header's syntax depends on, by its id."""
135 r = _Reader(nal)
136 profile = r.u(8)
137 r.u(16)
138 sps_id = r.ue()
139 sps = {"chroma_format_idc": 1, "separate_colour_plane": 0}
140 if profile in (100, 110, 122, 244, 44, 83, 86, 118, 128, 138, 139, 134, 135):
141 sps["chroma_format_idc"] = r.ue()
142 if sps["chroma_format_idc"] == 3:
143 sps["separate_colour_plane"] = r.u(1)
144 r.ue()
145 r.ue()
146 r.u(1)
147 if r.u(1):
148 for i in range(8 if sps["chroma_format_idc"] != 3 else 12):
149 if r.u(1):
150 _skip_scaling_list(r, 16 if i < 6 else 64)
151 sps["log2_max_frame_num"] = r.ue() + 4
152 sps["poc_type"] = r.ue()
153 if sps["poc_type"] == 0:
154 sps["log2_max_poc_lsb"] = r.ue() + 4
155 elif sps["poc_type"] == 1:
156 sps["delta_pic_order_always_zero"] = r.u(1)
157 r.se()
158 r.se()
159 for _ in range(r.ue()):
160 r.se()
161 r.ue()
162 r.u(1)
163 width = r.ue() + 1
164 height_units = r.ue() + 1
165 sps["frame_mbs_only"] = r.u(1)
166 sps["mbs"] = width * height_units * (2 - sps["frame_mbs_only"])
167 return sps_id, sps
168
169
170def parse_pps(nal: bytes) -> Tuple[int, dict]:
171 """The fields of a picture parameter set a slice header's syntax depends on, by its id."""
172 r = _Reader(nal)
173 pps_id = r.ue()
174 pps = {"sps_id": r.ue(), "cabac": r.u(1), "bottom_field_poc": r.u(1), "slice_groups": r.ue() + 1}
175 if pps["slice_groups"] > 1:
176 return pps_id, pps
177 pps["num_ref_idx_l0_default"] = r.ue() + 1
178 r.ue()
179 pps["weighted_pred"] = r.u(1)
180 r.u(2)
181 pps["pic_init_qp"] = 26 + r.se()
182 r.se()
183 r.se()
184 pps["deblocking_control"] = r.u(1)
185 r.u(1)
186 pps["redundant_pic_cnt"] = r.u(1)
187 return pps_id, pps
188
189
190def _cabac_skips(mbs: int) -> str:
191 """The slice data, from its first byte, of `mbs` macroblocks each skipped, the last ending the
192 slice: mb_skip_flag 1 and end_of_slice_flag 0 for each, the arithmetic coder flushed after the
193 last end_of_slice_flag 1 (9.3.4), its stop bit included."""
194 m, n = SKIP_CONTEXT
195 pre = max(1, min(126, ((m * SKIP_QP) >> 4) + n))
196 state, mps = (63 - pre, 0) if pre <= 63 else (pre - 64, 1)
197 low, rng, outstanding, first = 0, 510, 0, True
198 bits: List[str] = []
199
200 def put(bit: int) -> None:
201 nonlocal first, outstanding
202 if first:
203 first = False
204 else:
205 bits.append("1" if bit else "0")
206 if outstanding:
207 bits.append(("0" if bit else "1") * outstanding)
208 outstanding = 0
209
210 def renorm() -> None:
211 nonlocal low, rng, outstanding
212 while rng < 256:
213 if low < 256:
214 put(0)
215 elif low >= 512:
216 low -= 512
217 put(1)
218 else:
219 low -= 256
220 outstanding += 1
221 rng <<= 1
222 low <<= 1
223
224 for i in range(mbs):
225 lps = RANGE_LPS[state][(rng >> 6) & 3]
226 rng -= lps
227 if mps != 1:
228 low += rng
229 rng = lps
230 if state == 0:
231 mps = 1 - mps
232 state = TRANS_LPS[state]
233 else:
234 state = min(state + 1, 62)
235 renorm()
236 rng -= 2
237 if i < mbs - 1:
238 renorm()
239 low += rng
240 rng = 2
241 renorm()
242 put((low >> 9) & 1)
243 bits.append(format(((low >> 7) & 3) | 1, "02b"))
244 return "".join(bits)
245
246
247class SkipPictures:
248 """Stand-ins for the frames of one stream a peer is not sent (`stand_in`), read against the
249 parameter sets its key frames carried (`learn`)."""
250
251 def __init__(self) -> None:
252 self._sps: Dict[int, dict] = {}
253 self._pps: Dict[int, dict] = {}
254 self._cabac: Dict[int, str] = {}
255
256 def learn(self, nals: List[bytes]) -> None:
257 """Take the SPS and PPS among a frame's NAL units."""
258 for nal in nals:
259 kind = nal[0] & 0x1F if nal else 0
260 try:
261 if kind == NAL_SPS:
262 sps_id, sps = parse_sps(nal)
263 self._sps[sps_id] = sps
264 elif kind == NAL_PPS:
265 pps_id, pps = parse_pps(nal)
266 self._pps[pps_id] = pps
267 except ValueError:
268 continue
269
270 def stand_in(self, nals: List[bytes]) -> Optional[List[bytes]]:
271 """A skipped slice for each slice of the frame `nals`, or None where one cannot be made.
272 Only a slice's first HEAD_BYTES are read, so the caller may pass no more."""
273 slices = [nal[:HEAD_BYTES] for nal in nals if nal and nal[0] & 0x1F == NAL_SLICE]
274 if not slices or any(nal[0] & 0x1F == NAL_IDR for nal in nals):
275 return None
276 try:
277 starts = [_Reader(nal).ue() for nal in slices]
278 out = []
279 for k, nal in enumerate(slices):
280 built = self._skip_slice(nal, starts[k], starts[k + 1] if k + 1 < len(slices) else None)
281 if built is None:
282 return None
283 out.append(built)
284 return out
285 except (ValueError, KeyError, IndexError):
286 return None
287
288 def _skip_slice(self, nal: bytes, first_mb: int, next_mb: Optional[int]) -> Optional[bytes]:
289 r, w = _Reader(nal), _Writer()
290 w.ue(r.ue())
291 slice_type = r.ue()
292 if slice_type % 5 != P_SLICE:
293 return None
294 w.ue(slice_type)
295 pps_id = r.ue()
296 pps = self._pps[pps_id]
297 if pps["slice_groups"] > 1:
298 return None
299 sps = self._sps[pps["sps_id"]]
300 if sps["separate_colour_plane"]:
301 return None
302 w.ue(pps_id)
303 w.u(sps["log2_max_frame_num"], r.u(sps["log2_max_frame_num"]))
304 field = 0
305 if not sps["frame_mbs_only"]:
306 field = r.u(1)
307 w.u(1, field)
308 if field:
309 w.u(1, r.u(1))
310 if sps["poc_type"] == 0:
311 w.u(sps["log2_max_poc_lsb"], r.u(sps["log2_max_poc_lsb"]))
312 if pps["bottom_field_poc"] and not field:
313 w.se(r.se())
314 elif sps["poc_type"] == 1 and not sps["delta_pic_order_always_zero"]:
315 w.se(r.se())
316 if pps["bottom_field_poc"] and not field:
317 w.se(r.se())
318 if pps["redundant_pic_cnt"]:
319 redundant = r.ue()
320 if redundant:
321 return None
322 w.ue(redundant)
323 refs = pps["num_ref_idx_l0_default"]
324 override = r.u(1)
325 w.u(1, override)
326 if override:
327 refs = r.ue() + 1
328 w.ue(refs - 1)
329 modified = r.u(1)
330 w.u(1, modified)
331 while modified:
332 idc = r.ue()
333 w.ue(idc)
334 if idc in (0, 1, 2):
335 w.ue(r.ue())
336 elif idc == 3:
337 break
338 else:
339 return None
340 if pps["weighted_pred"]:
341 chroma = sps["chroma_format_idc"] != 0
342 w.ue(r.ue())
343 if chroma:
344 w.ue(r.ue())
345 for _ in range(refs):
346 flag = r.u(1)
347 w.u(1, flag)
348 if flag:
349 w.se(r.se())
350 w.se(r.se())
351 if chroma:
352 flag = r.u(1)
353 w.u(1, flag)
354 if flag:
355 for _ in range(4):
356 w.se(r.se())
357 if nal[0] & 0x60:
358 adaptive = r.u(1)
359 w.u(1, adaptive)
360 while adaptive:
361 op = r.ue()
362 w.ue(op)
363 if op == 0:
364 break
365 if op in (1, 3):
366 w.ue(r.ue())
367 if op == 2:
368 w.ue(r.ue())
369 if op in (3, 6):
370 w.ue(r.ue())
371 if op == 4:
372 w.ue(r.ue())
373 if pps["cabac"]:
374 r.ue()
375 w.ue(0)
376 r.se()
377 w.se(SKIP_QP - pps["pic_init_qp"])
378 if pps["deblocking_control"]:
379 if r.ue() != 1:
380 r.se()
381 r.se()
382 w.ue(1)
383 mbs = (next_mb if next_mb is not None else sps["mbs"]) - first_mb
384 if mbs <= 0:
385 return None
386 if pps["cabac"]:
387 w.raw("1" * (-w.length() % 8))
388 data = self._cabac.get(mbs)
389 if data is None:
390 data = self._cabac[mbs] = _cabac_skips(mbs)
391 w.raw(data)
392 else:
393 w.ue(mbs)
394 w.raw("1")
395 return w.nal(nal[0])