Do population stage 0 in parallel

Make HashMap extension fn like get_many_mut but for variable sizes
  See also rust-lang/hashbrown#332

fs_common/src/game/common/hashmap_ext.rs

@@ -0,0 +1,38 @@
+use std::{
+    collections::HashMap,
+    hash::{BuildHasher, Hash},
+};
+
+use itertools::Itertools;
+
+/// Workaround for <https://github.com/rust-lang/hashbrown/issues/332>
+#[allow(clippy::missing_safety_doc)]
+pub trait HashMapExt {
+    type K;
+    type V;
+
+    unsafe fn get_many_var_unchecked_mut(&mut self, keys: &[Self::K]) -> Option<Vec<&mut Self::V>>;
+    fn get_many_var_mut(&mut self, keys: &[Self::K]) -> Option<Vec<&mut Self::V>>;
+}
+
+impl<K: Eq + Hash, V, S: BuildHasher> HashMapExt for HashMap<K, V, S> {
+    type K = K;
+    type V = V;
+
+    unsafe fn get_many_var_unchecked_mut(&mut self, keys: &[K]) -> Option<Vec<&mut V>> {
+        let out = keys
+            .iter()
+            .map(|k| self.get_mut(k).map(|v| &mut *(v as &mut _ as *mut _)))
+            .collect();
+        out
+    }
+
+    fn get_many_var_mut(&mut self, keys: &[K]) -> Option<Vec<&mut V>> {
+        let unique = keys.iter().duplicates().next().is_none();
+        if unique {
+            unsafe { self.get_many_var_unchecked_mut(keys) }
+        } else {
+            None
+        }
+    }
+}

fs_common/src/game/common/mod.rs

@@ -2,6 +2,7 @@ pub mod networking;
 pub mod world;
 
 pub mod cli;
+pub mod hashmap_ext;
 mod settings;
 use std::ops::Range;
 

fs_common/src/game/common/world/chunk.rs

@@ -1,3 +1,4 @@
+use crate::game::common::hashmap_ext::HashMapExt;
 use crate::game::common::world::gen::populator::ChunkContext;
 use crate::game::common::world::particle::ParticleSystem;
 use crate::game::common::world::simulator::Simulator;
@@ -128,7 +129,7 @@ struct ChunkSaveFormat {
     colors: Vec<u8>,
 }
 
-impl<C: Chunk> ChunkHandlerGeneric for ChunkHandler<C> {
+impl<C: Chunk + Send> ChunkHandlerGeneric for ChunkHandler<C> {
     #[profiling::function]
     fn update_chunk_graphics(&mut self) {
         let keys = self.loaded_chunks.keys().copied().collect::<Vec<u32>>();
@@ -498,22 +499,30 @@ impl<C: Chunk> ChunkHandlerGeneric for ChunkHandler<C> {
                             })
                             .collect()
                     };
-                    for ch in generated {
-                        profiling::scope!("finish chunk");
-
-                        // TODO: this is very slow
-
-                        let chunk = self.loaded_chunks.get_mut(ch.0).unwrap();
-                        chunk.set_pixels(*ch.1);
-                        chunk.set_pixel_colors(*ch.2);
-
-                        // TODO: populating stage 0 should be able to be multithreaded
-                        self.generator.populators().populate(
-                            0,
-                            &mut [chunk as &mut dyn Chunk],
-                            seed,
-                            registries,
-                        );
+
+                    let keys: Vec<_> = generated
+                        .into_iter()
+                        .map(|ch| {
+                            profiling::scope!("finish chunk");
+
+                            let chunk = self.loaded_chunks.get_mut(ch.0).unwrap();
+                            chunk.set_pixels(*ch.1);
+                            chunk.set_pixel_colors(*ch.2);
+
+                            *ch.0
+                        })
+                        .collect();
+
+                    let pops = self.generator.populators();
+                    {
+                        profiling::scope!("populate stage 0");
+                        self.loaded_chunks
+                            .get_many_var_mut(&keys)
+                            .unwrap()
+                            .into_par_iter()
+                            .for_each(|chunk| {
+                                pops.populate(0, &mut [chunk as &mut dyn Chunk], seed, registries);
+                            });
                     }
                 }
             }
@@ -596,36 +605,31 @@ impl<C: Chunk> ChunkHandlerGeneric for ChunkHandler<C> {
                                         _ => false,
                                     }
                                 }) {
-                                    let mut chunks = Vec::new();
+                                    let mut keys = Vec::new();
 
                                     let range = i32::from(cur_stage + 1);
 
                                     // try to gather the nearby chunks needed to populate this one
-                                    let mut failed = false;
-                                    'outer: for y in -range..=range {
+                                    for y in -range..=range {
                                         for x in -range..=range {
-                                            let c = self
-                                                .loaded_chunks
-                                                .get_mut(&chunk_index(chunk_x + x, chunk_y + y));
-                                            if let Some(c) = c {
-                                                let c = c as &mut dyn Chunk as *mut _;
-                                                // this is just blatantly bypassing the borrow checker to get mutable refs to multiple unique hashmap entries
-                                                // TODO: can probably use get_many_mut once stabilized: https://github.com/rust-lang/rust/issues/97601
-                                                chunks.push(unsafe { &mut *c });
-                                            } else {
-                                                failed = true;
-                                                break 'outer;
-                                            }
+                                            keys.push(chunk_index(chunk_x + x, chunk_y + y));
                                         }
                                     }
 
+                                    let chunks = self.loaded_chunks.get_many_var_mut(&keys);
+
                                     // if we failed to get all nearby chunks, don't populate and don't go to the next stage
-                                    if !failed {
+                                    if let Some(chunks) = chunks {
+                                        let mut chunks_dyn: Vec<_> = chunks
+                                            .into_iter()
+                                            .map(|c| c as &mut dyn Chunk)
+                                            .collect();
+
                                         // TODO: make not hardcoded
 
                                         if cur_stage + 1 == 1 {
                                             let mut ctx =
-                                                ChunkContext::<1>::new(&mut chunks).unwrap();
+                                                ChunkContext::<1>::new(&mut chunks_dyn).unwrap();
                                             let mut rng = StdRng::seed_from_u64(
                                                 seed as u64
                                                     + u64::from(chunk_index(
@@ -640,7 +644,7 @@ impl<C: Chunk> ChunkHandlerGeneric for ChunkHandler<C> {
 
                                         self.generator.populators().populate(
                                             cur_stage + 1,
-                                            &mut chunks,
+                                            &mut chunks_dyn,
                                             seed,
                                             registries,
                                         );
@@ -742,7 +746,6 @@ impl<C: Chunk> ChunkHandlerGeneric for ChunkHandler<C> {
                                     .as_mut()
                                     .unwrap();
                                 // blatantly bypassing the borrow checker, see safety comment above
-                                // TODO: can probably use get_many_mut once stabilized: https://github.com/rust-lang/rust/issues/97601
                                 unsafe { &mut *raw }
                             });
 
@@ -766,7 +769,6 @@ impl<C: Chunk> ChunkHandlerGeneric for ChunkHandler<C> {
                                     .unwrap()
                                     .get_colors_mut();
                                 // blatantly bypassing the borrow checker, see safety comment above
-                                // TODO: can probably use get_many_mut once stabilized: https://github.com/rust-lang/rust/issues/97601
                                 unsafe { &mut *raw }
                             });
 

fs_common/src/game/common/world/copy_paste.rs

@@ -22,7 +22,7 @@ impl MaterialBuf {
         }
     }
 
-    pub fn copy<C: Chunk>(
+    pub fn copy<C: Chunk + Send>(
         chunk_handler: &ChunkHandler<C>,
         x: impl Into<i64>,
         y: impl Into<i64>,
@@ -47,7 +47,7 @@ impl MaterialBuf {
         Ok(Self { width, height, materials: buf })
     }
 
-    pub fn cut<C: Chunk>(
+    pub fn cut<C: Chunk + Send>(
         chunk_handler: &mut ChunkHandler<C>,
         x: impl Into<i64>,
         y: impl Into<i64>,
@@ -73,7 +73,7 @@ impl MaterialBuf {
         Ok(Self { width, height, materials: buf })
     }
 
-    pub fn paste<C: Chunk>(
+    pub fn paste<C: Chunk + Send>(
         &self,
         chunk_handler: &mut ChunkHandler<C>,
         x: impl Into<i64>,

fs_common/src/game/common/world/simulator.rs

@@ -233,7 +233,7 @@ struct SimulationHelperRigidBody<'a, C: Chunk> {
     physics: &'a mut Physics,
 }
 
-impl<C: Chunk> SimulationHelper for SimulationHelperRigidBody<'_, C> {
+impl<C: Chunk + Send> SimulationHelper for SimulationHelperRigidBody<'_, C> {
     fn get_pixel_local(&self, x: i32, y: i32) -> MaterialInstance {
         let world_mat = self.chunk_handler.get(i64::from(x), i64::from(y)); // TODO: consider changing the args to i64
         if let Ok(m) = world_mat {
@@ -402,7 +402,7 @@ impl Simulator {
     #[allow(clippy::unnecessary_unwrap)]
     #[allow(clippy::needless_range_loop)]
     #[profiling::function]
-    pub fn simulate_rigidbodies<C: Chunk>(
+    pub fn simulate_rigidbodies<C: Chunk + Send>(
         chunk_handler: &mut ChunkHandler<C>,
         rigidbodies: &mut Vec<FSRigidBody>,
         physics: &mut Physics,

fs_common/src/game/common/world/world.rs

@@ -55,7 +55,7 @@ pub struct World<C: Chunk> {
     pub seed: i32,
 }
 
-impl<C: Chunk> World<C> {
+impl<C: Chunk + Send> World<C> {
     #[profiling::function]
     pub fn create(path: Option<PathBuf>, seed: Option<i32>) -> Self {
         let mut ecs = specs::World::new();

fs_common/src/game/common/world/world_loading.rs

@@ -20,7 +20,7 @@ pub struct WorldMeta {
     pub last_played_time: toml::value::Datetime,
 }
 
-impl<C: Chunk> World<C> {
+impl<C: Chunk + Send> World<C> {
     pub fn find_files(root: PathBuf) -> Result<WorldTreeNode<PathBuf, PathBuf>, std::io::Error> {
         let mut res = Vec::new();
         for entry in fs::read_dir(&root)? {

fs_common/src/game/game.rs

@@ -62,7 +62,7 @@ pub struct FPSCounter {
     pub tick_physics_times: [f32; 200],
 }
 
-impl<C: Chunk> GameData<C> {
+impl<C: Chunk + Send> GameData<C> {
     #[profiling::function]
     pub fn new(file_helper: FileHelper, build_data: BuildData) -> Self {
         GameData {