From 734a0a67b1ec6445f21801fbef4af6ad3bedb49b Mon Sep 17 00:00:00 2001 From: Chaoscaot Date: Sat, 11 Mar 2023 11:34:40 +0100 Subject: [PATCH] Add Multithreading --- schemsearch-cli/Cargo.toml | 3 +- schemsearch-cli/src/main.rs | 148 +++++++++++++++++++++-------------- schemsearch-cli/src/types.rs | 14 ++-- schemsearch-files/src/lib.rs | 6 +- schemsearch-java/src/lib.rs | 6 +- 5 files changed, 103 insertions(+), 74 deletions(-) diff --git a/schemsearch-cli/Cargo.toml b/schemsearch-cli/Cargo.toml index e75738c..f328a1b 100644 --- a/schemsearch-cli/Cargo.toml +++ b/schemsearch-cli/Cargo.toml @@ -14,6 +14,7 @@ schemsearch-sql = { path = "../schemsearch-sql", optional = true } clap = { version = "4.1.8", features = ["cargo"] } futures = { version = "0.3", optional = true } sqlx = { version = "0.6", features = [ "runtime-async-std-native-tls" , "mysql" ], optional = true } +rayon = "1.7.0" [features] -sql = [ "dep:schemsearch-sql", "dep:futures", "dep:sqlx" ] +sql = ["dep:schemsearch-sql", "dep:futures", "dep:sqlx"] diff --git a/schemsearch-cli/src/main.rs b/schemsearch-cli/src/main.rs index 77f120f..a7d3f98 100644 --- a/schemsearch-cli/src/main.rs +++ b/schemsearch-cli/src/main.rs @@ -18,15 +18,18 @@ mod types; use std::fs::File; -use std::io::{BufWriter, StdoutLock, Write}; -use clap::{command, Arg, ArgAction, Command, ValueHint}; +use std::io; +use std::io::{BufWriter, Write}; +use clap::{command, Arg, ArgAction, ValueHint}; use schemsearch_files::Schematic; -use std::path::Path; +use std::path::PathBuf; use clap::error::ErrorKind; use schemsearch_lib::{search, SearchBehavior}; use crate::types::{PathSchematicSupplier, SchematicSupplierType}; #[cfg(feature = "sql")] use futures::executor::block_on; +use rayon::prelude::*; +use rayon::ThreadPoolBuilder; use schemsearch_sql::filter::SchematicFilter; #[cfg(feature = "sql")] use schemsearch_sql::load_all_schematics; @@ -110,6 +113,15 @@ fn main() { .default_value("0.9") .value_parser(|s: &str| s.parse::().map_err(|e| e.to_string())), ) + .arg( + Arg::new("threads") + .help("The number of threads to use [0 = Available Threads]") + .short('T') + .long("threads") + .action(ArgAction::Set) + .default_value("0") + .value_parser(|s: &str| s.parse::().map_err(|e| e.to_string())), + ) .about("Searches for a pattern in a schematic") .bin_name("schemsearch"); @@ -155,16 +167,35 @@ fn main() { threshold: *matches.get_one::("threshold").expect("Couldn't get threshold"), }; - let pattern = match Schematic::load(Path::new(matches.get_one::("pattern").unwrap())) { + let pattern = match Schematic::load(&PathBuf::from(matches.get_one::("pattern").unwrap())) { Ok(x) => x, Err(e) => { cmd.error(ErrorKind::Io, format!("Error while loading Pattern: {}", e.to_string())).exit(); } }; - let mut schematics: Vec = match matches.get_many::("schematic") { - None => vec![], - Some(x) => x.map(|x| SchematicSupplierType::PATH(Box::new(PathSchematicSupplier{path: Path::new(x) }))).collect() + let mut schematics: Vec = Vec::new(); + match matches.get_many::("schematic") { + None => {}, + Some(x) => { + let paths = x.map(|x| PathBuf::from(x)); + for path in paths { + if path.is_dir() { + path.read_dir() + .expect("Couldn't read directory") + .filter_map(|x| x.ok()) + .filter(|x| x.path().is_file()) + .filter(|x| x.path().extension().unwrap().to_str().unwrap() == "schem") + .for_each(|x| { + schematics.push(SchematicSupplierType::PATH(Box::new(PathSchematicSupplier { + path: x.path(), + }))) + }); + } else if path.extension().unwrap().to_str().unwrap() == "schem" { + schematics.push(SchematicSupplierType::PATH(Box::new(PathSchematicSupplier { path }))); + } + } + } }; #[cfg(feature = "sql")] @@ -201,10 +232,6 @@ fn main() { _ => {} } }; - - let stdout = std::io::stdout(); - let mut lock = stdout.lock(); - let file: Option; let mut file_out: Option> = None; @@ -224,81 +251,82 @@ fn main() { }; file_out = Some(BufWriter::new(file.unwrap())); } + ThreadPoolBuilder::new().num_threads(*matches.get_one::("threads").expect("Could not get threads")).build_global().unwrap(); - for schem in schematics { + let matches: Vec = schematics.par_iter().map(|schem| { match schem { SchematicSupplierType::PATH(schem) => { - let path = schem.path; - if path.is_dir() { - match path.read_dir() { - Ok(x) => { - for path in x { - match path { - Ok(x) => { - if x.path().extension().unwrap_or_default() == "schem" { - let schematic = load_schem(&mut cmd, &x.path()); - search_schempath(search_behavior, &pattern, &mut output_std, &mut output_std_csv, &mut output_file_csv, &mut output_file, &mut lock, &mut file_out, schematic, x.path().file_name().unwrap().to_str().unwrap().to_string()); - } - } - Err(e) => cmd.error(ErrorKind::Io, format!("Error while reading schem: {}", e.to_string())).exit() - } - } - } - Err(e) => cmd.error(ErrorKind::Io, format!("Expected to be a dir: {}", e.to_string())).exit() + let schematic = match load_schem(&schem.path) { + Some(x) => x, + None => return Result { + name: schem.get_name(), + matches: vec![] } - } else { - let schematic = load_schem(&mut cmd, &path); - search_schempath(search_behavior, &pattern, &mut output_std, &mut output_std_csv, &mut output_file_csv, &mut output_file, &mut lock, &mut file_out, schematic, schem.get_name()); + }; + Result { + name: schem.get_name(), + matches: search(schematic, &pattern, search_behavior) } } #[cfg(feature = "sql")] SchematicSupplierType::SQL(schem) => { match schem.get_schematic() { Ok(schematic) => { - search_schempath(search_behavior, &pattern, &mut output_std, &mut output_std_csv, &mut output_file_csv, &mut output_file, &mut lock, &mut file_out, schematic, schem.get_name()); + Result { + name: schem.get_name(), + matches: search(schematic, &pattern, search_behavior) + } } Err(e) => { if !output_std && !output_std_csv { println!("Error while loading schematic ({}): {}", schem.get_name(), e.to_string()); } + Result { + name: schem.get_name(), + matches: vec![] + } } } } } + }).collect(); + + let stdout = io::stdout(); + let mut lock = stdout.lock(); + + for matching in matches { + let schem_name = matching.name; + let matching = matching.matches; + for x in matching { + if output_std { + writeln!(lock, "Found match in '{}' at x: {}, y: {}, z: {}, % = {}", schem_name, x.0, x.1, x.2, x.3).unwrap(); + } + if output_std_csv { + writeln!(lock, "{},{},{},{},{}", schem_name, x.0, x.1, x.2, x.3).unwrap(); + } + if output_file { + writeln!(file_out.as_mut().unwrap(), "Found match in '{}' at x: {}, y: {}, z: {}, % = {}", schem_name, x.0, x.1, x.2, x.3).unwrap(); + } + if output_file_csv { + writeln!(file_out.as_mut().unwrap(), "{},{},{},{},{}", schem_name, x.0, x.1, x.2, x.3).unwrap(); + } + } } } -fn load_schem(cmd: &mut Command, schem_path: &Path) -> Schematic { +fn load_schem(schem_path: &PathBuf) -> Option { match Schematic::load(schem_path) { - Ok(x) => x, + Ok(x) => Some(x), Err(e) => { - cmd.error(ErrorKind::Io, format!("Error while loading Schematic ({}): {}", schem_path.file_name().unwrap().to_str().unwrap(), e.to_string())).exit(); + println!("Error while loading schematic ({}): {}", schem_path.to_str().unwrap(), e.to_string()); + None } } } -fn search_schempath(search_behavior: SearchBehavior, pattern: &Schematic, output_std: &mut bool, output_std_csv: &mut bool, output_file_csv: &mut bool, output_file: &mut bool, stdout: &mut StdoutLock, file_out: &mut Option>, schematic: Schematic, schem_name: String) { - if *output_std { - writeln!(stdout, "Searching in schematic: {}", schem_name).unwrap(); - } - if *output_file { - writeln!(file_out.as_mut().unwrap(), "Searching in schematic: {}", schem_name).unwrap(); - } - - let matches = search(schematic, pattern, search_behavior); - - for x in matches { - if *output_std { - writeln!(stdout, "Found match at x: {}, y: {}, z: {}, % = {}", x.0, x.1, x.2, x.3).unwrap(); - } - if *output_std_csv { - writeln!(stdout, "{},{},{},{},{}", schem_name, x.0, x.1, x.2, x.3).unwrap(); - } - if *output_file { - writeln!(file_out.as_mut().unwrap(), "Found match at x: {}, y: {}, z: {}, % = {}", x.0, x.1, x.2, x.3).unwrap(); - } - if *output_file_csv { - writeln!(file_out.as_mut().unwrap(), "{},{},{},{},{}", schem_name, x.0, x.1, x.2, x.3).unwrap(); - } - } +#[derive(Debug, Clone)] +struct Result { + name: String, + matches: Vec<(u16, u16, u16, f32)>, } + diff --git a/schemsearch-cli/src/types.rs b/schemsearch-cli/src/types.rs index 6760184..c2127c7 100644 --- a/schemsearch-cli/src/types.rs +++ b/schemsearch-cli/src/types.rs @@ -15,26 +15,26 @@ * along with this program. If not, see . */ -use std::path::Path; +use std::path::PathBuf; #[cfg(feature = "sql")] use futures::executor::block_on; use schemsearch_files::Schematic; #[cfg(feature = "sql")] use schemsearch_sql::{load_schemdata, SchematicNode}; -pub enum SchematicSupplierType<'local> { - PATH(Box>), +pub enum SchematicSupplierType { + PATH(Box), #[cfg(feature = "sql")] SQL(SqlSchematicSupplier), } -pub struct PathSchematicSupplier<'local> { - pub path: &'local Path, +pub struct PathSchematicSupplier { + pub path: PathBuf, } -impl PathSchematicSupplier<'_> { +impl PathSchematicSupplier { pub fn get_name(&self) -> String { - self.path.file_name().unwrap().to_str().unwrap().to_string() + self.path.file_stem().unwrap().to_str().unwrap().to_string() } } diff --git a/schemsearch-files/src/lib.rs b/schemsearch-files/src/lib.rs index fb1a662..b900f74 100644 --- a/schemsearch-files/src/lib.rs +++ b/schemsearch-files/src/lib.rs @@ -16,7 +16,7 @@ */ use std::io::Read; -use std::path::Path; +use std::path::PathBuf; use nbt::{Map, Value}; use serde::{Deserialize, Deserializer, Serialize}; @@ -81,10 +81,10 @@ impl Schematic { Ok(schematic) } - pub fn load(path: &Path) -> Result { + pub fn load(path: &PathBuf) -> Result { let file = match std::fs::File::open(path) { Ok(x) => x, - Err(_) => return Err(format!("Failed to open file: {}", path.display())) + Err(_) => return Err(format!("Failed to open file: {}", path.to_str().unwrap())) }; Schematic::load_data(file) } diff --git a/schemsearch-java/src/lib.rs b/schemsearch-java/src/lib.rs index f1db3e9..4323c55 100644 --- a/schemsearch-java/src/lib.rs +++ b/schemsearch-java/src/lib.rs @@ -15,7 +15,7 @@ * along with this program. If not, see . */ -use std::path::Path; +use std::path::PathBuf; use jni::JNIEnv; use jni::objects::{JClass, JString}; @@ -32,8 +32,8 @@ pub extern "system" fn Java_SchemSearch_search<'local>(mut env: JNIEnv<'local>, pattern_path: JString<'local>) -> jstring { let schematic_path: String = env.get_string(&schematic_path).expect("Couldn't get java string!").into(); let pattern_path: String = env.get_string(&pattern_path).expect("Couldn't get java string!").into(); - let schematic = Schematic::load(Path::new(&schematic_path)).unwrap(); - let pattern = Schematic::load(Path::new(&pattern_path)).unwrap(); + let schematic = Schematic::load(&PathBuf::from(&schematic_path)).unwrap(); + let pattern = Schematic::load(&PathBuf::from(&pattern_path)).unwrap(); let matches = search(schematic, &pattern, SearchBehavior { ignore_block_data: true,