h264_parser/
eg.rs

1use crate::bitreader::BitReader;
2use crate::{Error, Result};
3
4pub fn read_ue(reader: &mut BitReader) -> Result<u32> {
5    let mut leading_zeros = 0;
6
7    while !reader.read_bit()? {
8        leading_zeros += 1;
9        if leading_zeros > 31 {
10            return Err(Error::BitstreamError("Invalid exp-golomb code".into()));
11        }
12    }
13
14    if leading_zeros == 0 {
15        return Ok(0);
16    }
17
18    let code_value = reader.read_bits(leading_zeros)?;
19    Ok((1 << leading_zeros) - 1 + code_value)
20}
21
22pub fn read_se(reader: &mut BitReader) -> Result<i32> {
23    let code_num = read_ue(reader)?;
24
25    // H.264 spec mapping:
26    // code_num = 0 => 0
27    // code_num = 1 => 1
28    // code_num = 2 => -1
29    // code_num = 3 => 2
30    // code_num = 4 => -2
31    // Pattern: odd values are positive, even values (except 0) are negative
32
33    if code_num == 0 {
34        Ok(0)
35    } else if code_num & 1 == 1 {
36        // Odd: positive value
37        Ok(((code_num + 1) >> 1) as i32)
38    } else {
39        // Even: negative value
40        Ok(-((code_num >> 1) as i32))
41    }
42}
43
44pub fn read_me(reader: &mut BitReader, chroma_format_idc: u8) -> Result<u32> {
45    match chroma_format_idc {
46        1 | 2 => {
47            let code_num = read_ue(reader)?;
48            if code_num > 2 {
49                return Err(Error::BitstreamError(
50                    "Invalid mapped exp-golomb code".into(),
51                ));
52            }
53            Ok(code_num)
54        }
55        _ => read_ue(reader),
56    }
57}
58
59pub fn read_te(reader: &mut BitReader, max_value: u32) -> Result<u32> {
60    if max_value == 0 {
61        return Ok(0);
62    }
63
64    if max_value == 1 {
65        let bit = reader.read_bit()?;
66        return Ok(1 - bit as u32);
67    }
68
69    read_ue(reader)
70}
71
72pub fn write_ue(value: u32) -> Vec<bool> {
73    if value == 0 {
74        return vec![true];
75    }
76
77    let code_num = value + 1;
78    let num_bits = 32 - code_num.leading_zeros();
79    let total_bits = 2 * num_bits - 1;
80
81    let mut bits = Vec::with_capacity(total_bits as usize);
82
83    for _ in 0..(num_bits - 1) {
84        bits.push(false);
85    }
86
87    for i in (0..num_bits).rev() {
88        bits.push((code_num >> i) & 1 != 0);
89    }
90
91    bits
92}
93
94pub fn write_se(value: i32) -> Vec<bool> {
95    // Inverse of read_se mapping:
96    // 0 => code_num = 0
97    // 1 => code_num = 1
98    // -1 => code_num = 2
99    // 2 => code_num = 3
100    // -2 => code_num = 4
101
102    let code_num = if value == 0 {
103        0
104    } else if value > 0 {
105        (value as u32) * 2 - 1
106    } else {
107        ((-value) as u32) * 2
108    };
109
110    write_ue(code_num)
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn test_read_ue() {
119        let data = vec![0b10100000];
120        let mut reader = BitReader::new(&data);
121        assert_eq!(read_ue(&mut reader).unwrap(), 0);
122
123        let data = vec![0b01010000];
124        let mut reader = BitReader::new(&data);
125        assert_eq!(read_ue(&mut reader).unwrap(), 1);
126
127        let data = vec![0b01100000];
128        let mut reader = BitReader::new(&data);
129        assert_eq!(read_ue(&mut reader).unwrap(), 2);
130
131        let data = vec![0b00101100];
132        let mut reader = BitReader::new(&data);
133        assert_eq!(read_ue(&mut reader).unwrap(), 4);
134
135        let data = vec![0b00011110];
136        let mut reader = BitReader::new(&data);
137        assert_eq!(read_ue(&mut reader).unwrap(), 14);
138    }
139
140    #[test]
141    fn test_read_se() {
142        // SE(0) = UE(0) = 1 => 0
143        let data = vec![0b10100000];
144        let mut reader = BitReader::new(&data);
145        assert_eq!(read_se(&mut reader).unwrap(), 0);
146
147        // SE(1) = UE(1) = 010 => 1
148        let data = vec![0b01010000];
149        let mut reader = BitReader::new(&data);
150        assert_eq!(read_se(&mut reader).unwrap(), 1);
151
152        // SE(-1) = UE(2) = 011 => -1
153        let data = vec![0b01100000];
154        let mut reader = BitReader::new(&data);
155        assert_eq!(read_se(&mut reader).unwrap(), -1);
156
157        // SE(2) = UE(3) = 00100 => 2
158        let data = vec![0b00100000];
159        let mut reader = BitReader::new(&data);
160        assert_eq!(read_se(&mut reader).unwrap(), 2);
161
162        // SE(-2) = UE(4) = 00101 => -2
163        let data = vec![0b00101000];
164        let mut reader = BitReader::new(&data);
165        assert_eq!(read_se(&mut reader).unwrap(), -2);
166    }
167
168    #[test]
169    fn test_write_ue() {
170        assert_eq!(write_ue(0), vec![true]);
171        assert_eq!(write_ue(1), vec![false, true, false]);
172        assert_eq!(write_ue(2), vec![false, true, true]);
173        assert_eq!(write_ue(3), vec![false, false, true, false, false]);
174    }
175
176    #[test]
177    fn test_write_se() {
178        assert_eq!(write_se(0), vec![true]);
179        assert_eq!(write_se(1), vec![false, true, false]);
180        assert_eq!(write_se(-1), vec![false, true, true]);
181        assert_eq!(write_se(2), vec![false, false, true, false, false]);
182        assert_eq!(write_se(-2), vec![false, false, true, false, true]);
183    }
184
185    #[test]
186    fn test_read_te() {
187        let data = vec![0b10000000];
188        let mut reader = BitReader::new(&data);
189        assert_eq!(read_te(&mut reader, 0).unwrap(), 0);
190
191        let data = vec![0b00000000];
192        let mut reader = BitReader::new(&data);
193        assert_eq!(read_te(&mut reader, 1).unwrap(), 1);
194
195        let data = vec![0b10000000];
196        let mut reader = BitReader::new(&data);
197        assert_eq!(read_te(&mut reader, 1).unwrap(), 0);
198    }
199}