Skip to main content

atmos/os_lib/
h264_slice.rs

1// h264_slice.rs - H.264 スライスヘッダの解析と POC(表示順)の計算(純粋モジュール)。
2//
3// 【2026-10-17】B スライス対応の土台として `h264_decode.rs` から切り出した。B フレームは
4// 復号順と表示順が違い、参照リスト(L0/L1)も POC の大小で組むため、スライスヘッダの値から
5// POC を求める必要がある。ハードウェアにもグローバル状態にも依存しないので単体試験で固定する。
6//
7// 対応: I/P/B、POC type 0・2、CAVLC/CABAC(cabac_init_idc を読む)、短期参照の
8// 参照リスト並べ替え(ref_pic_list_modification)と MMCO 1(短期参照を外す)、
9// 明示的な重み付き予測の表(pred_weight_table)。
10// 未対応(明示的にエラー): SP/SI、POC type 1、長期参照(並べ替え・MMCO 2〜6)、
11// フィールド(frame_mbs_only_flag = 0)。
12
13use super::h264::{BitReader, Pps, Sps};
14use alloc::vec::Vec;
15
16/// スライス種別(slice_type % 5)。
17#[derive(Clone, Copy, PartialEq, Eq, Debug)]
18pub enum SliceKind {
19    P,
20    B,
21    I,
22}
23
24/// 解析済みのスライスヘッダ(7.3.3)。
25#[derive(Clone, Debug)]
26pub struct SliceHeader {
27    pub first_mb: u32,
28    pub kind: SliceKind,
29    pub frame_num: u32,
30    pub pic_order_cnt_lsb: u32,
31    pub delta_pic_order_cnt_bottom: i32,
32    pub direct_spatial_mv_pred: bool,
33    /// num_ref_idx_l0/l1_active(上書き反映後)。
34    pub num_ref_idx_active: [u32; 2],
35    pub cabac_init_idc: u32,
36    pub slice_qp: i32,
37    pub disable_deblocking_filter_idc: u32,
38    pub slice_alpha_c0_offset: i32,
39    pub slice_beta_offset: i32,
40    /// ref_pic_list_modification(L0/L1)。(modification_of_pic_nums_idc, abs_diff_pic_num_minus1)、
41    /// idc は 0(引く)か 1(足す)だけ。
42    pub list_mod: [Vec<(u32, u32)>; 2],
43    /// adaptive_ref_pic_marking_mode_flag。真ならスライディングウィンドウを使わず `mmco` に従う。
44    pub adaptive_marking: bool,
45    /// MMCO 1 の difference_of_pic_nums_minus1 の並び(その PicNum の短期参照を外す)。
46    pub mmco_unmark: Vec<u32>,
47    /// 明示的な重み付き予測の表(P で weighted_pred_flag、B で weighted_bipred_idc = 1 のとき)。
48    pub pwt: Option<PredWeightTable>,
49}
50
51/// pred_weight_table(7.3.3.2)。`w[list][ref_idx][成分]` = (重み, オフセット)、成分は Y/Cb/Cr。
52/// フラグの立っていない参照は既定値(重み 2^denom、オフセット 0)で埋める。4:2:0 前提。
53#[derive(Clone, Debug)]
54pub struct PredWeightTable {
55    pub luma_log2_denom: u32,
56    pub chroma_log2_denom: u32,
57    pub w: [Vec<[(i32, i32); 3]>; 2],
58}
59
60fn parse_pred_weight_table(r: &mut BitReader, kind: SliceKind, num_active: [u32; 2]) -> Result<PredWeightTable, &'static str> {
61    let luma_log2_denom = r.read_ue();
62    let chroma_log2_denom = r.read_ue();
63    if luma_log2_denom > 7 || chroma_log2_denom > 7 {
64        return Err("bad log2_weight_denom");
65    }
66    let lists = if kind == SliceKind::B { 2 } else { 1 };
67    let mut w = [Vec::new(), Vec::new()];
68    for (list, wl) in w.iter_mut().enumerate().take(lists) {
69        for _ in 0..num_active[list] {
70            let mut e = [(1i32 << luma_log2_denom, 0i32), (1 << chroma_log2_denom, 0), (1 << chroma_log2_denom, 0)];
71            if r.read_bit() == 1 {
72                e[0] = (r.read_se(), r.read_se());
73            }
74            if r.read_bit() == 1 {
75                e[1] = (r.read_se(), r.read_se());
76                e[2] = (r.read_se(), r.read_se());
77            }
78            wl.push(e);
79        }
80    }
81    Ok(PredWeightTable { luma_log2_denom, chroma_log2_denom, w })
82}
83
84/// ref_pic_list_modification の 1 リスト分(7.3.3.1)。短期参照(idc 0/1)だけ受け付ける。
85fn parse_list_mod(r: &mut BitReader) -> Result<Vec<(u32, u32)>, &'static str> {
86    let mut v = Vec::new();
87    if r.read_bit() == 0 {
88        return Ok(v);
89    }
90    loop {
91        let idc = r.read_ue();
92        match idc {
93            0 | 1 => v.push((idc, r.read_ue())),
94            3 => return Ok(v),
95            _ => return Err("long-term ref_pic_list_modification not supported"),
96        }
97        if v.len() > 32 {
98            return Err("too many ref_pic_list_modification entries");
99        }
100    }
101}
102
103/// スライスヘッダを読む。`r` はスライス RBSP の先頭から。戻った時点で `r` は slice_data の先頭。
104pub fn parse_slice_header(
105    r: &mut BitReader,
106    sps: &Sps,
107    pps: &Pps,
108    is_idr: bool,
109    nal_ref_idc: u8,
110) -> Result<SliceHeader, &'static str> {
111    if sps.frame_mbs_only_flag == 0 {
112        return Err("field coding not supported");
113    }
114    let first_mb = r.read_ue();
115    let kind = match r.read_ue() % 5 {
116        0 => SliceKind::P,
117        1 => SliceKind::B,
118        2 => SliceKind::I,
119        _ => return Err("SP/SI slices not supported"),
120    };
121    let _pps_id = r.read_ue();
122    let frame_num = r.read_bits(sps.log2_max_frame_num);
123    if is_idr {
124        let _idr_pic_id = r.read_ue();
125    }
126    let mut pic_order_cnt_lsb = 0;
127    let mut delta_pic_order_cnt_bottom = 0;
128    match sps.pic_order_cnt_type {
129        0 => {
130            pic_order_cnt_lsb = r.read_bits(sps.log2_max_pic_order_cnt_lsb);
131            if pps.bottom_field_pic_order_in_frame_present_flag != 0 {
132                delta_pic_order_cnt_bottom = r.read_se();
133            }
134        }
135        2 => {}
136        _ => return Err("pic_order_cnt_type 1 not supported"),
137    }
138    if pps.redundant_pic_cnt_present_flag != 0 {
139        let _redundant_pic_cnt = r.read_ue();
140    }
141    let mut direct_spatial_mv_pred = false;
142    if kind == SliceKind::B {
143        direct_spatial_mv_pred = r.read_bit() == 1;
144    }
145    let mut num_ref_idx_active = [
146        pps.num_ref_idx_l0_default_active.max(1),
147        pps.num_ref_idx_l1_default_active.max(1),
148    ];
149    if kind != SliceKind::I && r.read_bit() == 1 {
150        num_ref_idx_active[0] = r.read_ue().min(31) + 1;
151        if kind == SliceKind::B {
152            num_ref_idx_active[1] = r.read_ue().min(31) + 1;
153        }
154    }
155    let mut list_mod = [Vec::new(), Vec::new()];
156    if kind != SliceKind::I {
157        list_mod[0] = parse_list_mod(r)?;
158    }
159    if kind == SliceKind::B {
160        list_mod[1] = parse_list_mod(r)?;
161    }
162    let mut pwt = None;
163    if (kind == SliceKind::P && pps.weighted_pred_flag != 0)
164        || (kind == SliceKind::B && pps.weighted_bipred_idc == 1)
165    {
166        pwt = Some(parse_pred_weight_table(r, kind, num_ref_idx_active)?);
167    }
168    // dec_ref_pic_marking(nal_ref_idc != 0 のときだけ)。長期参照は未対応。
169    let mut adaptive_marking = false;
170    let mut mmco_unmark = Vec::new();
171    if nal_ref_idc != 0 {
172        if is_idr {
173            let _no_output_of_prior_pics = r.read_bit();
174            if r.read_bit() == 1 {
175                return Err("long-term reference not supported");
176            }
177        } else if r.read_bit() == 1 {
178            adaptive_marking = true;
179            loop {
180                match r.read_ue() {
181                    0 => break,
182                    1 => mmco_unmark.push(r.read_ue()),
183                    _ => return Err("MMCO other than 1 not supported"),
184                }
185                if mmco_unmark.len() > 32 {
186                    return Err("too many MMCO entries");
187                }
188            }
189        }
190    }
191    let mut cabac_init_idc = 0;
192    if pps.entropy_coding_mode_flag != 0 && kind != SliceKind::I {
193        cabac_init_idc = r.read_ue();
194        if cabac_init_idc > 2 {
195            return Err("bad cabac_init_idc");
196        }
197    }
198    let slice_qp = pps.pic_init_qp + r.read_se();
199    let mut disable_deblocking_filter_idc = 0;
200    let mut slice_alpha_c0_offset = 0;
201    let mut slice_beta_offset = 0;
202    if pps.deblocking_filter_control_present_flag != 0 {
203        disable_deblocking_filter_idc = r.read_ue();
204        if disable_deblocking_filter_idc != 1 {
205            slice_alpha_c0_offset = r.read_se() * 2;
206            slice_beta_offset = r.read_se() * 2;
207        }
208    }
209    Ok(SliceHeader {
210        first_mb,
211        kind,
212        frame_num,
213        pic_order_cnt_lsb,
214        delta_pic_order_cnt_bottom,
215        direct_spatial_mv_pred,
216        num_ref_idx_active,
217        cabac_init_idc,
218        slice_qp,
219        disable_deblocking_filter_idc,
220        slice_alpha_c0_offset,
221        slice_beta_offset,
222        list_mod,
223        adaptive_marking,
224        mmco_unmark,
225        pwt,
226    })
227}
228
229/// POC の計算状態(8.2.1)。フレームの復号のたびに `compute` を呼ぶ。
230#[derive(Default, Clone)]
231pub struct PocState {
232    prev_poc_msb: i32,
233    prev_poc_lsb: i32,
234    prev_frame_num: u32,
235    prev_frame_num_offset: i32,
236}
237
238impl PocState {
239    pub fn new() -> Self {
240        Self::default()
241    }
242
243    /// 現在のフレームの POC(フレームなので TopFieldOrderCnt と BottomFieldOrderCnt の小さい方)。
244    pub fn compute(&mut self, sps: &Sps, h: &SliceHeader, is_idr: bool, nal_ref_idc: u8) -> i32 {
245        match sps.pic_order_cnt_type {
246            0 => {
247                // 8.2.1.1
248                if is_idr {
249                    self.prev_poc_msb = 0;
250                    self.prev_poc_lsb = 0;
251                }
252                let max_lsb = 1i32 << sps.log2_max_pic_order_cnt_lsb;
253                let lsb = h.pic_order_cnt_lsb as i32;
254                let msb = if lsb < self.prev_poc_lsb && self.prev_poc_lsb - lsb >= max_lsb / 2 {
255                    self.prev_poc_msb + max_lsb
256                } else if lsb > self.prev_poc_lsb && lsb - self.prev_poc_lsb > max_lsb / 2 {
257                    self.prev_poc_msb - max_lsb
258                } else {
259                    self.prev_poc_msb
260                };
261                let top = msb + lsb;
262                let bottom = top + h.delta_pic_order_cnt_bottom;
263                // 参照ピクチャだけが次の prevPicOrderCnt になる(MMCO5 は未対応)。
264                if nal_ref_idc != 0 {
265                    self.prev_poc_msb = msb;
266                    self.prev_poc_lsb = lsb;
267                }
268                top.min(bottom)
269            }
270            _ => {
271                // 8.2.1.3(type 2): 表示順=復号順。
272                let max_frame_num = 1i32 << sps.log2_max_frame_num;
273                let offset = if is_idr {
274                    0
275                } else if self.prev_frame_num > h.frame_num {
276                    self.prev_frame_num_offset + max_frame_num
277                } else {
278                    self.prev_frame_num_offset
279                };
280                let poc = if is_idr {
281                    0
282                } else if nal_ref_idc == 0 {
283                    2 * (offset + h.frame_num as i32) - 1
284                } else {
285                    2 * (offset + h.frame_num as i32)
286                };
287                self.prev_frame_num = h.frame_num;
288                self.prev_frame_num_offset = offset;
289                poc
290            }
291        }
292    }
293}