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 if code_num == 0 {
34 Ok(0)
35 } else if code_num & 1 == 1 {
36 Ok(((code_num + 1) >> 1) as i32)
38 } else {
39 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 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 let data = vec![0b10100000];
144 let mut reader = BitReader::new(&data);
145 assert_eq!(read_se(&mut reader).unwrap(), 0);
146
147 let data = vec![0b01010000];
149 let mut reader = BitReader::new(&data);
150 assert_eq!(read_se(&mut reader).unwrap(), 1);
151
152 let data = vec![0b01100000];
154 let mut reader = BitReader::new(&data);
155 assert_eq!(read_se(&mut reader).unwrap(), -1);
156
157 let data = vec![0b00100000];
159 let mut reader = BitReader::new(&data);
160 assert_eq!(read_se(&mut reader).unwrap(), 2);
161
162 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}