From f9c263faf6e39f1326a14a9655856237f201a3f3 Mon Sep 17 00:00:00 2001 From: Avi Weinstock Date: Thu, 1 Jul 2021 13:41:37 -0400 Subject: [PATCH] Add benchmark for color quantization using the rtree crate. --- Cargo.lock | 33 +++- common/net/src/msg/compression.rs | 47 +++++- world/Cargo.toml | 3 +- .../examples/chunk_compression_benchmarks.rs | 155 ++++++++++++++---- 4 files changed, 195 insertions(+), 43 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b8e5e4307c..f35a75a634 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -106,7 +106,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "01aa95c33b750236bd69526dd773d618c5c8ad5e156daad1fc760dfd4be8def3" dependencies = [ - "heapless", + "heapless 0.5.6", "nom 4.2.3", ] @@ -2269,6 +2269,18 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "heapless" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "634bd4d29cbf24424d0a4bfcbf80c6960129dc24424752a7d1d1390607023422" +dependencies = [ + "as-slice", + "generic-array 0.14.4", + "hash32", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.3.3" @@ -3690,6 +3702,12 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "pdqselect" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec91767ecc0a0bbe558ce8c9da33c068066c57ecc8bb8477ef8c1ad3ef77c27" + [[package]] name = "peeking_take_while" version = "0.1.2" @@ -4374,6 +4392,18 @@ dependencies = [ "petgraph 0.5.1", ] +[[package]] +name = "rstar" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce61d743ebe516592df4dd542dfe823577b811299f7bee1106feb1bbb993dbac" +dependencies = [ + "heapless 0.6.1", + "num-traits", + "pdqselect", + "smallvec", +] + [[package]] name = "rusqlite" version = "0.24.2" @@ -6167,6 +6197,7 @@ dependencies = [ "rand_chacha 0.3.0", "rayon", "ron", + "rstar", "rusqlite", "serde", "structopt", diff --git a/common/net/src/msg/compression.rs b/common/net/src/msg/compression.rs index 98136476ec..3d5a90e002 100644 --- a/common/net/src/msg/compression.rs +++ b/common/net/src/msg/compression.rs @@ -123,7 +123,7 @@ impl PackingFormula for GridLtrPacking { } } -pub trait VoxelImageEncoding: Copy { +pub trait VoxelImageEncoding { type Workspace; type Output; fn create(width: u32, height: u32) -> Self::Workspace; @@ -154,6 +154,37 @@ pub fn image_from_bytes<'a, I: ImageDecoder<'a>, P: 'static + Pixel VoxelImageEncoding for &'a VIE { + type Output = VIE::Output; + type Workspace = VIE::Workspace; + + fn create(width: u32, height: u32) -> Self::Workspace { VIE::create(width, height) } + + fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb) { + (*self).put_solid(ws, x, y, kind, rgb) + } + + fn put_sprite( + &self, + ws: &mut Self::Workspace, + x: u32, + y: u32, + kind: BlockKind, + sprite: SpriteKind, + ori: Option, + ) { + (*self).put_sprite(ws, x, y, kind, sprite, ori) + } + + fn finish(ws: &Self::Workspace) -> Option { VIE::finish(ws) } +} +impl<'a, VIE: VoxelImageDecoding> VoxelImageDecoding for &'a VIE { + fn start(ws: &Self::Output) -> Option { VIE::start(ws) } + + fn get_block(ws: &Self::Workspace, x: u32, y: u32, is_border: bool) -> Block { + VIE::get_block(ws, x, y, is_border) + } +} #[derive(Debug, Clone, Copy, Serialize, Deserialize)] pub struct QuadPngEncoding(); @@ -626,7 +657,7 @@ impl VoxelImageDecoding for TriPngEncoding( - vie: VIE, + vie: &VIE, packing: P, chonk: &Chonk, ) -> Option { @@ -645,7 +676,7 @@ pub fn image_terrain_volgrid< P: PackingFormula, VIE: VoxelImageEncoding, >( - vie: VIE, + vie: &VIE, packing: P, volgrid: &VolGrid2d>, ) -> Option { @@ -669,7 +700,7 @@ pub fn image_terrain< P: PackingFormula, VIE: VoxelImageEncoding, >( - vie: VIE, + vie: &VIE, packing: P, vol: &V, lo: Vec3, @@ -700,10 +731,10 @@ pub fn image_terrain< .unwrap_or(&Block::empty()); match (block.get_color(), block.get_sprite()) { (Some(rgb), None) => { - VIE::put_solid(&vie, &mut image, i, j, *block, rgb); + VIE::put_solid(vie, &mut image, i, j, *block, rgb); }, (None, Some(sprite)) => { - VIE::put_sprite(&vie, &mut image, i, j, *block, sprite, block.get_ori()); + VIE::put_sprite(vie, &mut image, i, j, *block, sprite, block.get_ori()); }, _ => panic!( "attr being used for color vs sprite is mutually exclusive (and that's \ @@ -772,7 +803,7 @@ impl { pub fn from_chonk(vie: VIE, packing: P, chonk: &Chonk) -> Option { - let data = image_terrain_chonk(vie, packing, chonk)?; + let data = image_terrain_chonk(&vie, packing, chonk)?; Some(Self { zmin: chonk.get_min_z(), zmax: chonk.get_max_z(), @@ -791,7 +822,7 @@ impl Option> { let mut chonk = Chonk::new(self.zmin, self.below, self.above, self.meta.clone()); write_image_terrain( - self.vie, + &self.vie, self.packing, &mut chonk, &self.data, diff --git a/world/Cargo.toml b/world/Cargo.toml index b6d4ef4666..4269b85634 100644 --- a/world/Cargo.toml +++ b/world/Cargo.toml @@ -6,7 +6,7 @@ edition = "2018" [features] simd = ["vek/platform_intrinsics"] -bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "kiddo", "clap"] +bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "kiddo", "clap", "rstar"] default = ["simd"] @@ -43,6 +43,7 @@ flate2 = { version = "1.0.20", optional = true } num-traits = { version = "0.2", optional = true } fallible-iterator = { version = "0.2.0", optional = true } kiddo = { version = "0.1.4", optional = true } +rstar = { version = "0.8.3", optional = true } clap = { version = "2.33.3", optional = true } diff --git a/world/examples/chunk_compression_benchmarks.rs b/world/examples/chunk_compression_benchmarks.rs index 538612618e..e216c1badf 100644 --- a/world/examples/chunk_compression_benchmarks.rs +++ b/world/examples/chunk_compression_benchmarks.rs @@ -494,9 +494,61 @@ impl VoxelImageEncoding for MixedEncodingDenseSprites { } use kiddo::KdTree; +use rstar::{PointDistance, RTree, RTreeObject, RTreeParams}; + +#[derive(Debug)] +struct ColorPoint { + rgb: Rgb, + index: u8, +} + +impl RTreeObject for ColorPoint { + type Envelope = <[i32; 3] as RTreeObject>::Envelope; + + fn envelope(&self) -> Self::Envelope { + [self.rgb.r as i32, self.rgb.g as i32, self.rgb.b as i32].envelope() + } +} + +impl PointDistance for ColorPoint { + fn distance_2(&self, other: &[i32; 3]) -> i32 { + (self.rgb.r as i32 - other[0]).pow(2) + + (self.rgb.g as i32 - other[1]).pow(2) + + (self.rgb.b as i32 - other[2]).pow(2) + } + + fn contains_point(&self, other: &[i32; 3]) -> bool { + &[self.rgb.r as i32, self.rgb.g as i32, self.rgb.b as i32] == other + } +} + +struct TestParams; +impl RTreeParams for TestParams { + type DefaultInsertionStrategy = rstar::RStarInsertionStrategy; + + const MAX_SIZE: usize = 4; + const MIN_SIZE: usize = 2; + const REINSERTION_COUNT: usize = 1; +} lazy_static::lazy_static! { - pub static ref PALETTE: HashMap> = { + static ref PALETTE_RTREE: HashMap> = { + let ron_bytes = include_bytes!("palettes.ron"); + let palettes: HashMap>> = + ron::de::from_bytes(&*ron_bytes).expect("palette should parse"); + palettes + .into_iter() + .map(|(k, v)| { + let tree = RTree::bulk_load_with_params(v.into_iter() + .enumerate() + .map(|(index, rgb)| ColorPoint { rgb, index: index as u8 }) + .collect() + ); + (k, tree) + }) + .collect() + }; + pub static ref PALETTE_KDTREE: HashMap> = { let ron_bytes = include_bytes!("palettes.ron"); let palettes: HashMap>> = ron::de::from_bytes(&*ron_bytes).expect("palette should parse"); @@ -514,10 +566,32 @@ lazy_static::lazy_static! { }; } -#[derive(Debug, Clone, Copy)] -pub struct PaletteEncoding<'a, const N: u32>(&'a HashMap>); +pub trait NearestNeighbor { + fn nearest_neighbor(&self, x: &Rgb) -> Option; +} -impl<'a, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, N> { +impl NearestNeighbor for KdTree { + fn nearest_neighbor(&self, x: &Rgb) -> Option { + self.nearest_one( + &[x.r as f32, x.g as f32, x.b as f32], + &kiddo::distance::squared_euclidean, + ) + .map(|(_, i)| *i) + .ok() + } +} + +impl NearestNeighbor for RTree { + fn nearest_neighbor(&self, x: &Rgb) -> Option { + self.nearest_neighbor(&[x.r as i32, x.g as i32, x.b as i32]) + .map(|p| p.index) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct PaletteEncoding<'a, NN: NearestNeighbor, const N: u32>(&'a HashMap); + +impl<'a, NN: NearestNeighbor, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, NN, N> { #[allow(clippy::type_complexity)] type Output = CompressedData<(Vec, [usize; 4])>; #[allow(clippy::type_complexity)] @@ -539,13 +613,7 @@ impl<'a, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, N> { fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb) { ws.0.put_pixel(x, y, image::Luma([kind as u8])); - let i = self.0[&kind] - .nearest_one( - &[rgb.r as f32, rgb.g as f32, rgb.b as f32], - &kiddo::distance::squared_euclidean, - ) - .map(|(_, i)| *i) - .unwrap_or(0); + let i = self.0[&kind].nearest_neighbor(&rgb).unwrap_or(0); ws.3.put_pixel(x / N, y / N, image::Luma([i])); } @@ -877,7 +945,7 @@ fn main() { if !SKIP_IMAGECHONK { let jpegchonkgrid_pre = Instant::now(); let jpegchonkgrid = - image_terrain_chonk(JpegEncoding, GridLtrPacking, &chunk).unwrap(); + image_terrain_chonk(&JpegEncoding, GridLtrPacking, &chunk).unwrap(); let jpegchonkgrid_post = Instant::now(); if false { @@ -892,19 +960,19 @@ fn main() { let jpegchonktall_pre = Instant::now(); let jpegchonktall = - image_terrain_chonk(JpegEncoding, TallPacking { flip_y: false }, &chunk) + image_terrain_chonk(&JpegEncoding, TallPacking { flip_y: false }, &chunk) .unwrap(); let jpegchonktall_post = Instant::now(); let jpegchonkflip_pre = Instant::now(); let jpegchonkflip = - image_terrain_chonk(JpegEncoding, TallPacking { flip_y: true }, &chunk) + image_terrain_chonk(&JpegEncoding, TallPacking { flip_y: true }, &chunk) .unwrap(); let jpegchonkflip_post = Instant::now(); let pngchonk_pre = Instant::now(); let pngchonk = - image_terrain_chonk(PngEncoding, GridLtrPacking, &chunk).unwrap(); + image_terrain_chonk(&PngEncoding, GridLtrPacking, &chunk).unwrap(); let pngchonk_post = Instant::now(); sizes.extend_from_slice(&[ @@ -924,7 +992,7 @@ fn main() { if !SKIP_MIXED { let mixedchonk_pre = Instant::now(); let mixedchonk = - image_terrain_chonk(MixedEncoding, TallPacking { flip_y: true }, &chunk) + image_terrain_chonk(&MixedEncoding, TallPacking { flip_y: true }, &chunk) .unwrap(); let mixedchonk_post = Instant::now(); @@ -933,7 +1001,7 @@ fn main() { let mixeddense_pre = Instant::now(); let mixeddense = image_terrain_chonk( - MixedEncodingDenseSprites, + &MixedEncodingDenseSprites, TallPacking { flip_y: true }, &chunk, ) @@ -954,7 +1022,7 @@ fn main() { let quadpngfull_pre = Instant::now(); let quadpngfull = image_terrain_chonk( - QuadPngEncoding::<1>(), + &QuadPngEncoding::<1>(), TallPacking { flip_y: true }, &chunk, ) @@ -963,7 +1031,7 @@ fn main() { let quadpnghalf_pre = Instant::now(); let quadpnghalf = image_terrain_chonk( - QuadPngEncoding::<2>(), + &QuadPngEncoding::<2>(), TallPacking { flip_y: true }, &chunk, ) @@ -972,7 +1040,7 @@ fn main() { let quadpngquarttall_pre = Instant::now(); let quadpngquarttall = image_terrain_chonk( - QuadPngEncoding::<4>(), + &QuadPngEncoding::<4>(), TallPacking { flip_y: true }, &chunk, ) @@ -981,30 +1049,39 @@ fn main() { let quadpngquartwide_pre = Instant::now(); let quadpngquartwide = - image_terrain_chonk(QuadPngEncoding::<4>(), WidePacking::(), &chunk) + image_terrain_chonk(&QuadPngEncoding::<4>(), WidePacking::(), &chunk) .unwrap(); let quadpngquartwide_post = Instant::now(); let tripngaverage_pre = Instant::now(); let tripngaverage = - image_terrain_chonk(TriPngEncoding::(), WidePacking::(), &chunk) + image_terrain_chonk(&TriPngEncoding::(), WidePacking::(), &chunk) .unwrap(); let tripngaverage_post = Instant::now(); let tripngconst_pre = Instant::now(); let tripngconst = - image_terrain_chonk(TriPngEncoding::(), WidePacking::(), &chunk) + image_terrain_chonk(&TriPngEncoding::(), WidePacking::(), &chunk) .unwrap(); let tripngconst_post = Instant::now(); - let palettepng_pre = Instant::now(); - let palettepng = image_terrain_chonk( - PaletteEncoding::<4>(&PALETTE), + let palette_kdtree_pre = Instant::now(); + let palette_kdtree = image_terrain_chonk( + &PaletteEncoding::<_, 4>(&PALETTE_KDTREE), WidePacking::(), &chunk, ) .unwrap(); - let palettepng_post = Instant::now(); + let palette_kdtree_post = Instant::now(); + + let palette_rtree_pre = Instant::now(); + let palette_rtree = image_terrain_chonk( + &PaletteEncoding::<_, 4>(&PALETTE_RTREE), + WidePacking::(), + &chunk, + ) + .unwrap(); + let palette_rtree_post = Instant::now(); #[rustfmt::skip] sizes.extend_from_slice(&[ @@ -1014,7 +1091,8 @@ fn main() { ("quadpngquartwide", quadpngquartwide.data.len() as f32 / n as f32), ("tripngaverage", tripngaverage.data.len() as f32 / n as f32), ("tripngconst", tripngconst.data.len() as f32 / n as f32), - ("palettepng", palettepng.data.len() as f32 / n as f32), + ("palette_kdtree", palette_kdtree.data.len() as f32 / n as f32), + ("palette_rtree", palette_rtree.data.len() as f32 / n as f32), ]); let best_idx = sizes .iter() @@ -1035,7 +1113,8 @@ fn main() { ("quadpngquartwide", (quadpngquartwide_post - quadpngquartwide_pre).subsec_nanos()), ("tripngaverage", (tripngaverage_post - tripngaverage_pre).subsec_nanos()), ("tripngconst", (tripngconst_post - tripngconst_pre).subsec_nanos()), - ("palettepng", (palettepng_post - palettepng_pre).subsec_nanos()), + ("palette_kdtree", (palette_kdtree_post - palette_kdtree_pre).subsec_nanos()), + ("palette_rtree", (palette_rtree_post - palette_rtree_pre).subsec_nanos()), ]); if false { let bucket = z_buckets @@ -1077,12 +1156,21 @@ fn main() { } if true { let bucket = z_buckets - .entry("palettepng") + .entry("palette_kdtree") .or_default() .entry(chunk.get_max_z() - chunk.get_min_z()) .or_insert((0, 0.0)); bucket.0 += 1; - bucket.1 += (palettepng_post - palettepng_pre).subsec_nanos() as f32; + bucket.1 += (palette_kdtree_post - palette_kdtree_pre).subsec_nanos() as f32; + } + if true { + let bucket = z_buckets + .entry("palette_rtree") + .or_default() + .entry(chunk.get_max_z() - chunk.get_min_z()) + .or_insert((0, 0.0)); + bucket.0 += 1; + bucket.1 += (palette_rtree_post - palette_rtree_pre).subsec_nanos() as f32; } trace!( "{} {}: uncompressed: {}, {:?} {} {:?}", @@ -1111,12 +1199,13 @@ fn main() { let mut f = File::create(&format!("chonkjpegs/{}_{}.jpg", sitename, count)) .unwrap(); let jpeg_volgrid = - image_terrain_volgrid(JpegEncoding, GridLtrPacking, &volgrid).unwrap(); + image_terrain_volgrid(&JpegEncoding, GridLtrPacking, &volgrid).unwrap(); f.write_all(&*jpeg_volgrid).unwrap(); let mixedgrid_pre = Instant::now(); let (mixed_volgrid, indices) = - image_terrain_volgrid(MixedEncoding, GridLtrPacking, &volgrid).unwrap(); + image_terrain_volgrid(&MixedEncoding, GridLtrPacking, &volgrid) + .unwrap(); let mixedgrid_post = Instant::now(); let seconds = (mixedgrid_post - mixedgrid_pre).as_secs_f64(); println!(