From fd37e74b12f2446ecfe9f2b59242099d911072df Mon Sep 17 00:00:00 2001 From: Greg Lamberson Date: Fri, 25 Sep 2026 05:20:37 -0500 Subject: [PATCH] fix(graphics): keep ZGFX compression linear once the history is full Once the 2.5 MB history filled, add_to_history drained the front of a Vec and rebased every match-table position on each literal and match, making every token cost O(history + table). Keep the history in a fixed ring addressed by absolute stream position, store absolute positions in the match table and skip candidates beyond the reachable window at lookup, so appends never move or rebase anything. The encoded format is unchanged. --- .../src/zgfx/circular_buffer.rs | 9 + .../ironrdp-graphics/src/zgfx/compressor.rs | 288 ++++++++++++++---- 2 files changed, 233 insertions(+), 64 deletions(-) diff --git a/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs b/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs index 56b503a756..68ca8ce0ce 100644 --- a/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs +++ b/crates/ironrdp-graphics/src/zgfx/circular_buffer.rs @@ -59,6 +59,15 @@ impl FixedCircularBuffer { Ok(()) } + + /// The byte `distance` positions before the write position. + /// + /// INVARIANT: `0 < distance <= self.buffer.len()`. The callers are the compressor's history, + /// which only asks for positions it computed itself, so this is checked in debug builds only. + pub(crate) fn byte_back(&self, distance: usize) -> u8 { + debug_assert!(0 < distance && distance <= self.buffer.len()); + self.buffer[(self.position + self.buffer.len() - distance) % self.buffer.len()] + } } impl io::Write for FixedCircularBuffer { diff --git a/crates/ironrdp-graphics/src/zgfx/compressor.rs b/crates/ironrdp-graphics/src/zgfx/compressor.rs index f9c11585d0..309d5eb930 100644 --- a/crates/ironrdp-graphics/src/zgfx/compressor.rs +++ b/crates/ironrdp-graphics/src/zgfx/compressor.rs @@ -7,11 +7,11 @@ //! //! [\[MS-RDPEGFX\] 2.2.1.1.1]: https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-rdpegfx/ -use std::collections::HashMap; +use std::{collections::HashMap, io::Write as _}; use bitvec::prelude::*; -use super::{HISTORY_SIZE, TOKEN_TABLE, ZgfxError}; +use super::{FixedCircularBuffer, HISTORY_SIZE, TOKEN_TABLE, ZgfxError}; const MIN_MATCH_LENGTH: usize = 3; const MAX_MATCH_LENGTH: usize = 65535; /// Maximum back-reference distance (last token in MS-RDPEGFX table) @@ -32,15 +32,17 @@ const COMPACT_TARGET_ENTRIES: usize = MAX_HASH_TABLE_ENTRIES / 2; /// ZGFX compressor maintaining a 2.5 MB history buffer and prefix hash table. pub struct Compressor { - history: Vec, - /// 3-byte prefix → history positions, for O(1) match candidate lookup - match_table: HashMap<[u8; 3], Vec>, + history: History, + /// 3-byte prefix → absolute stream positions, for O(1) match candidate lookup. + /// Positions are never rebased; candidates that have fallen out of reach are + /// skipped at lookup time. + match_table: HashMap<[u8; 3], Vec>, } impl Compressor { pub fn new() -> Self { Self { - history: Vec::with_capacity(HISTORY_SIZE), + history: History::new(), match_table: HashMap::new(), } } @@ -71,47 +73,17 @@ impl Compressor { Ok(bit_writer.finish()) } - /// Extend the sliding window, evicting oldest bytes when full. + /// Extend the sliding window; the ring overwrites the oldest bytes once full. fn add_to_history(&mut self, bytes: &[u8]) { - if self.history.len() + bytes.len() > HISTORY_SIZE { - let overflow = (self.history.len() + bytes.len()) - HISTORY_SIZE; - - self.history.drain(..overflow); - - // Shift all stored positions to account for evicted bytes - for positions in self.match_table.values_mut() { - positions.retain_mut(|pos| { - if *pos >= overflow { - *pos -= overflow; - true - } else { - false - } - }); - } - self.match_table.retain(|_, positions| !positions.is_empty()); - } - - let base_pos = self.history.len(); - self.history.extend_from_slice(bytes); + let base_pos = self.history.total(); + self.history.push(bytes); // For large chunks (match replays), sample every 4th position to keep // the hash table manageable without sacrificing much compression ratio let step_size = if bytes.len() > 256 { 4 } else { 1 }; for i in (0..bytes.len().saturating_sub(MIN_MATCH_LENGTH - 1)).step_by(step_size) { - let pos = base_pos + i; - let prefix = [self.history[pos], self.history[pos + 1], self.history[pos + 2]]; - - let entry = self.match_table.entry(prefix).or_default(); - - if entry.len() < MAX_POSITIONS_PER_PREFIX { - entry.push(pos); - } else { - // Evict oldest position to keep recent data preferred - entry.remove(0); - entry.push(pos); - } + self.record_position([bytes[i], bytes[i + 1], bytes[i + 2]], base_pos + as_u64(i)); } if self.match_table.len() > MAX_HASH_TABLE_ENTRIES { @@ -120,23 +92,34 @@ impl Compressor { // Index positions spanning the old/new data boundary so matches // that straddle the append point can be found - for offset in [2, 1] { - if base_pos >= offset && bytes.len() + offset > 2 { - let pos = base_pos - offset; - if pos + MIN_MATCH_LENGTH <= self.history.len() { - let prefix = [self.history[pos], self.history[pos + 1], self.history[pos + 2]]; - let entry = self.match_table.entry(prefix).or_default(); - if entry.last() != Some(&pos) { - if entry.len() >= MAX_POSITIONS_PER_PREFIX { - entry.remove(0); - } - entry.push(pos); - } - } - } + for offset in [2u64, 1] { + let Some(pos) = base_pos.checked_sub(offset) else { + continue; + }; + let Some(prefix) = self.history.prefix_at(pos) else { + continue; + }; + self.record_position(prefix, pos); } } + /// Remember that `prefix` occurs at the absolute stream position `pos`, dropping the + /// prefix's oldest position when it already has `MAX_POSITIONS_PER_PREFIX`. + /// + /// INVARIANT: each position is recorded once, when its third byte arrives. The main + /// loop covers the positions whose three bytes are all in the new chunk, and the + /// boundary loop the two positions that straddle the append point, so no position + /// is in both. + fn record_position(&mut self, prefix: [u8; 3], pos: u64) { + let positions = self.match_table.entry(prefix).or_default(); + debug_assert!(positions.last() != Some(&pos)); + if MAX_POSITIONS_PER_PREFIX <= positions.len() { + // Evict oldest position to keep recent data preferred + positions.remove(0); + } + positions.push(pos); + } + /// Halve stored positions per prefix, then evict whole prefixes down to /// `COMPACT_TARGET_ENTRIES` when the table is over the cap. /// @@ -164,7 +147,7 @@ impl Compressor { // MAX_MATCH_DISTANCE regardless. Each history position belongs to a // single prefix, so these newest positions are distinct and the cutoff // retains exactly COMPACT_TARGET_ENTRIES entries. - let mut newest: Vec = self + let mut newest: Vec = self .match_table .values() .map(|positions| positions.last().copied().unwrap_or(0)) @@ -178,7 +161,7 @@ impl Compressor { /// Search hash table for the longest match at `input[pos..]`. fn find_best_match(&self, input: &[u8], pos: usize) -> Option { let remaining = input.len() - pos; - if remaining < MIN_MATCH_LENGTH || self.history.is_empty() { + if remaining < MIN_MATCH_LENGTH { return None; } @@ -187,22 +170,26 @@ impl Compressor { let max_match_len = remaining.min(MAX_MATCH_LENGTH); let mut best_match: Option = None; - let search_limit = self.history.len().min(MAX_MATCH_DISTANCE); + let search_limit = as_u64(self.history.len().min(MAX_MATCH_DISTANCE)); + let total = self.history.total(); // Most recent candidates first — better locality, often longer matches for &hist_pos in candidates.iter().rev().take(MAX_CANDIDATES) { - let distance = self.history.len() - hist_pos; - - if distance > search_limit { + // Positions are absolute, so anything older than the reachable + // window is simply skipped rather than rebased on every append. + let Some(distance) = total.checked_sub(hist_pos).filter(|d| *d <= search_limit) else { continue; - } + }; + let distance = usize::try_from(distance).expect("distance is at most MAX_MATCH_DISTANCE"); - // Prefix already matched via hash table; extend from byte 3 onward + // Prefix already matched via hash table; extend from byte 3 onward. + // INVARIANT: match_len < distance, so the match reads only bytes + // already in the history. let mut match_len = MIN_MATCH_LENGTH; while match_len < max_match_len - && hist_pos + match_len < self.history.len() - && self.history[hist_pos + match_len] == input[pos + match_len] + && match_len < distance + && self.history.byte_back(distance - match_len) == input[pos + match_len] { match_len += 1; } @@ -339,6 +326,71 @@ impl Default for Compressor { } } +fn as_u64(value: usize) -> u64 { + u64::try_from(value).expect("usize fits in u64") +} + +/// The last `HISTORY_SIZE` bytes of the stream, kept in a fixed ring so that +/// appending never moves existing bytes. +struct History { + ring: FixedCircularBuffer, + /// Bytes appended since the compressor was created; the absolute stream + /// position of the next byte. + total: u64, +} + +impl History { + fn new() -> Self { + Self { + ring: FixedCircularBuffer::new(HISTORY_SIZE), + total: 0, + } + } + + /// Number of valid bytes in the ring, at most `HISTORY_SIZE`. + fn len(&self) -> usize { + usize::try_from(self.total.min(as_u64(HISTORY_SIZE))).expect("at most HISTORY_SIZE") + } + + fn total(&self) -> u64 { + self.total + } + + fn push(&mut self, bytes: &[u8]) { + // The ring keeps only the last HISTORY_SIZE bytes of an input larger than itself. + self.ring.write_all(bytes).expect("writing to the ring cannot fail"); + self.total += as_u64(bytes.len()); + } + + /// The byte `distance` positions before the end of the stream. + /// + /// INVARIANT: `0 < distance <= self.len()`. Both callers establish it + /// themselves, which is why the check is a `debug_assert!` and not a branch + /// on every extended byte: `find_best_match` only keeps candidates within + /// the window and extends while `match_len < distance`, and `prefix_at` + /// checks the window before it reads. Tests run in debug builds, where a + /// violation panics, and the round-trip tests would catch a wrong byte. + fn byte_back(&self, distance: usize) -> u8 { + debug_assert!(0 < distance && distance <= self.len()); + self.ring.byte_back(distance) + } + + /// The three bytes starting at absolute position `pos`, or `None` when they + /// are not all in the ring: fewer than three bytes follow `pos`, or `pos` + /// has fallen out of the window. + fn prefix_at(&self, pos: u64) -> Option<[u8; 3]> { + let distance = usize::try_from(self.total.checked_sub(pos)?).ok()?; + if !(MIN_MATCH_LENGTH..=self.len()).contains(&distance) { + return None; + } + Some([ + self.byte_back(distance), + self.byte_back(distance - 1), + self.byte_back(distance - 2), + ]) + } +} + #[derive(Debug, Clone, Copy)] struct Match { distance: usize, @@ -523,6 +575,114 @@ mod tests { ); } + #[test] + fn compress_round_trips_across_history_wrap() { + use super::super::Decompressor; + + // More than twice HISTORY_SIZE through one compressor, in segments that + // repeat earlier content, so matches reach back across the ring wrap + // point and the decompressor's history has to agree byte for byte. + let mut state: u32 = 0x0BAD_F00D; + let mut next_byte = move || { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + u8::try_from(state >> 24).unwrap() + }; + let blocks: Vec> = + core::iter::repeat_with(|| core::iter::repeat_with(&mut next_byte).take(4096).collect()) + .take(64) + .collect(); + + let mut compressor = Compressor::new(); + let mut decompressor = Decompressor::new(); + let mut total = 0; + let mut segment_index = 0usize; + while total < 2 * HISTORY_SIZE + 100_000 { + let mut segment = Vec::with_capacity(60_000); + for k in 0..14 { + let block = &blocks[(segment_index * 7 + k * 5) % blocks.len()]; + segment.extend_from_slice(block); + segment.push(next_byte()); + } + segment_index += 1; + + let compressed = compressor.compress(&segment).unwrap(); + let mut output = Vec::new(); + decompressor.decompress_segment(&compressed, &mut output).unwrap(); + assert_eq!(output, segment, "segment {segment_index} after {total} bytes"); + total += segment.len(); + } + assert!(compressor.history.total() > u64::try_from(2 * HISTORY_SIZE).unwrap()); + } + + #[test] + fn a_prefix_keeps_only_its_newest_positions() { + let mut compressor = Compressor::new(); + let recorded = MAX_POSITIONS_PER_PREFIX + 8; + for pos in 0..recorded { + compressor.record_position(*b"abc", as_u64(pos)); + } + + let expected: Vec = (8..recorded).map(as_u64).collect(); + assert_eq!(compressor.match_table[b"abc"], expected); + } + + #[test] + fn positions_are_recorded_inside_a_chunk_and_across_the_append_point() { + let mut compressor = Compressor::new(); + compressor.add_to_history(b"abcd"); + compressor.add_to_history(b"ef"); + + // The first chunk holds "abc" at 0 and "bcd" at 1. The second chunk completes "cde" at 2 and + // "def" at 3, which start before the append point. + let table = &compressor.match_table; + assert_eq!(table[b"abc"], [0]); + assert_eq!(table[b"bcd"], [1]); + assert_eq!(table[b"cde"], [2]); + assert_eq!(table[b"def"], [3]); + assert_eq!(table.len(), 4); + } + + #[test] + fn history_ring_keeps_the_last_bytes_in_order() { + let mut history = History::new(); + let chunk: Vec = (0..=255u8).cycle().take(HISTORY_SIZE - 10).collect(); + history.push(&chunk); + history.push(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]); + + assert_eq!(history.len(), HISTORY_SIZE); + assert_eq!(history.byte_back(1), 15); + assert_eq!(history.byte_back(15), 1); + assert_eq!(history.byte_back(16), *chunk.last().unwrap()); + assert_eq!(history.prefix_at(history.total() - 3), Some([13, 14, 15])); + assert_eq!( + history.prefix_at(history.total() - 2), + None, + "fewer than three bytes follow the position" + ); + let oldest = history.total() - as_u64(HISTORY_SIZE); + assert_eq!( + history.prefix_at(oldest), + Some([ + history.byte_back(HISTORY_SIZE), + history.byte_back(HISTORY_SIZE - 1), + history.byte_back(HISTORY_SIZE - 2), + ]), + "the oldest position still in the ring is reachable" + ); + assert_eq!( + history.prefix_at(oldest - 1), + None, + "one position older has left the ring" + ); + + // A push larger than the ring keeps only its tail. + let big: Vec = (0..HISTORY_SIZE + 7).map(|i| u8::try_from(i % 251).unwrap()).collect(); + history.push(&big); + assert_eq!(history.len(), HISTORY_SIZE); + assert_eq!(history.byte_back(1), *big.last().unwrap()); + assert_eq!(history.byte_back(HISTORY_SIZE), big[7]); + } + #[test] fn bit_writer_basic() { let mut writer = BitWriter::new();