oarust 0.8.6

GPU-first semantic computing engine written in Rust
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

use crate::{Error, Result};

pub(super) fn remove_emulation_prevention(data: &[u8]) -> Vec<u8> {
	let mut output = Vec::with_capacity(data.len());
	let mut zeros = 0;
	for &byte in data {
		if zeros >= 2 && byte == 3 {
			zeros = 0;
			continue;
		}
		output.push(byte);
		zeros = if byte == 0 { zeros + 1 } else { 0 };
	}
	output
}

pub(super) struct BitReader<'a> {
	data: &'a [u8],
	bit: usize,
}

impl<'a> BitReader<'a> {
	pub(super) const fn new(data: &'a [u8]) -> Self {
		Self { data, bit: 0 }
	}

	pub(super) const fn position(&self) -> usize {
		self.bit
	}

	pub(super) fn read_bits(&mut self, count: usize) -> Result<u32> {
		if count > 32
			|| self
				.bit
				.checked_add(count)
				.is_none_or(|end| end > self.data.len() * 8)
		{
			return Err(Error::data_loss("truncated codec bitstream"));
		}
		let mut value = 0_u32;
		for _ in 0..count {
			value = (value << 1) | u32::from((self.data[self.bit / 8] >> (7 - self.bit % 8)) & 1);
			self.bit += 1;
		}
		Ok(value)
	}

	pub(super) fn skip_bits(&mut self, mut count: usize) -> Result<()> {
		while count != 0 {
			let chunk = count.min(32);
			self.read_bits(chunk)?;
			count -= chunk;
		}
		Ok(())
	}

	pub(super) fn read_ue(&mut self) -> Result<u32> {
		let mut leading_zeros = 0_usize;
		while self.read_bits(1)? == 0 {
			leading_zeros += 1;
			if leading_zeros > 31 {
				return Err(Error::data_loss("Exp-Golomb value exceeds u32"));
			}
		}
		let suffix = self.read_bits(leading_zeros)?;
		Ok(((1_u32 << leading_zeros) - 1) + suffix)
	}

	pub(super) fn read_se(&mut self) -> Result<i32> {
		let code = self.read_ue()?;
		let magnitude = i32::try_from(code.div_ceil(2))
			.map_err(|_| Error::data_loss("signed Exp-Golomb value exceeds i32"))?;
		Ok(if code.is_multiple_of(2) {
			-magnitude
		} else {
			magnitude
		})
	}

	pub(super) fn more_rbsp_data(&self) -> bool {
		let Some((byte_index, byte)) = self
			.data
			.iter()
			.enumerate()
			.rev()
			.find(|(_, byte)| **byte != 0)
		else {
			return false;
		};
		let stop_bit = byte_index * 8 + (u8::BITS - 1 - byte.trailing_zeros()) as usize;
		self.bit < stop_bit
	}
}