use super::Encoder;
use crate::error::Result;
use flate2::{Compress, Compression, FlushCompress};
use jpeg_encoder::{ColorType, Encoder as JpegEncoder};
const TIGHT_EXPLICIT_FILTER: u8 = 0x04;
const TIGHT_FILL: u8 = 0x80;
const TIGHT_JPEG: u8 = 0x90;
const TIGHT_FILTER_PALETTE: u8 = 0x01;
const MAX_TIGHT_PALETTE: usize = 16;
const MIN_BYTES_TO_COMPRESS: usize = 12;
pub struct TightEncoder {
quality: u8,
zlib_compressor: Compress,
uncompressed_buf: Vec<u8>,
rgb_buf: Vec<u8>,
jpeg_buf: Vec<u8>,
pal_keys: [u32; MAX_TIGHT_PALETTE],
pal_rgb: [[u8; 3]; MAX_TIGHT_PALETTE],
pal_count: usize,
}
impl Default for TightEncoder {
fn default() -> Self {
Self::new()
}
}
impl TightEncoder {
pub fn new() -> Self {
Self {
quality: 82, zlib_compressor: Compress::new(Compression::fast(), true),
uncompressed_buf: Vec::with_capacity(128 * 1024),
rgb_buf: Vec::with_capacity(128 * 1024),
jpeg_buf: Vec::with_capacity(64 * 1024),
pal_keys: [0; MAX_TIGHT_PALETTE],
pal_rgb: [[0; 3]; MAX_TIGHT_PALETTE],
pal_count: 0,
}
}
pub fn set_quality(&mut self, quality: u8) {
self.quality = quality.clamp(1, 100);
}
#[inline]
fn write_compact_len(out: &mut Vec<u8>, len: usize) {
if len < 128 {
out.push(len as u8);
} else if len < 16384 {
out.push(((len & 0x7F) | 0x80) as u8);
out.push((len >> 7) as u8);
} else {
out.push(((len & 0x7F) | 0x80) as u8);
out.push((((len >> 7) & 0x7F) | 0x80) as u8);
out.push((len >> 14) as u8);
}
}
#[inline(always)]
fn pixel_rgb(pixels: &[u8], off: usize, swap_rb: bool) -> ([u8; 3], u32) {
let (r, g, b) = if swap_rb {
(pixels[off], pixels[off + 1], pixels[off + 2])
} else {
(pixels[off + 2], pixels[off + 1], pixels[off])
};
let key = (r as u32) | ((g as u32) << 8) | ((b as u32) << 16);
([r, g, b], key)
}
#[inline]
fn is_solid_rect(
pixels: &[u8],
stride: usize,
rx: usize,
ry: usize,
rw: usize,
rh: usize,
) -> Option<[u8; 4]> {
let first = ry * stride + rx * 4;
if first + 4 > pixels.len() {
return None;
}
let first_px = [
pixels[first],
pixels[first + 1],
pixels[first + 2],
pixels[first + 3],
];
let px_u32 = u32::from_ne_bytes(first_px);
let px_u64 = ((px_u32 as u64) << 32) | (px_u32 as u64);
let corners = [
ry * stride + (rx + rw - 1) * 4,
(ry + rh - 1) * stride + rx * 4,
(ry + rh - 1) * stride + (rx + rw - 1) * 4,
(ry + rh / 2) * stride + (rx + rw / 2) * 4,
];
for &c in &corners {
if c + 4 <= pixels.len()
&& u32::from_ne_bytes([pixels[c], pixels[c + 1], pixels[c + 2], pixels[c + 3]])
!= px_u32
{
return None;
}
}
let row_bytes = rw * 4;
for row in ry..ry + rh {
let rs = row * stride + rx * 4;
let re = rs + row_bytes;
if re > pixels.len() {
return None;
}
let row_slice = &pixels[rs..re];
let mut u64_chunks = row_slice.chunks_exact(8);
for chunk in u64_chunks.by_ref() {
if u64::from_ne_bytes(chunk.try_into().unwrap()) != px_u64 {
return None;
}
}
for chunk in u64_chunks.remainder().chunks_exact(4) {
if u32::from_ne_bytes(chunk.try_into().unwrap()) != px_u32 {
return None;
}
}
}
Some(first_px)
}
#[allow(clippy::too_many_arguments)]
fn analyze_palette(
&mut self,
pixels: &[u8],
stride: usize,
rx: usize,
ry: usize,
rw: usize,
rh: usize,
swap_rb: bool,
) -> bool {
self.pal_count = 0;
for row in ry..ry + rh {
let base = row * stride + rx * 4;
for col in 0..rw {
let off = base + col * 4;
let (rgb, key) = Self::pixel_rgb(pixels, off, swap_rb);
let mut found = false;
for i in 0..self.pal_count {
if self.pal_keys[i] == key {
found = true;
break;
}
}
if !found {
if self.pal_count >= MAX_TIGHT_PALETTE {
return false;
}
self.pal_keys[self.pal_count] = key;
self.pal_rgb[self.pal_count] = rgb;
self.pal_count += 1;
}
}
}
true
}
#[inline]
fn pal_index(&self, key: u32) -> u8 {
for i in 0..self.pal_count {
if self.pal_keys[i] == key {
return i as u8;
}
}
0
}
#[allow(clippy::too_many_arguments)]
fn encode_palette_indexed(
&mut self,
pixels: &[u8],
stride: usize,
rx: usize,
ry: usize,
rw: usize,
rh: usize,
swap_rb: bool,
out: &mut Vec<u8>,
) -> Result<()> {
let num_colors = self.pal_count;
out.push(TIGHT_EXPLICIT_FILTER); out.push(TIGHT_FILTER_PALETTE);
out.push((num_colors - 1) as u8);
for i in 0..num_colors {
out.extend_from_slice(&self.pal_rgb[i]);
}
self.uncompressed_buf.clear();
if num_colors == 2 {
for row in ry..ry + rh {
let base = row * stride + rx * 4;
let mut byte = 0u8;
let mut bits = 0usize;
for col in 0..rw {
let off = base + col * 4;
let (_, key) = Self::pixel_rgb(pixels, off, swap_rb);
let idx = self.pal_index(key);
byte = (byte << 1) | (idx & 1);
bits += 1;
if bits == 8 {
self.uncompressed_buf.push(byte);
byte = 0;
bits = 0;
}
}
if bits > 0 {
byte <<= 8 - bits;
self.uncompressed_buf.push(byte);
}
}
} else {
for row in ry..ry + rh {
let base = row * stride + rx * 4;
for col in 0..rw {
let off = base + col * 4;
let (_, key) = Self::pixel_rgb(pixels, off, swap_rb);
let idx = self.pal_index(key);
self.uncompressed_buf.push(idx);
}
}
}
self.compress_and_append(out)
}
#[allow(clippy::too_many_arguments)]
fn encode_jpeg(
&mut self,
pixels: &[u8],
stride: usize,
rx: usize,
ry: usize,
rw: usize,
rh: usize,
swap_rb: bool,
out: &mut Vec<u8>,
) -> Result<()> {
self.rgb_buf.clear();
let total_rgb = rw * rh * 3;
self.rgb_buf.reserve(total_rgb);
for row in ry..ry + rh {
let base = row * stride + rx * 4;
for col in 0..rw {
let off = base + col * 4;
let (rgb, _) = Self::pixel_rgb(pixels, off, swap_rb);
self.rgb_buf.extend_from_slice(&rgb);
}
}
self.jpeg_buf.clear();
{
let encoder = JpegEncoder::new(&mut self.jpeg_buf, self.quality);
encoder
.encode(&self.rgb_buf, rw as u16, rh as u16, ColorType::Rgb)
.map_err(|e| crate::error::VncError::Encoding(format!("JPEG error: {}", e)))?;
}
out.push(TIGHT_JPEG);
Self::write_compact_len(out, self.jpeg_buf.len());
out.extend_from_slice(&self.jpeg_buf);
Ok(())
}
fn compress_and_append(&mut self, out: &mut Vec<u8>) -> Result<()> {
let raw_len = self.uncompressed_buf.len();
if raw_len < MIN_BYTES_TO_COMPRESS {
Self::write_compact_len(out, raw_len);
out.extend_from_slice(&self.uncompressed_buf);
return Ok(());
}
let mut temp_compressed = Vec::with_capacity(raw_len);
self.zlib_compressor
.compress_vec(&self.uncompressed_buf, &mut temp_compressed, FlushCompress::Sync)
.map_err(|e| crate::error::VncError::Encoding(e.to_string()))?;
Self::write_compact_len(out, temp_compressed.len());
out.extend_from_slice(&temp_compressed);
Ok(())
}
}
impl Encoder for TightEncoder {
fn encoding_id(&self) -> i32 {
7
}
#[allow(clippy::too_many_arguments)]
fn encode_rect_into(
&mut self,
pixels: &[u8],
stride: usize,
x: u16,
y: u16,
w: u16,
h: u16,
swap_rb: bool,
out: &mut Vec<u8>,
) -> Result<()> {
let rx = x as usize;
let ry = y as usize;
let rw = w as usize;
let rh = h as usize;
if rw == 0 || rh == 0 {
out.push(TIGHT_FILL);
out.extend_from_slice(&[0, 0, 0]);
return Ok(());
}
if let Some(mut color) = Self::is_solid_rect(pixels, stride, rx, ry, rw, rh) {
if swap_rb {
color.swap(0, 2);
}
out.push(TIGHT_FILL);
out.push(color[2]); out.push(color[1]); out.push(color[0]); return Ok(());
}
let is_palette = self.analyze_palette(pixels, stride, rx, ry, rw, rh, swap_rb);
if is_palette {
if self.pal_count <= 1 {
let color = if self.pal_count == 1 {
self.pal_rgb[0]
} else {
[0, 0, 0]
};
out.push(TIGHT_FILL);
out.extend_from_slice(&color);
return Ok(());
}
return self.encode_palette_indexed(pixels, stride, rx, ry, rw, rh, swap_rb, out);
}
self.encode_jpeg(pixels, stride, rx, ry, rw, rh, swap_rb, out)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn solid_frame(w: usize, h: usize, color: [u8; 4]) -> (Vec<u8>, usize) {
let stride = w * 4;
let mut buf = vec![0u8; stride * h];
for pixel in buf.chunks_exact_mut(4) {
pixel.copy_from_slice(&color);
}
(buf, stride)
}
#[test]
fn test_tight_solid_fill() {
let mut enc = TightEncoder::new();
let (frame, stride) = solid_frame(64, 64, [10, 20, 30, 255]);
let mut out = Vec::new();
enc.encode_rect_into(&frame, stride, 0, 0, 64, 64, false, &mut out).unwrap();
assert_eq!(out.len(), 4);
assert_eq!(out[0], TIGHT_FILL);
assert_eq!(out[1], 30); assert_eq!(out[2], 20); assert_eq!(out[3], 10); }
#[test]
fn test_tight_palette_two_colors() {
let mut enc = TightEncoder::new();
let w = 64usize;
let h = 64usize;
let stride = w * 4;
let mut frame = vec![0u8; stride * h];
for y in 0..h {
let c = if y < h / 2 { [255, 0, 0, 255] } else { [0, 255, 0, 255] };
for x in 0..w {
let off = y * stride + x * 4;
frame[off..off + 4].copy_from_slice(&c);
}
}
let mut out = Vec::new();
enc.encode_rect_into(&frame, stride, 0, 0, w as u16, h as u16, false, &mut out).unwrap();
assert_eq!(out[0], TIGHT_EXPLICIT_FILTER);
assert_eq!(out[1], TIGHT_FILTER_PALETTE);
assert_eq!(out[2], 1); assert!(out.len() < 200);
}
#[test]
fn test_tight_jpeg_high_colors() {
let mut enc = TightEncoder::new();
let w = 64usize;
let h = 64usize;
let stride = w * 4;
let mut frame = vec![0u8; stride * h];
for (i, b) in frame.iter_mut().enumerate() {
*b = ((i * 17 + 53) % 256) as u8;
}
let mut out = Vec::new();
enc.encode_rect_into(&frame, stride, 0, 0, w as u16, h as u16, false, &mut out).unwrap();
assert_eq!(out[0], TIGHT_JPEG);
assert!(out.len() > 10);
}
}