Add filters and reduce to one generation function

Author: SteveCookTU

src/filter.rs

@@ -0,0 +1,131 @@
+use crate::{Advance, Pokemon};
+
+#[derive(Default)]
+pub struct Filter {
+    pub shiny: bool,
+    pub species: Option<u16>,
+    pub min_ivs: [u8; 6],
+    pub max_ivs: [u8; 6],
+    pub ability: Option<u8>,
+    pub nature: Option<u8>,
+    pub item: Option<u16>,
+    pub egg_move: Option<u16>,
+    pub gender: Option<u8>,
+}
+
+impl Filter {
+    pub fn shiny(mut self, shiny: bool) -> Filter {
+        self.shiny = shiny;
+        self
+    }
+
+    pub fn species(mut self, species: u16) -> Filter {
+        self.species = Some(species);
+        self
+    }
+
+    pub fn min_ivs(mut self, min_ivs: [u8; 6]) -> Filter {
+        self.min_ivs = min_ivs;
+        self
+    }
+
+    pub fn max_ivs(mut self, max_ivs: [u8; 6]) -> Filter {
+        self.max_ivs = max_ivs;
+        self
+    }
+
+    pub fn ability(mut self, ability: u8) -> Filter {
+        self.ability = Some(ability);
+        self
+    }
+
+    pub fn nature(mut self, nature: u8) -> Filter {
+        self.nature = Some(nature);
+        self
+    }
+
+    pub fn item(mut self, item: u16) -> Filter {
+        self.item = Some(item);
+        self
+    }
+
+    pub fn egg_move(mut self, egg_move: u16) -> Filter {
+        self.egg_move = Some(egg_move);
+        self
+    }
+
+    pub fn gender(mut self, gender: u8) -> Filter {
+        self.gender = Some(gender);
+        self
+    }
+
+    pub fn passes_filter(&self, advance: &Advance) -> bool {
+        for pokemon in advance.regular_pokemon.iter() {
+            if self.check_pokemon(pokemon) {
+                return true;
+            }
+        }
+
+        if let Some(pokemon) = &advance.rare_pokemon {
+            if self.check_pokemon(pokemon) {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    fn check_pokemon(&self, pokemon: &Pokemon) -> bool {
+        if let Some(species) = self.species {
+            if pokemon.species != species {
+                return false;
+            }
+        }
+
+        if self.shiny && !pokemon.shiny {
+            return false;
+        }
+
+        let mut passes_ivs = true;
+        for (i, iv) in pokemon.ivs.iter().enumerate() {
+            if !(self.min_ivs[i]..=self.max_ivs[i]).contains(iv) {
+                passes_ivs = false;
+                break;
+            }
+        }
+
+        if !passes_ivs {
+            return false;
+        }
+
+        if let Some(ability) = self.ability {
+            if pokemon.ability != ability {
+                return false;
+            }
+        }
+
+        if let Some(nature) = self.nature {
+            if pokemon.nature != nature {
+                return false;
+            }
+        }
+
+        if let Some(item) = self.item {
+            if pokemon.item != item {
+                return false;
+            }
+        }
+
+        if pokemon.egg_move != self.egg_move {
+            return false;
+        }
+
+        if let Some(gender) = self.gender {
+            if pokemon.gender != gender {
+                return false;
+            }
+        }
+
+        true
+    }
+}

src/lib.rs

@@ -1,14 +1,14 @@
 mod flag_util;
-mod personal_info;
-mod personal_info_bdsp;
-mod personal_table;
+pub mod personal_info;
+pub mod personal_info_bdsp;
+pub mod personal_table;
 mod resource_util;
-mod run_print;
 mod run_results;
 mod xorshift;
+mod filter;
 
-pub use run_print::run_print;
 pub use run_results::*;
+pub use filter::*;
 
 use crate::personal_info::PersonalInfo;
 use crate::xorshift::XorShift;

src/main.rs

@@ -1,7 +1,8 @@
 #![allow(unused)]
 
-use bdsp_ug_generator::{run_print, RoomType, Version};
+use bdsp_ug_generator::{run_results, Pokemon, RoomType, Version, ABILITIES_EN, GENDER_SYMBOLS, ITEMS_EN, MOVES_EN, NATURES_EN, SPECIES_EN, personal_table, Filter};
 use clap::Parser;
+use std::fmt::Write;
 
 #[derive(Parser)]
 struct Cli {
@@ -13,13 +14,44 @@ struct Cli {
     story_flag: u8,
     #[clap(short = 's', long)]
     shiny_only: bool,
+    #[clap(long, default_value = "0/0/0/0/0/0")]
+    min_ivs: String,
+    #[clap(long, default_value = "31/31/31/31/31/31")]
+    max_ivs: String,
+    #[clap(long)]
+    species: Option<u16>,
+    #[clap(long)]
+    nature: Option<u8>,
+    #[clap(long)]
+    ability: Option<u8>,
+    #[clap(long)]
+    item: Option<u16>,
+    #[clap(long)]
+    egg_move: Option<u16>,
+    #[clap(long)]
+    gender: Option<u8>,
     advances: u32,
     s0: String,
     s1: String,
     s2: String,
     s3: String,
 }
 
+fn write_pokemon(pokemon: &Pokemon, string: &mut String) {
+    let personal_info = personal_table::BDSP.get_form_entry(pokemon.species as usize, 0);
+    let ability = match pokemon.ability {
+        0 => personal_info.get_ability_1(),
+        _ => personal_info.get_ability_2()
+    };
+    writeln!(string, "Species: {}\nPID: {:08X} EC: {:08X} Shiny: {}\nIVs: {:?} Ability: {} Gender: {}\nNature: {} Item: {}{}\n", SPECIES_EN[pokemon.species as usize], pokemon.pid, pokemon.ec, pokemon.shiny, pokemon.ivs, ABILITIES_EN[ability],  GENDER_SYMBOLS[pokemon.gender as usize], NATURES_EN[pokemon.nature as usize].trim(),
+             ITEMS_EN[pokemon.item as usize].trim(),
+             if let Some(no) = pokemon.egg_move {
+                 format!(" Egg Move: {}", MOVES_EN[no as usize].trim())
+             } else {
+                 "".to_string()
+             }).unwrap();
+}
+
 fn main() {
     let cli: Cli = Cli::parse();
 
@@ -38,7 +70,38 @@ fn main() {
     println!("s3: {:#08X}", s3);
     println!();
 
-    run_print(
+    let min_split = cli.min_ivs.split('/');
+    let max_split = cli.max_ivs.split('/');
+
+    let mut min_ivs = [0, 0, 0, 0, 0, 0];
+
+    for (i, val) in min_split.take(6).enumerate() {
+        if !val.is_empty() {
+            min_ivs[i] = val.parse::<u8>().expect(&format!("Failed to parse min iv {}", i));
+        }
+    }
+
+    let mut max_ivs = [31, 31, 31, 31, 31, 31];
+
+    for (i, val) in max_split.take(6).enumerate() {
+        if !val.is_empty() {
+            max_ivs[i] = val.parse::<u8>().expect(&format!("Failed to parse max iv {}", i));
+        }
+    }
+
+    let filter = Filter {
+        shiny: cli.shiny_only,
+        species: cli.species,
+        min_ivs,
+        max_ivs,
+        ability: cli.ability,
+        nature: cli.nature,
+        item: cli.item,
+        egg_move: cli.egg_move,
+        gender: cli.gender
+    };
+
+    let results = run_results(
         cli.advances,
         s0,
         s1,
@@ -47,6 +110,26 @@ fn main() {
         cli.version,
         cli.story_flag,
         cli.room,
-        cli.shiny_only,
+        filter
     );
+
+    let mut print = String::new();
+
+    for (advance, result) in results.iter().enumerate() {
+        writeln!(
+            print,
+            "-------------------------------------------\nAdvances: {}",
+            advance
+        )
+        .unwrap();
+        for pokemon in result.regular_pokemon.iter() {
+            write_pokemon(pokemon, &mut print);
+        }
+
+        if let Some(pokemon) = &result.rare_pokemon {
+            write_pokemon(pokemon, &mut print);
+        }
+    }
+
+    println!("{}", print);
 }