Migrate off rayon to shin_tasks

sdu/Cargo.toml

@@ -14,7 +14,7 @@ shin-core = { path = "../shin-core" }
 anyhow = { workspace = true }
 clap = { version = "4.0.15", features = ["derive"] }
 tracing-subscriber = { version = "0.3.16", features = ["env-filter"] }
-image = { workspace = true }
+image = { workspace = true, features = ["png"] }
 bytes = "1.2.1"
 itertools = "0.10.5"
 clap_complete = "4.0.6"

sdu/src/main.rs

@@ -836,6 +836,7 @@ fn main() -> Result<()> {
         // .with_span_events(tracing_subscriber::fmt::format::FmtSpan::NEW)
         .compact()
         .init();
+    shin_core::create_task_pools();
     let args = Args::parse();
     match args.action {
         SduAction::GenerateCompletion(command) => generate_command(command),

shin-core/Cargo.toml

@@ -9,6 +9,7 @@ authors = ["DCNick3"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [dependencies]
+shin-tasks = { path = "../shin-tasks" }
 
 # derive macros
 num-traits = "0.2.15"
@@ -31,7 +32,6 @@ glam = { workspace = true }
 float-ord = "0.3.2"
 image = { workspace = true, default-features = false }
 itertools = "0.10.5"
-rayon = "1.5.3"
 smallvec = "1.10.0"
 tracing = "0.1.37"
 smartstring = "1.0.1"

shin-core/src/format/bustup/mod.rs

@@ -5,10 +5,8 @@ use anyhow::{bail, Result};
 use binrw::{BinRead, BinWrite};
 use bitvec::bitbox;
 use image::RgbaImage;
-use rayon::prelude::*;
+use shin_tasks::ParallelSlice;
 use std::collections::HashMap;
-use std::ops::DerefMut;
-use std::sync::Mutex;
 
 use crate::format::text::ZeroString;
 
@@ -130,6 +128,7 @@ pub fn read_bustup(source: &[u8]) -> Result<Bustup> {
             "Two chunks have the same ID, but different contents"
         );
     }
+    let base_chunks = base_chunks.into_iter().collect::<Vec<_>>();
 
     let mut additional_chunks = HashMap::new();
     for chunk in header.iter_additional_chunk_descs() {
@@ -139,46 +138,53 @@ pub fn read_bustup(source: &[u8]) -> Result<Bustup> {
             "Two chunks have the same ID, but different contents"
         );
     }
+    let additional_chunks = additional_chunks.into_iter().collect::<Vec<_>>();
 
     // TODO: ditch rayon?
     // TODO: actually, collecting all of these is not the most efficient in terms of memory...
     // It might be better to first collect the "base" chunks into the base picture (the same way it's done in picture)
     // and then we can start reading the expressions and their mouths.
 
-    let base_image = RgbaImage::new(header.viewport_width as u32, header.viewport_height as u32);
-    let base_image = Mutex::new(base_image);
+    let mut base_image =
+        RgbaImage::new(header.viewport_width as u32, header.viewport_height as u32);
 
     fn par_decode_chunks(
-        chunks: HashMap<u32, BustupChunkDesc>,
+        chunks: Vec<(u32, BustupChunkDesc)>,
         source: &[u8],
-    ) -> impl ParallelIterator<Item = Result<(u32, PictureChunk)>> + '_ {
-        chunks.into_par_iter().map(|(id, desc)| -> Result<_> {
-            let data = &source[desc.offset as usize..(desc.offset + desc.size) as usize];
-            let mut chunk = read_picture_chunk(data)?;
-            cleanup_unused_areas(&mut chunk);
-            Ok((id, chunk))
-        })
+    ) -> Vec<Result<(u32, PictureChunk)>> {
+        chunks.par_chunk_map(
+            shin_tasks::AsyncComputeTaskPool::get(),
+            1,
+            |chunk| -> Result<_> {
+                let &[(id, desc)] = chunk else { unreachable!() };
+                let data = &source[desc.offset as usize..(desc.offset + desc.size) as usize];
+                let mut chunk = read_picture_chunk(data)?;
+                cleanup_unused_areas(&mut chunk);
+                Ok((id, chunk))
+            },
+        )
     }
 
-    par_decode_chunks(base_chunks, source.get_ref()).try_for_each(|res| -> Result<()> {
-        let (_, chunk) = res?;
+    par_decode_chunks(base_chunks, source.get_ref())
+        .into_iter()
+        .try_for_each(|res| -> Result<()> {
+            let (_, chunk) = res?;
 
-        let mut base_image = base_image.lock().unwrap();
-
-        image::imageops::overlay(
-            base_image.deref_mut(),
-            &chunk.data,
-            chunk.offset_x as i64,
-            chunk.offset_y as i64,
-        );
-        Ok(())
-    })?;
+            image::imageops::overlay(
+                &mut base_image,
+                &chunk.data,
+                chunk.offset_x as i64,
+                chunk.offset_y as i64,
+            );
+            Ok(())
+        })?;
 
     let additional_chunks = par_decode_chunks(additional_chunks, source.get_ref())
+        .into_iter()
         .collect::<Result<HashMap<_, _>>>()?;
 
     Ok(Bustup {
-        base_image: base_image.into_inner().unwrap(),
+        base_image,
         origin: (header.origin_x, header.origin_y),
         expressions: header
             .expressions

shin-core/src/format/picture/mod.rs

@@ -7,6 +7,7 @@ use bitflags::bitflags;
 use bytemuck::{Pod, Zeroable};
 use image::{ImageBuffer, RgbaImage};
 use itertools::Itertools;
+use shin_tasks::ParallelSlice;
 use std::borrow::Cow;
 use std::io;
 use std::sync::Mutex;
@@ -462,8 +463,6 @@ pub fn read_picture<'a, B: PictureBuilder<'a>>(
         chunks.push(((chunk_desc.x as usize, chunk_desc.y as usize), chunk_data));
     }
 
-    use rayon::prelude::*;
-
     let builder = B::new(
         builder_args,
         header.effective_width as u32,
@@ -477,9 +476,13 @@ pub fn read_picture<'a, B: PictureBuilder<'a>>(
     // ideally we want to be generic over the parallelization strategy
     let builder = Mutex::new(builder);
     chunks
-        .par_iter()
-        .cloned()
-        .map(|(pos, data)| (pos, read_picture_chunk(data)))
+        .par_chunk_map(shin_tasks::AsyncComputeTaskPool::get(), 1, |chunk| {
+            let &[(pos, data)] = chunk else {
+                unreachable!()
+            };
+            (pos, read_picture_chunk(data))
+        })
+        .into_iter()
         .try_for_each(|(pos, chunk)| {
             builder
                 .lock()

shin-core/src/format/texture_archive/mod.rs

@@ -3,7 +3,7 @@
 use anyhow::Result;
 use binrw::{BinRead, BinWrite};
 use image::RgbaImage;
-use rayon::prelude::*;
+use shin_tasks::ParallelSlice;
 use std::collections::HashMap;
 
 use crate::format::text::ZeroString;
@@ -83,8 +83,8 @@ pub fn read_texture_archive(source: &[u8]) -> Result<TextureArchive> {
 
     let textures = header
         .index
-        .par_iter()
-        .map(|v| {
+        .par_chunk_map(shin_tasks::AsyncComputeTaskPool::get(), 1, |chunk| {
+            let [v] = chunk else { unreachable!() };
             let size = if v.data_compressed_size != 0 {
                 v.data_compressed_size
             } else {
@@ -96,6 +96,7 @@ pub fn read_texture_archive(source: &[u8]) -> Result<TextureArchive> {
                 header.use_dict_encoding != 0,
             )
         })
+        .into_iter()
         .collect::<Result<Vec<_>>>()?;
 
     let name_to_index = header

shin-core/src/lib.rs

@@ -6,6 +6,9 @@
 
 extern crate self as shin_core;
 
+// re-export for convenience
+pub use shin_tasks::create_task_pools;
+
 pub mod format;
 pub mod layout;
 pub mod time;

shin/Cargo.toml

@@ -40,7 +40,7 @@ glam = { workspace = true, features = [
     "bytemuck", "scalar-math"
 ] }
 bytemuck = { version = "1.12.3", features = ["derive"] }
-image = { workspace = true, default-features = false, features = ["png"] }
+image = { workspace = true, default-features = false }
 
 # Theese part of bevy does not depend on ECS or the reflection, so it's not a big problem to use them
 bevy_utils = { workspace = true }

shin/src/window.rs

@@ -32,7 +32,6 @@ use shin_render::GpuCommonResources;
 use shin_render::Pillarbox;
 use shin_render::Pipelines;
 use shin_render::{RenderTarget, Renderable};
-use shin_tasks::create_task_pools;
 
 #[cfg(target_arch = "wasm32")]
 use wasm_bindgen::prelude::*;
@@ -309,7 +308,7 @@ pub async fn run(cli: Cli) {
         }
     }
 
-    create_task_pools();
+    shin_tasks::create_task_pools();
 
     let event_loop = EventLoop::new();
     let window = WindowBuilder::new()