diff --git a/crates/ironrdp-graphics/src/rlgr.rs b/crates/ironrdp-graphics/src/rlgr.rs index c3b734d9ed..7980701b8f 100644 --- a/crates/ironrdp-graphics/src/rlgr.rs +++ b/crates/ironrdp-graphics/src/rlgr.rs @@ -6,8 +6,6 @@ use bitvec::prelude::*; use ironrdp_pdu::codecs::rfx::EntropyAlgorithm; use yuv::YuvError; -use crate::utils::Bits; - const KP_MAX: u32 = 80; const LS_GR: u32 = 3; const UP_GR: u32 = 4; @@ -26,12 +24,11 @@ macro_rules! write_byte { }; } -macro_rules! try_split_bits { +macro_rules! try_read_bits { ($bits:ident, $n:expr) => { - if $bits.len() < $n { - break; - } else { - $bits.split_to($n) + match $bits.read($n) { + Some(value) => value, + None => break, } }; } @@ -268,21 +265,21 @@ pub fn decode(mode: EntropyAlgorithm, tile: &[u8], mut output: &mut [i16]) -> Re let mut kp: u32 = k << LS_GR; let mut krp: u32 = kr << LS_GR; - let mut bits = Bits::new(BitSlice::from_slice(tile)); + let mut bits = BitReader::new(tile); - while !bits.is_empty() && !output.is_empty() { + while bits.remaining() != 0 && !output.is_empty() { match CompressionMode::from(k) { CompressionMode::RunLength => { - let number_of_zeros = truncate_leading_value(&mut bits, false); - try_split_bits!(bits, 1); - let run = count_run(number_of_zeros, &mut k, &mut kp) + load_be_u32(try_split_bits!(bits, k as usize)); + let number_of_zeros = bits.skip_run(false); + try_read_bits!(bits, 1); + let run = count_run(number_of_zeros, &mut k, &mut kp) + try_read_bits!(bits, k as usize); - let sign_bit = try_split_bits!(bits, 1).load_be::(); + let sign_bit = try_read_bits!(bits, 1) as u8; - let number_of_ones = truncate_leading_value(&mut bits, true); - try_split_bits!(bits, 1); + let number_of_ones = bits.skip_run(true); + try_read_bits!(bits, 1); - let code_remainder = load_be_u32(try_split_bits!(bits, kr as usize)) + ((number_of_ones as u32) << kr); + let code_remainder = try_read_bits!(bits, kr as usize) + ((number_of_ones as u32) << kr); update_parameters_according_to_number_of_ones(number_of_ones, &mut kr, &mut krp); kp = kp.saturating_sub(DN_GR); @@ -296,10 +293,10 @@ pub fn decode(mode: EntropyAlgorithm, tile: &[u8], mut output: &mut [i16]) -> Re write_byte!(output, magnitude); } CompressionMode::GolombRice => { - let number_of_ones = truncate_leading_value(&mut bits, true); - try_split_bits!(bits, 1); + let number_of_ones = bits.skip_run(true); + try_read_bits!(bits, 1); - let code_remainder = load_be_u32(try_split_bits!(bits, kr as usize)) + ((number_of_ones as u32) << kr); + let code_remainder = try_read_bits!(bits, kr as usize) + ((number_of_ones as u32) << kr); update_parameters_according_to_number_of_ones(number_of_ones, &mut kr, &mut krp); @@ -311,7 +308,7 @@ pub fn decode(mode: EntropyAlgorithm, tile: &[u8], mut output: &mut [i16]) -> Re EntropyAlgorithm::Rlgr3 => { let n_index = compute_n_index(code_remainder); - let val1 = load_be_u32(try_split_bits!(bits, n_index)); + let val1 = try_read_bits!(bits, n_index); let val2 = code_remainder - val1; if val1 != 0 && val2 != 0 { kp = kp.saturating_sub(2 * DQ_GR); @@ -344,19 +341,87 @@ fn fill(buffer: &mut [i16], value: i16) { } } -fn load_be_u32(s: &BitSlice) -> u32 { - if s.is_empty() { 0 } else { s.load_be::() } +/// Reads a tile's bits most significant first, through a 64-bit window +/// instead of a bit slice: RLGR decoding is most of a progressive tile's cost. +struct BitReader<'a> { + data: &'a [u8], + /// Position in bits. + pos: usize, + /// Length in bits. + len: usize, } -// Returns number of truncated bits -fn truncate_leading_value(bits: &mut Bits<'_>, value: bool) -> usize { - let leading_values = if value { - bits.leading_ones() - } else { - bits.leading_zeros() - }; - bits.split_to(leading_values); - leading_values +impl<'a> BitReader<'a> { + fn new(data: &'a [u8]) -> Self { + Self { + data, + pos: 0, + len: data.len() * 8, + } + } + + fn remaining(&self) -> usize { + self.len - self.pos + } + + /// The next 64 bits, left-aligned, zero past the end. + fn peek64(&self) -> u64 { + let byte = self.pos / 8; + let shift = self.pos % 8; + let mut window = match self.data.get(byte..byte + 8) { + // A plain load; copying into a buffer calls memmove. + Some(bytes) => u64::from_be_bytes(bytes.try_into().expect("8 bytes")), + None => { + let mut buf = [0u8; 8]; + let tail = &self.data[byte..]; + buf[..tail.len()].copy_from_slice(tail); + u64::from_be_bytes(buf) + } + } << shift; + if shift != 0 && byte + 8 < self.data.len() { + window |= u64::from(self.data[byte + 8]) >> (8 - shift); + } + window + } + + /// Reads `n` bits (at most 32) as a big-endian number, or `None` when + /// fewer than `n` remain. Zero bits read as 0. + fn read(&mut self, n: usize) -> Option { + #![expect(clippy::as_conversions, clippy::cast_possible_truncation, reason = "at most 32 bits")] + + debug_assert!(n <= 32, "RLGR reads at most 32 bits"); + if n > self.remaining() { + return None; + } + if n == 0 { + return Some(0); + } + let value = (self.peek64() >> (64 - n)) as u32; + self.pos += n; + Some(value) + } + + /// Consumes the run of `value` bits at the current position, up to the + /// end of the tile, and returns its length. + fn skip_run(&mut self, value: bool) -> usize { + #![expect(clippy::as_conversions, reason = "u32-to-usize")] + + let start = self.pos; + while self.pos < self.len { + let window = self.peek64(); + let run = if value { + window.leading_ones() + } else { + window.leading_zeros() + } as usize; + let run = min(run, self.remaining()); + self.pos += run; + if run < 64 { + break; + } + } + self.pos - start + } } fn count_run(number_of_zeros: usize, k: &mut u32, kp: &mut u32) -> u32 { @@ -409,16 +474,13 @@ fn compute_rlgr3_magnitude(val: u32) -> Result { } } +#[expect(clippy::as_conversions, reason = "u32-to-usize")] fn compute_n_index(code_remainder: u32) -> usize { if code_remainder == 0 { return 0; } - let code_bytes = code_remainder.to_be_bytes(); - let code_bits = BitSlice::::from_slice(code_bytes.as_ref()); - let leading_zeros = code_bits.leading_zeros(); - - 32 - leading_zeros + (32 - code_remainder.leading_zeros()) as usize } fn update_parameters_according_to_number_of_ones(number_of_ones: usize, kr: &mut u32, krp: &mut u32) { @@ -496,3 +558,60 @@ impl From for RlgrError { Self::Io(err) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bit_reader_reads_across_byte_boundaries() { + let mut bits = BitReader::new(&[0b1011_0011, 0b1100_0101, 0xFF, 0x00, 0xAA, 0x55, 0x0F, 0xF0, 0x81]); + + assert_eq!(bits.read(3), Some(0b101)); + assert_eq!(bits.read(0), Some(0)); + assert_eq!(bits.read(7), Some(0b100_1111)); + assert_eq!(bits.read(10), Some(0b00_0101_1111)); + // Straddles the end of the first 64-bit window. + assert_eq!(bits.read(32), Some(0b1111_0000_0000_1010_1010_0101_0101_0000)); + assert_eq!(bits.remaining(), 20); + assert_eq!(bits.read(20), Some(0b1111_1111_0000_1000_0001)); + assert_eq!(bits.remaining(), 0); + } + + #[test] + fn bit_reader_refuses_reads_past_the_end() { + let mut bits = BitReader::new(&[0xAB, 0xCD]); + + assert_eq!(bits.read(5), Some(0b10101)); + assert_eq!(bits.read(12), None); + // A refused read consumes nothing. + assert_eq!(bits.remaining(), 11); + assert_eq!(bits.read(11), Some(0b011_1100_1101)); + assert_eq!(bits.read(1), None); + assert_eq!(bits.read(0), Some(0)); + } + + #[test] + fn bit_reader_zero_run_stops_at_the_end() { + let mut bits = BitReader::new(&[0xFF, 0x00, 0x00]); + + assert_eq!(bits.read(4), Some(0xF)); + assert_eq!(bits.skip_run(true), 4); + assert_eq!(bits.skip_run(false), 16); + assert_eq!(bits.remaining(), 0); + assert_eq!(bits.skip_run(false), 0); + } + + #[test] + fn bit_reader_one_run_stops_at_the_end() { + // A run longer than one 64-bit window. + let mut data = vec![0x0F]; + data.extend([0xFF; 11]); + let mut bits = BitReader::new(&data); + + assert_eq!(bits.skip_run(false), 4); + assert_eq!(bits.skip_run(true), 92); + assert_eq!(bits.remaining(), 0); + assert_eq!(bits.read(1), None); + } +}