Add benchmark for color quantization using the rtree crate.

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",

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<Subpixel = u
     ImageBuffer::from_raw(w, h, buf)
 }
 
+impl<'a, VIE: VoxelImageEncoding> 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<u8>) {
+        (*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<u8>,
+    ) {
+        (*self).put_sprite(ws, x, y, kind, sprite, ori)
+    }
+
+    fn finish(ws: &Self::Workspace) -> Option<Self::Output> { VIE::finish(ws) }
+}
+impl<'a, VIE: VoxelImageDecoding> VoxelImageDecoding for &'a VIE {
+    fn start(ws: &Self::Output) -> Option<Self::Workspace> { 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<const RESOLUTION_DIVIDER: u32>();
 
@@ -626,7 +657,7 @@ impl<const AVERAGE_PALETTE: bool> VoxelImageDecoding for TriPngEncoding<AVERAGE_
 }
 
 pub fn image_terrain_chonk<S: RectVolSize, M: Clone, P: PackingFormula, VIE: VoxelImageEncoding>(
-    vie: VIE,
+    vie: &VIE,
     packing: P,
     chonk: &Chonk<Block, S, M>,
 ) -> Option<VIE::Output> {
@@ -645,7 +676,7 @@ pub fn image_terrain_volgrid<
     P: PackingFormula,
     VIE: VoxelImageEncoding,
 >(
-    vie: VIE,
+    vie: &VIE,
     packing: P,
     volgrid: &VolGrid2d<Chonk<Block, S, M>>,
 ) -> Option<VIE::Output> {
@@ -669,7 +700,7 @@ pub fn image_terrain<
     P: PackingFormula,
     VIE: VoxelImageEncoding,
 >(
-    vie: VIE,
+    vie: &VIE,
     packing: P,
     vol: &V,
     lo: Vec3<u32>,
@@ -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<VIE: VoxelImageEncoding + VoxelImageDecoding, P: PackingFormula, M: Clone,
     WireChonk<VIE, P, M, S>
 {
     pub fn from_chonk(vie: VIE, packing: P, chonk: &Chonk<Block, S, M>) -> Option<Self> {
-        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<VIE: VoxelImageEncoding + VoxelImageDecoding, P: PackingFormula, M: Clone,
     pub fn to_chonk(&self) -> Option<Chonk<Block, S, M>> {
         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,

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 }
 
 

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<u8>,
+    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<BlockKind, KdTree<f32, u8, 3>> = {
+    static ref PALETTE_RTREE: HashMap<BlockKind, RTree<ColorPoint, TestParams>> = {
+        let ron_bytes = include_bytes!("palettes.ron");
+        let palettes: HashMap<BlockKind, Vec<Rgb<u8>>> =
+            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<BlockKind, KdTree<f32, u8, 3>> = {
         let ron_bytes = include_bytes!("palettes.ron");
         let palettes: HashMap<BlockKind, Vec<Rgb<u8>>> =
             ron::de::from_bytes(&*ron_bytes).expect("palette should parse");
@@ -514,10 +566,32 @@ lazy_static::lazy_static! {
     };
 }
 
-#[derive(Debug, Clone, Copy)]
+pub trait NearestNeighbor {
-pub struct PaletteEncoding<'a, const N: u32>(&'a HashMap<BlockKind, KdTree<f32, u8, 3>>);
+    fn nearest_neighbor(&self, x: &Rgb<u8>) -> Option<u8>;
+}
 
-impl<'a, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, N> {
+impl NearestNeighbor for KdTree<f32, u8, 3> {
+    fn nearest_neighbor(&self, x: &Rgb<u8>) -> Option<u8> {
+        self.nearest_one(
+            &[x.r as f32, x.g as f32, x.b as f32],
+            &kiddo::distance::squared_euclidean,
+        )
+        .map(|(_, i)| *i)
+        .ok()
+    }
+}
+
+impl<P: RTreeParams> NearestNeighbor for RTree<ColorPoint, P> {
+    fn nearest_neighbor(&self, x: &Rgb<u8>) -> Option<u8> {
+        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<BlockKind, NN>);
+
+impl<'a, NN: NearestNeighbor, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, NN, N> {
     #[allow(clippy::type_complexity)]
     type Output = CompressedData<(Vec<u8>, [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<u8>) {
         ws.0.put_pixel(x, y, image::Luma([kind as u8]));
-        let i = self.0[&kind]
+        let i = self.0[&kind].nearest_neighbor(&rgb).unwrap_or(0);
-            .nearest_one(
-                &[rgb.r as f32, rgb.g as f32, rgb.b as f32],
-                &kiddo::distance::squared_euclidean,
-            )
-            .map(|(_, i)| *i)
-            .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::<true>(), &chunk)
+                    image_terrain_chonk(&QuadPngEncoding::<4>(), WidePacking::<true>(), &chunk)
                         .unwrap();
                 let quadpngquartwide_post = Instant::now();
 
                 let tripngaverage_pre = Instant::now();
                 let tripngaverage =
-                    image_terrain_chonk(TriPngEncoding::<true>(), WidePacking::<true>(), &chunk)
+                    image_terrain_chonk(&TriPngEncoding::<true>(), WidePacking::<true>(), &chunk)
                         .unwrap();
                 let tripngaverage_post = Instant::now();
 
                 let tripngconst_pre = Instant::now();
                 let tripngconst =
-                    image_terrain_chonk(TriPngEncoding::<false>(), WidePacking::<true>(), &chunk)
+                    image_terrain_chonk(&TriPngEncoding::<false>(), WidePacking::<true>(), &chunk)
                         .unwrap();
                 let tripngconst_post = Instant::now();
 
-                let palettepng_pre = Instant::now();
+                let palette_kdtree_pre = Instant::now();
-                let palettepng = image_terrain_chonk(
+                let palette_kdtree = image_terrain_chonk(
-                    PaletteEncoding::<4>(&PALETTE),
+                    &PaletteEncoding::<_, 4>(&PALETTE_KDTREE),
                     WidePacking::<true>(),
                     &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::<true>(),
+                    &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!(