From ac4145ee40bcc86bd3f6159f5735ed0db6532d7e Mon Sep 17 00:00:00 2001 From: prescientmoon Date: Mon, 9 Sep 2024 18:06:07 +0200 Subject: [PATCH] Revampt jacket loading system --- .gitignore | 2 +- Cargo.lock | 215 +++++++++++++++++++++++++++-- Cargo.toml | 4 + flake.nix | 14 ++ scripts/copy-chart-info.sh | 13 ++ shimmering/config/shorthands.csv | 12 ++ src/arcaea/achievement.rs | 6 +- src/arcaea/chart.rs | 38 ++++- src/arcaea/jacket.rs | 114 ++++++--------- src/assets.rs | 1 + src/bitmap.rs | 3 +- src/cli/mod.rs | 15 ++ src/cli/prepare_jackets.rs | 155 +++++++++++++++++++++ src/commands/chart.rs | 77 ++++++++--- src/commands/discord.rs | 4 +- src/commands/score.rs | 48 ++----- src/commands/stats.rs | 8 +- src/context.rs | 80 ++++++++--- src/main.rs | 135 ++++++++++-------- src/recognition/fuzzy_song_name.rs | 25 ++-- src/recognition/hyperglass.rs | 18 ++- src/recognition/recognize.rs | 14 +- src/recognition/ui.rs | 3 +- src/user.rs | 7 +- 24 files changed, 750 insertions(+), 261 deletions(-) create mode 100755 scripts/copy-chart-info.sh create mode 100644 src/cli/mod.rs create mode 100644 src/cli/prepare_jackets.rs diff --git a/.gitignore b/.gitignore index c338dc9..fc99c01 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,7 @@ shimmering/data shimmering/logs shimmering/assets/fonts -shimmering/assets/songs +shimmering/assets/songs* shimmering/assets/b30_background.* target diff --git a/Cargo.lock b/Cargo.lock index 11bb81b..53c275f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -76,10 +76,59 @@ dependencies = [ ] [[package]] -name = "anyhow" -version = "1.0.86" +name = "anstream" +version = "0.6.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" +checksum = "64e15c1ab1f89faffbf04a634d5e1962e9074f2741eef6d97f3c4e322426d526" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bec1de6f59aedf83baf9ff929c98f2ad654b97c9510f4e70cf6f661d49fd5b1" + +[[package]] +name = "anstyle-parse" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb47de1e80c2b463c735db5b217a0ddc39d612e7ac9e2e96a5aed1f57616c1cb" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d36fc52c7f6c869915e99412912f22093507da8d9e942ceaf66fe4b7c14422a" +dependencies = [ + "windows-sys 0.52.0", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5bf74e1b6e971609db8ca7a9ce79fd5768ab6ae46441c572e46cf596f59e57f8" +dependencies = [ + "anstyle", + "windows-sys 0.52.0", +] + +[[package]] +name = "anyhow" +version = "1.0.87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10f00e1f6e58a40e807377c75c6a7f97bf9044fab57816f2414e6f5f4499d7b8" [[package]] name = "approx" @@ -331,12 +380,64 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "clap" +version = "4.5.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e5a21b8495e732f1b3c364c9949b201ca7bae518c502c80256c96ad79eaf6ac" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.5.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cf2dd12af7a047ad9d6da2b6b249759a22a7abc0f474c1dae1777afa4b21a73" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.5.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "501d359d5f3dcaf6ecdeee48833ae73ec6e42723a1e52419c79abf9507eec0a0" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.66", +] + +[[package]] +name = "clap_lex" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1462739cb27611015575c0c11df5df7601141071f07518d56fcc1be504cbec97" + +[[package]] +name = "cobs" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67ba02a97a2bd10f4b59b25c7973101c79642302776489e030cd13cdab09ed15" + [[package]] name = "color_quant" version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b" +[[package]] +name = "colorchoice" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fd119d74b830634cea2a0f58bbd0d54540518a14397557951e79340abc28c0" + [[package]] name = "core-foundation" version = "0.9.4" @@ -509,7 +610,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" dependencies = [ "cfg-if", - "hashbrown", + "hashbrown 0.14.5", "lock_api", "once_cell", "parking_lot_core", @@ -601,6 +702,18 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3dca9240753cf90908d7e4aac30f630662b02aebaa1b58a3cadabdb23385b58b" +[[package]] +name = "embedded-io" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef1a6892d9eef45c8fa6b9e0086428a2cca8491aca8f787c534a3d6d0bcb3ced" + +[[package]] +name = "embedded-io" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d" + [[package]] name = "encoding_rs" version = "0.8.34" @@ -945,7 +1058,7 @@ dependencies = [ "futures-sink", "futures-util", "http 0.2.12", - "indexmap", + "indexmap 2.2.6", "slab", "tokio", "tokio-util", @@ -962,6 +1075,12 @@ dependencies = [ "crunchy", ] +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" + [[package]] name = "hashbrown" version = "0.14.5" @@ -977,7 +1096,7 @@ version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af" dependencies = [ - "hashbrown", + "hashbrown 0.14.5", ] [[package]] @@ -992,6 +1111,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "http" version = "0.2.12" @@ -1214,6 +1339,17 @@ dependencies = [ "quote", ] +[[package]] +name = "indexmap" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" +dependencies = [ + "autocfg", + "hashbrown 0.12.3", + "serde", +] + [[package]] name = "indexmap" version = "2.2.6" @@ -1221,7 +1357,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.14.5", + "serde", ] [[package]] @@ -1241,6 +1378,12 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f518f335dce6725a761382244631d86cf0ccb2863413590b31338feb467f9c3" +[[package]] +name = "is_terminal_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" + [[package]] name = "itertools" version = "0.12.1" @@ -1316,7 +1459,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.52.6", ] [[package]] @@ -1791,6 +1934,18 @@ dependencies = [ "syn 2.0.66", ] +[[package]] +name = "postcard" +version = "1.0.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f7f0a8d620d71c457dd1d47df76bb18960378da56af4527aaa10f515eee732e" +dependencies = [ + "cobs", + "embedded-io 0.4.0", + "embedded-io 0.6.1", + "serde", +] + [[package]] name = "powerfmt" version = "0.2.0" @@ -2366,6 +2521,36 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "3.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cecfa94848272156ea67b2b1a53f20fc7bc638c4a46d2f8abde08f05f4b857" +dependencies = [ + "base64 0.22.1", + "chrono", + "hex", + "indexmap 1.9.3", + "indexmap 2.2.6", + "serde", + "serde_derive", + "serde_json", + "serde_with_macros", + "time", +] + +[[package]] +name = "serde_with_macros" +version = "3.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8fee4991ef4f274617a51ad4af30519438dacb2f56ac773b08a1922ff743350" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.66", +] + [[package]] name = "serenity" version = "0.12.2" @@ -2414,7 +2599,9 @@ dependencies = [ name = "shimmeringmoon" version = "0.1.0" dependencies = [ + "anyhow", "chrono", + "clap", "freetype-rs", "hypertesseract", "image 0.25.2", @@ -2423,11 +2610,13 @@ dependencies = [ "num", "plotters", "poise", + "postcard", "r2d2", "r2d2_sqlite", "rusqlite", "rusqlite_migration", "serde", + "serde_with", "tempfile", "tokio", "toml", @@ -2815,7 +3004,7 @@ version = "0.22.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "583c44c02ad26b0c3f3066fe629275e50627026c51ac2e595cca4c230ce1ce1d" dependencies = [ - "indexmap", + "indexmap 2.2.6", "serde", "serde_spanned", "toml_datetime", @@ -2925,7 +3114,7 @@ checksum = "eb704842c709bc76f63e99e704cb208beeccca2abbabd0d9aec02e48ca1cee0f" dependencies = [ "chrono", "dashmap", - "hashbrown", + "hashbrown 0.14.5", "mini-moka", "parking_lot", "secrecy", @@ -3000,6 +3189,12 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.10.0" diff --git a/Cargo.toml b/Cargo.toml index df25aa4..84408a3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,10 @@ include_dir = "0.7.4" serde = "1.0.209" toml = "0.8.19" tempfile = "3.12.0" +clap = { version = "4.5.17", features = ["derive"] } +postcard = { version = "1.0.10", features = ["use-std"], default-features = false } +serde_with = "3.9.0" +anyhow = "1.0.87" [profile.dev.package."*"] opt-level = 3 diff --git a/flake.nix b/flake.nix index 9b80ee3..a2d34a2 100644 --- a/flake.nix +++ b/flake.nix @@ -15,6 +15,20 @@ inherit (pkgs) lib; in { + packages.shimmeringmoon = pkgs.rustPlatform.buildRustPackage { + pname = "shimmeringmoon"; + version = "unstable-2024-09-06"; + + src = lib.cleanSource ./.; + + cargoLock = { + lockFile = ./Cargo.lock; + outputHashes = { + "hypertesseract-0.1.0" = "sha256-G0dos5yvvcfBKznAo1IIzLgXqRDxmyZwB93QQ6hVZSo="; + "plotters-0.4.0" = "sha256-9wtd7lig1vQ2RJVaEHdicfPZy2AyuoNav8shPMZ1EuE="; + }; + }; + }; devShell = pkgs.mkShell rec { packages = with pkgs; [ (fenix.complete.withComponents [ diff --git a/scripts/copy-chart-info.sh b/scripts/copy-chart-info.sh new file mode 100755 index 0000000..8ceed3d --- /dev/null +++ b/scripts/copy-chart-info.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash + +if [ "$#" != 2 ]; then + echo "Usage: $0 " + echo "This script copies the chart/song data from a db to another. Useful for creating new dbs for testing." + exit 1 +fi + +from="$1/db.sqlite" +to ="$2/db.sqlite" + +sqlite3 $from ".dump songs" | sqlite3 $to +sqlite3 $from ".dump charts" | sqlite3 $to diff --git a/shimmering/config/shorthands.csv b/shimmering/config/shorthands.csv index 183b6c9..9ebcded 100644 --- a/shimmering/config/shorthands.csv +++ b/shimmering/config/shorthands.csv @@ -47,3 +47,15 @@ Mistempered Malignance,,,mismal Twilight Concerto,,,tasogare Heart,,,kokoro Dancin' on a Cat's Paw,,,nekonote +Bookmaker (2D Version),,,bookmaker +Dement ~after legend~,,,dement +Einherjar Joker,,,einherjar +GOODTEK (Arcaea Edit),,,goodtek +Kanagawa Cyber Culvert,,,kanagawa +La'qryma of the Wasteland,,,laqryma +PRAGMATISM -RESURRECTION-,,,pragmatism +qualia -ideaesthesia-,,,qualia +Shades of Light in a Transcendent Realm,,,shadesoflight +trappola bewitching,,,trappola +Vicious [ANTi] Heroism,,,viciousheroism +eden,,,edenwacca diff --git a/src/arcaea/achievement.rs b/src/arcaea/achievement.rs index b9eac88..d5b5f76 100644 --- a/src/arcaea/achievement.rs +++ b/src/arcaea/achievement.rs @@ -1,3 +1,4 @@ +use anyhow::anyhow; use image::RgbaImage; use crate::{ @@ -121,7 +122,8 @@ impl GoalStats { user: &User, scoring_system: ScoringSystem, ) -> Result { - let plays = get_best_plays(ctx, user.id, scoring_system, 0, usize::MAX, None)??; + let plays = get_best_plays(ctx, user.id, scoring_system, 0, usize::MAX, None)? + .map_err(|s| anyhow!("{s}"))?; let conn = ctx.db.get()?; // {{{ PM count @@ -157,7 +159,7 @@ impl GoalStats { ), |row| row.get(0), ) - .map_err(|_| "No ptt history data found")?; + .map_err(|_| anyhow!("No ptt history data found"))?; // }}} // {{{ Peak PM relay let peak_pm_relay = { diff --git a/src/arcaea/chart.rs b/src/arcaea/chart.rs index 202f8fd..b641fa8 100644 --- a/src/arcaea/chart.rs +++ b/src/arcaea/chart.rs @@ -1,5 +1,6 @@ use std::{fmt::Display, num::NonZeroU16, path::PathBuf}; +use anyhow::anyhow; use image::{ImageBuffer, Rgb}; use rusqlite::types::{FromSql, FromSqlError, FromSqlResult, ValueRef}; @@ -237,6 +238,14 @@ impl CachedSong { chart_ids: [None; 5], } } + + #[inline] + pub fn charts(&self) -> impl Iterator { + self.chart_ids + .into_iter() + .filter_map(|i| i) + .map(|i| i.get() as u32) + } } // }}} // {{{ Song cache @@ -252,7 +261,7 @@ impl SongCache { self.songs .get(id as usize) .and_then(|i| i.as_ref()) - .ok_or_else(|| format!("Could not find song with id {}", id).into()) + .ok_or_else(|| anyhow!("Could not find song with id {}", id)) } #[inline] @@ -261,7 +270,7 @@ impl SongCache { .charts .get(chart_id as usize) .and_then(|i| i.as_ref()) - .ok_or_else(|| format!("Could not find chart with id {}", chart_id))?; + .ok_or_else(|| anyhow!("Could not find chart with id {}", chart_id))?; let song = &self.lookup_song(chart.song_id)?.song; Ok((song, chart)) @@ -272,7 +281,7 @@ impl SongCache { self.songs .get_mut(id as usize) .and_then(|i| i.as_mut()) - .ok_or_else(|| format!("Could not find song with id {}", id).into()) + .ok_or_else(|| anyhow!("Could not find song with id {}", id)) } #[inline] @@ -280,7 +289,7 @@ impl SongCache { self.charts .get_mut(chart_id as usize) .and_then(|i| i.as_mut()) - .ok_or_else(|| format!("Could not find chart with id {}", chart_id).into()) + .ok_or_else(|| anyhow!("Could not find chart with id {}", chart_id)) } #[inline] @@ -292,7 +301,7 @@ impl SongCache { let cached_song = self.lookup_song(id)?; let chart_id = cached_song.chart_ids[difficulty.to_index()] .ok_or_else(|| { - format!( + anyhow!( "Cannot find chart {} [{difficulty:?}]", cached_song.song.title ) @@ -302,6 +311,25 @@ impl SongCache { Ok((&cached_song.song, chart)) } + #[inline] + pub fn lookup_by_difficulty_mut( + &mut self, + id: u32, + difficulty: Difficulty, + ) -> Result<&mut Chart, Error> { + let cached_song = self.lookup_song(id)?; + let chart_id = cached_song.chart_ids[difficulty.to_index()] + .ok_or_else(|| { + anyhow!( + "Cannot find chart {} [{difficulty:?}]", + cached_song.song.title + ) + })? + .get() as u32; + let chart = self.lookup_chart_mut(chart_id)?; + Ok(chart) + } + #[inline] pub fn charts(&self) -> impl Iterator { self.charts.iter().filter_map(|i| i.as_ref()) diff --git a/src/arcaea/jacket.rs b/src/arcaea/jacket.rs index 78fa625..744f0e8 100644 --- a/src/arcaea/jacket.rs +++ b/src/arcaea/jacket.rs @@ -1,13 +1,15 @@ -use std::{fs, io::Cursor}; +use std::fs; +use anyhow::Context; use image::{imageops::FilterType, GenericImageView, Rgba}; use num::Integer; +use serde::{Deserialize, Serialize}; +use serde_with::serde_as; use crate::{ arcaea::chart::{Difficulty, Jacket, SongCache}, - assets::{get_asset_dir, should_blur_jacket_art, should_skip_jacket_art}, + assets::{get_asset_dir, should_skip_jacket_art}, context::Error, - recognition::fuzzy_song_name::guess_chart_name, }; /// How many sub-segments to split each side into @@ -15,14 +17,16 @@ pub const SPLIT_FACTOR: u32 = 8; pub const IMAGE_VEC_DIM: usize = (SPLIT_FACTOR * SPLIT_FACTOR * 3) as usize; pub const BITMAP_IMAGE_SIZE: u32 = 174; -#[derive(Debug, Clone)] +#[serde_as] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct ImageVec { + #[serde_as(as = "[_; IMAGE_VEC_DIM]")] pub colors: [f32; IMAGE_VEC_DIM], } impl ImageVec { // {{{ (Image => vector) encoding - fn from_image(image: &impl GenericImageView>) -> Self { + pub fn from_image(image: &impl GenericImageView>) -> Self { let mut colors = [0.0; IMAGE_VEC_DIM]; let chunk_width = image.width() / SPLIT_FACTOR; let chunk_height = image.height() / SPLIT_FACTOR; @@ -101,94 +105,64 @@ impl JacketCache { Vec::new() } else { + let songs_dir = get_asset_dir().join("songs/by_id"); let entries = - fs::read_dir(get_asset_dir().join("songs")).expect("Couldn't read songs directory"); - let mut jacket_vectors = vec![]; + fs::read_dir(songs_dir).with_context(|| "Couldn't read songs directory")?; + let bytes = fs::read(get_asset_dir().join("songs/recognition_matrix")) + .with_context(|| "Could not read jacket recognition matrix")?; + let jacket_vectors = postcard::from_bytes(&bytes) + .with_context(|| "Could not decode jacket recognition matrix")?; for entry in entries { let dir = entry?; let raw_dir_name = dir.file_name(); let dir_name = raw_dir_name.to_str().unwrap(); - for entry in fs::read_dir(dir.path()).expect("Couldn't read song directory") { + let song_id = dir_name.parse().with_context(|| { + format!("Dir name {dir_name} could not be parsed as `u32` song id") + })?; + + let entries = + fs::read_dir(dir.path()).with_context(|| "Couldn't read song directory")?; + for entry in entries { let file = entry?; let raw_name = file.file_name(); let name = raw_name.to_str().unwrap().strip_suffix(".jpg").unwrap(); - if !name.ends_with("_256") { - continue; - } + let difficulty = Difficulty::DIFFICULTY_SHORTHANDS + .iter() + .zip(Difficulty::DIFFICULTIES) + .find_map(|(s, d)| Some(d).filter(|_| name == s.to_lowercase())); - let name = name.strip_suffix("_256").unwrap(); + let contents: &'static _ = fs::read(file.path()) + .with_context(|| "Coult not read prepared jacket image")? + .leak(); - let difficulty = match name { - "0" => Some(Difficulty::PST), - "1" => Some(Difficulty::PRS), - "2" => Some(Difficulty::FTR), - "3" => Some(Difficulty::BYD), - "4" => Some(Difficulty::ETR), - "base" => None, - "base_night" => None, - "base_ja" => None, - _ => Err(format!("Unknown jacket suffix {}", name))?, - }; - - let (song_id, chart_id) = { - let (song, chart) = - guess_chart_name(dir_name, &song_cache, difficulty, true)?; - (song.id, chart.id) - }; - - let contents: &'static _ = fs::read(file.path())?.leak(); - - let image = image::load_from_memory(contents)?; - jacket_vectors.push((song_id, ImageVec::from_image(&image))); - let mut image = - image.resize(BITMAP_IMAGE_SIZE, BITMAP_IMAGE_SIZE, FilterType::Nearest); - - if should_blur_jacket_art() { - image = image.blur(40.0); - } - - let encoded_pic = { - let mut processed_pic = Vec::new(); - image.write_to( - &mut Cursor::new(&mut processed_pic), - image::ImageFormat::Jpeg, - )?; - processed_pic.leak() - }; + let image = image::load_from_memory(contents) + .with_context(|| "Could not load jacket image from prepared bytes")?; let bitmap: &'static _ = Box::leak(Box::new(image.into_rgb8())); - if name == "base" { - // Inefficiently iterates over everything, but it's fine for ~1k entries - for chart in song_cache.charts_mut() { - if chart.song_id == song_id && chart.cached_jacket.is_none() { + if let Some(difficulty) = difficulty { + let chart = song_cache + .lookup_by_difficulty_mut(song_id, difficulty) + .unwrap(); + chart.cached_jacket = Some(Jacket { + raw: contents, + bitmap, + }); + } else { + for chart_id in song_cache.lookup_song(song_id)?.charts() { + let chart = song_cache.lookup_chart_mut(chart_id)?; + if chart.cached_jacket.is_none() { chart.cached_jacket = Some(Jacket { - raw: encoded_pic, + raw: contents, bitmap, }); } } - } else if difficulty.is_some() { - let chart = song_cache.lookup_chart_mut(chart_id).unwrap(); - chart.cached_jacket = Some(Jacket { - raw: encoded_pic, - bitmap, - }); } } } - for chart in song_cache.charts() { - if chart.cached_jacket.is_none() { - println!( - "No jacket found for '{} [{:?}]'", - song_cache.lookup_song(chart.song_id)?.song.title, - chart.difficulty - ) - } - } - jacket_vectors }; diff --git a/src/assets.rs b/src/assets.rs index 7a6c91c..379113f 100644 --- a/src/assets.rs +++ b/src/assets.rs @@ -85,6 +85,7 @@ pub fn should_skip_jacket_art() -> bool { } #[inline] +#[allow(dead_code)] pub fn should_blur_jacket_art() -> bool { var("SHIMMERING_BLUR_JACKETS").unwrap_or_default() == "1" } diff --git a/src/bitmap.rs b/src/bitmap.rs index b2070cf..d082dba 100644 --- a/src/bitmap.rs +++ b/src/bitmap.rs @@ -6,6 +6,7 @@ //! There's still stuff to be implemented here, like a cache for glyphs and //! whatnot, but this does run pretty stably for the b30 renderer. +use anyhow::anyhow; use freetype::{ bitmap::PixelMode, face::{KerningMode, LoadFlag}, @@ -355,7 +356,7 @@ impl BitmapCanvas { Some((i, glyph_index)) }) .ok_or_else(|| { - format!("Could not get glyph index for char '{}' in \"{}\"", c, text) + anyhow!("Could not get glyph index for char '{}' in \"{}\"", c, text) })?; let face = &mut faces[face_index]; diff --git a/src/cli/mod.rs b/src/cli/mod.rs new file mode 100644 index 0000000..a4090bd --- /dev/null +++ b/src/cli/mod.rs @@ -0,0 +1,15 @@ +pub mod prepare_jackets; + +#[derive(clap::Parser)] +#[command(author, version, about, long_about = None)] +pub struct Cli { + #[command(subcommand)] + pub command: Command, +} + +#[derive(clap::Subcommand)] +pub enum Command { + /// Start the discord bot + Discord {}, + PrepareJackets {}, +} diff --git a/src/cli/prepare_jackets.rs b/src/cli/prepare_jackets.rs new file mode 100644 index 0000000..427c5ab --- /dev/null +++ b/src/cli/prepare_jackets.rs @@ -0,0 +1,155 @@ +use std::{ + fs, + io::{stdout, Write}, +}; + +use anyhow::{anyhow, bail, Context}; +use image::imageops::FilterType; + +use crate::{ + arcaea::{ + chart::{Difficulty, SongCache}, + jacket::{ImageVec, BITMAP_IMAGE_SIZE}, + }, + assets::{get_asset_dir, get_data_dir}, + context::{connect_db, Error}, + recognition::fuzzy_song_name::guess_chart_name, +}; + +pub fn prepare_jackets() -> Result<(), Error> { + let db = connect_db(&get_data_dir()); + let song_cache = SongCache::new(&db)?; + + let songs_dir = get_asset_dir().join("songs"); + let raw_songs_dir = songs_dir.join("raw"); + + let by_id_dir = songs_dir.join("by_id"); + if by_id_dir.exists() { + fs::remove_dir_all(&by_id_dir).with_context(|| "Could not remove `by_id` dir")?; + } + fs::create_dir_all(&by_id_dir).with_context(|| "Could not create `by_id` dir")?; + + let mut jacket_vectors = vec![]; + + let entries = fs::read_dir(&raw_songs_dir) + .with_context(|| "Couldn't read songs directory")? + .collect::, _>>() + .with_context(|| format!("Could not read member of `songs/raw`"))?; + + for (i, dir) in entries.iter().enumerate() { + let raw_dir_name = dir.file_name(); + let dir_name = raw_dir_name.to_str().unwrap(); + + // {{{ Update progress live + print!( + "{}/{}: {dir_name} \r", + i, + entries.len() + ); + if i % 5 == 0 { + stdout().flush()?; + } + // }}} + + let entries = fs::read_dir(dir.path()) + .with_context(|| "Couldn't read song directory")? + .map(|f| f.unwrap()) + .filter(|f| f.file_name().to_str().unwrap().ends_with("_256.jpg")) + .collect::>(); + + for file in &entries { + let raw_name = file.file_name(); + let name = raw_name + .to_str() + .unwrap() + .strip_suffix("_256.jpg") + .ok_or_else(|| { + anyhow!("No '_256.jpg' suffix to remove from filename {raw_name:?}") + })?; + + let difficulty = match name { + "0" => Some(Difficulty::PST), + "1" => Some(Difficulty::PRS), + "2" => Some(Difficulty::FTR), + "3" => Some(Difficulty::BYD), + "4" => Some(Difficulty::ETR), + "base" => None, + "base_night" => None, + "base_ja" => None, + _ => bail!("Unknown jacket suffix {}", name), + }; + + // Sometimes it's useful to distinguish between separate (but related) + // charts like "Vicious Heroism" and "Vicious [ANTi] Heroism" being in + // the same directory. To do this, we only allow the base jacket to refer + // to the FUTURE difficulty, unless it's the only jacket present + // (or unless we are parsing the tutorial) + let search_difficulty = + if entries.len() > 1 && difficulty.is_none() && dir_name != "tutorial" { + Some(Difficulty::FTR) + } else { + difficulty + }; + + let (song, _) = guess_chart_name(dir_name, &song_cache, search_difficulty, true) + .with_context(|| format!("Could not recognise chart name from '{dir_name}'"))?; + + // {{{ Set up `out_dir` paths + let out_dir = { + let out = by_id_dir.join(song.id.to_string()); + if !out.exists() { + fs::create_dir_all(&out).with_context(|| { + format!( + "Could not create parent dir for song '{}' inside `by_id`", + song.title + ) + })?; + } + + out + }; + // }}} + + let difficulty_string = if let Some(difficulty) = difficulty { + &Difficulty::DIFFICULTY_SHORTHANDS[difficulty.to_index()].to_lowercase() + } else { + "def" + }; + + let contents: &'static _ = fs::read(file.path()) + .with_context(|| format!("Could not read image for file {:?}", file.path()))? + .leak(); + let image = image::load_from_memory(contents)?; + + jacket_vectors.push((song.id, ImageVec::from_image(&image))); + + let image = image.resize(BITMAP_IMAGE_SIZE, BITMAP_IMAGE_SIZE, FilterType::Gaussian); + let image_out_path = + out_dir.join(format!("{difficulty_string}_{BITMAP_IMAGE_SIZE}.jpg")); + image + .save(&image_out_path) + .with_context(|| format!("Could not save image to {image_out_path:?}"))?; + } + } + + // NOTE: this is N^2, but it's a one-off warning thing, so it's fine + for chart in song_cache.charts() { + if jacket_vectors.iter().all(|(i, _)| chart.song_id != *i) { + println!( + "No jacket found for '{} [{:?}]'", + song_cache.lookup_song(chart.song_id)?.song.title, + chart.difficulty + ) + } + } + + { + println!("Encoded {} images", jacket_vectors.len()); + let bytes = postcard::to_allocvec(&jacket_vectors) + .with_context(|| format!("Coult not encode jacket matrix"))?; + fs::write(songs_dir.join("recognition_matrix"), bytes) + .with_context(|| format!("Could not write jacket matrix"))?; + } + + Ok(()) +} diff --git a/src/commands/chart.rs b/src/commands/chart.rs index 0e3cc6c..662a2e5 100644 --- a/src/commands/chart.rs +++ b/src/commands/chart.rs @@ -1,3 +1,4 @@ +use anyhow::anyhow; use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage}; use crate::{ @@ -25,6 +26,8 @@ use crate::{ user::discord_id_to_discord_user, }; +use super::discord::MessageContext; + // {{{ Top command /// Chart-related stats #[poise::command( @@ -38,15 +41,9 @@ pub async fn chart(_ctx: Context<'_>) -> Result<(), Error> { } // }}} // {{{ Info -/// Show a chart given it's name -#[poise::command(prefix_command, slash_command, user_cooldown = 1)] -async fn info( - ctx: Context<'_>, - #[rest] - #[description = "Name of chart to show (difficulty at the end)"] - name: String, -) -> Result<(), Error> { - let (song, chart) = guess_song_and_chart(&ctx.data(), &name)?; +// {{{ Implementation +async fn info_impl(ctx: &mut impl MessageContext, name: &str) -> Result<(), Error> { + let (song, chart) = guess_song_and_chart(&ctx.data(), name)?; let attachement_name = "chart.png"; let icon_attachement = match chart.cached_jacket.as_ref() { @@ -95,17 +92,62 @@ async fn info( embed = embed.thumbnail(format!("attachment://{}", &attachement_name)); } - ctx.channel_id() - .send_files( - ctx.http(), - icon_attachement, - CreateMessage::new().embed(embed), - ) + ctx.send_files(icon_attachement, CreateMessage::new().embed(embed)) .await?; Ok(()) } // }}} +// {{{ Tests +#[cfg(test)] +mod info_tests { + use crate::{commands::discord::mock::MockContext, with_test_ctx}; + + use super::*; + + #[tokio::test] + async fn no_suffix() -> Result<(), Error> { + with_test_ctx!("test/commands/chart/info/no_suffix", async |ctx| { + info_impl(ctx, "Pentiment").await?; + Ok(()) + }) + } + + #[tokio::test] + async fn specify_difficulty() -> Result<(), Error> { + with_test_ctx!("test/commands/chart/info/specify_difficulty", async |ctx| { + info_impl(ctx, "Hellohell [ETR]").await?; + Ok(()) + }) + } + + #[tokio::test] + async fn last_byd() -> Result<(), Error> { + with_test_ctx!( + "test/commands/chart/info/last_byd", + async |ctx: &mut MockContext| { + info_impl(ctx, "Last | Moment [BYD]").await?; + info_impl(ctx, "Last | Eternity [BYD]").await?; + Ok(()) + } + ) + } +} +// }}} + +/// Show a chart given it's name +#[poise::command(prefix_command, slash_command, user_cooldown = 1)] +async fn info( + mut ctx: Context<'_>, + #[rest] + #[description = "Name of chart to show (difficulty at the end)"] + name: String, +) -> Result<(), Error> { + info_impl(&mut ctx, &name).await?; + + Ok(()) +} +// }}} // {{{ Best score /// Show the best score on a given chart #[poise::command(prefix_command, slash_command, user_cooldown = 1)] @@ -138,9 +180,10 @@ async fn best( )? .query_row((user.id, chart.id), |row| Play::from_sql(chart, row)) .map_err(|_| { - format!( + anyhow!( "Could not find any scores for {} [{:?}]", - song.title, chart.difficulty + song.title, + chart.difficulty ) })?; diff --git a/src/commands/discord.rs b/src/commands/discord.rs index 63c112a..cc75fb5 100644 --- a/src/commands/discord.rs +++ b/src/commands/discord.rs @@ -53,9 +53,7 @@ pub trait MessageContext { } // }}} // {{{ Poise implementation -impl<'a, 'b> MessageContext - for poise::Context<'a, UserContext, Box> -{ +impl<'a> MessageContext for poise::Context<'a, UserContext, Error> { type Attachment = poise::serenity_prelude::Attachment; fn data(&self) -> &UserContext { diff --git a/src/commands/score.rs b/src/commands/score.rs index b569494..0c4eac6 100644 --- a/src/commands/score.rs +++ b/src/commands/score.rs @@ -4,6 +4,7 @@ use crate::context::{Context, Error}; use crate::recognition::recognize::{ImageAnalyzer, ScoreKind}; use crate::user::{discord_id_to_discord_user, User}; use crate::{get_user, timed}; +use anyhow::anyhow; use image::DynamicImage; use poise::serenity_prelude as serenity; use poise::serenity_prelude::CreateMessage; @@ -90,9 +91,10 @@ async fn magic_impl( analyzer .read_score(ctx.data(), Some(chart.note_count), &grayscale_image, kind) .map_err(|err| { - format!( + anyhow!( "Could not read score for chart {} [{:?}]: {err}", - song.title, chart.difficulty + song.title, + chart.difficulty ) })? }); @@ -137,52 +139,24 @@ async fn magic_impl( // {{{ Tests #[cfg(test)] mod magic_tests { - use std::{path::PathBuf, process::Command, str::FromStr}; - use r2d2_sqlite::SqliteConnectionManager; + use std::path::PathBuf; - use crate::{ - commands::discord::mock::MockContext, - context::{connect_db, get_shared_context}, - }; + use crate::with_test_ctx; use super::*; - macro_rules! with_ctx { - ($test_path:expr, $f:expr) => {{ - let mut data = (*get_shared_context().await).clone(); - let dir = tempfile::tempdir()?; - let path = dir.path().join("db.sqlite"); - println!("path {path:?}"); - data.db = connect_db(SqliteConnectionManager::file(path)); - - Command::new("scripts/import-charts.py") - .env("SHIMMERING_DATA_DIR", dir.path().to_str().unwrap()) - .output() - .unwrap(); - - let mut ctx = MockContext::new(data); - User::create_from_context(&ctx)?; - - let res: Result<(), Error> = $f(&mut ctx).await; - res?; - - ctx.write_to(&PathBuf::from_str($test_path)?)?; - Ok(()) - }}; - } - #[tokio::test] async fn no_pics() -> Result<(), Error> { - with_ctx!("test/commands/score/magic/no_pics", async |ctx| { + with_test_ctx!("test/commands/score/magic/no_pics", async |ctx| { magic_impl(ctx, vec![]).await?; Ok(()) }) } #[tokio::test] - async fn basic_pic() -> Result<(), Error> { - with_ctx!("test/commands/score/magic/single_pic", async |ctx| { + async fn simple_pic() -> Result<(), Error> { + with_test_ctx!("test/commands/score/magic/single_pic", async |ctx| { magic_impl( ctx, vec![PathBuf::from_str("test/screenshots/alter_ego.jpg")?], @@ -194,7 +168,7 @@ mod magic_tests { #[tokio::test] async fn weird_kerning() -> Result<(), Error> { - with_ctx!("test/commands/score/magic/weird_kerning", async |ctx| { + with_test_ctx!("test/commands/score/magic/weird_kerning", async |ctx| { magic_impl( ctx, vec![ @@ -298,7 +272,7 @@ pub async fn show( Ok((song, chart, play, discord_id)) })? .next() - .ok_or_else(|| format!("Could not find play with id {}", id))??; + .ok_or_else(|| anyhow!("Could not find play with id {}", id))??; let author = discord_id_to_discord_user(&ctx, &discord_id).await?; let user = User::by_id(ctx.data(), play.user_id)?; diff --git a/src/commands/stats.rs b/src/commands/stats.rs index 6143f00..166fc05 100644 --- a/src/commands/stats.rs +++ b/src/commands/stats.rs @@ -1,5 +1,6 @@ use std::io::Cursor; +use anyhow::anyhow; use image::{DynamicImage, ImageBuffer}; use poise::{ serenity_prelude::{CreateAttachment, CreateEmbed}, @@ -194,9 +195,10 @@ async fn best_plays( // }}} // {{{ Display jacket let jacket = chart.cached_jacket.as_ref().ok_or_else(|| { - format!( + anyhow!( "Cannot find jacket for chart {} [{:?}]", - song.title, chart.difficulty + song.title, + chart.difficulty ) })?; @@ -289,7 +291,7 @@ async fn best_plays( // {{{ Display status text with_font(&EXO_FONT, |faces| { let status = play.short_status(scoring_system, chart).ok_or_else(|| { - format!( + anyhow!( "Could not get status for score {}", play.score(scoring_system) ) diff --git a/src/context.rs b/src/context.rs index e9b9f14..0a77c31 100644 --- a/src/context.rs +++ b/src/context.rs @@ -3,6 +3,7 @@ use r2d2::Pool; use r2d2_sqlite::SqliteConnectionManager; use rusqlite_migration::Migrations; use std::fs; +use std::path::Path; use std::sync::LazyLock; use crate::{ @@ -13,7 +14,7 @@ use crate::{ }; // Types used by all command functions -pub type Error = Box; +pub type Error = anyhow::Error; pub type Context<'a> = poise::Context<'a, UserContext, Error>; pub type DbConnection = r2d2::Pool; @@ -33,21 +34,24 @@ pub struct UserContext { pub kazesawa_bold_measurements: CharMeasurements, } -pub fn connect_db(manager: SqliteConnectionManager) -> DbConnection { +pub fn connect_db(data_dir: &Path) -> DbConnection { timed!("create_sqlite_pool", { - Pool::new(manager.with_init(|conn| { - static MIGRATIONS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/migrations"); - static MIGRATIONS: LazyLock = LazyLock::new(|| { - Migrations::from_directory(&MIGRATIONS_DIR).expect("Could not load migrations") - }); + fs::create_dir_all(data_dir).expect("Could not create $SHIMMERING_DATA_DIR"); - MIGRATIONS - .to_latest(conn) - .expect("Could not run migrations"); + let data_dir = data_dir.to_str().unwrap().to_owned(); - Ok(()) - })) - .expect("Could not open sqlite database.") + let db_path = format!("{}/db.sqlite", data_dir); + let mut conn = rusqlite::Connection::open(&db_path).unwrap(); + static MIGRATIONS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/migrations"); + static MIGRATIONS: LazyLock = LazyLock::new(|| { + Migrations::from_directory(&MIGRATIONS_DIR).expect("Could not load migrations") + }); + + MIGRATIONS + .to_latest(&mut conn) + .expect("Could not run migrations"); + + Pool::new(SqliteConnectionManager::file(&db_path)).expect("Could not open sqlite database.") }) } @@ -55,12 +59,7 @@ impl UserContext { #[inline] pub async fn new() -> Result { timed!("create_context", { - fs::create_dir_all(get_data_dir())?; - - let db = connect_db(SqliteConnectionManager::file(&format!( - "{}/db.sqlite", - get_data_dir().to_str().unwrap() - ))); + let db = connect_db(&get_data_dir()); let mut song_cache = timed!("make_song_cache", { SongCache::new(&db)? }); let jacket_cache = timed!("make_jacket_cache", { JacketCache::new(&mut song_cache)? }); @@ -93,8 +92,45 @@ impl UserContext { } } -pub async fn get_shared_context() -> &'static UserContext { - static CELL: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); - CELL.get_or_init(async || UserContext::new().await.unwrap()) +#[cfg(test)] +pub mod testing { + use super::*; + + pub async fn get_shared_context() -> &'static UserContext { + static CELL: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); + CELL.get_or_init(async || { + // env::set_var("SHIMMERING_DATA_DIR", "") + UserContext::new().await.unwrap() + }) .await + } + + pub fn import_songs_and_jackets_from(to: &Path) -> () { + std::process::Command::new("scripts/copy-chart-info.sh") + .arg(get_data_dir()) + .arg(to) + .output() + .expect("Could not run sh chart info copy script"); + } + + #[macro_export] + macro_rules! with_test_ctx { + ($test_path:expr, $f:expr) => {{ + use std::str::FromStr; + + let mut data = (*crate::context::testing::get_shared_context().await).clone(); + let dir = tempfile::tempdir()?; + data.db = crate::context::connect_db(dir.path()); + crate::context::testing::import_songs_and_jackets_from(dir.path()); + + let mut ctx = crate::commands::discord::mock::MockContext::new(data); + crate::user::User::create_from_context(&ctx)?; + + let res: Result<(), Error> = $f(&mut ctx).await; + res?; + + ctx.write_to(&std::path::PathBuf::from_str($test_path)?)?; + Ok(()) + }}; + } } diff --git a/src/main.rs b/src/main.rs index 9d77321..66eb916 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,10 +7,12 @@ #![feature(thread_local)] #![feature(generic_arg_infer)] #![feature(lazy_cell_consume)] +#![feature(iter_collect_into)] mod arcaea; mod assets; mod bitmap; +mod cli; mod commands; mod context; mod levenshtein; @@ -21,6 +23,8 @@ mod transform; mod user; use arcaea::play::generate_missing_scores; +use clap::Parser; +use cli::{prepare_jackets::prepare_jackets, Cli, Command}; use context::{Error, UserContext}; use poise::serenity_prelude::{self as serenity}; use std::{env::var, sync::Arc, time::Duration}; @@ -39,70 +43,79 @@ async fn on_error(error: poise::FrameworkError<'_, UserContext, Error>) { #[tokio::main] async fn main() { - // {{{ Poise options - let options = poise::FrameworkOptions { - commands: vec![ - commands::help(), - commands::score::score(), - commands::stats::stats(), - commands::chart::chart(), - ], - prefix_options: poise::PrefixFrameworkOptions { - stripped_dynamic_prefix: Some(|_ctx, message, _user_ctx| { - Box::pin(async { - if message.author.bot || Into::::into(message.author.id) == 1 { - Ok(None) - } else if message.content.starts_with("!") { - Ok(Some(message.content.split_at(1))) - } else if message.guild_id.is_none() { - if message.content.trim().len() == 0 { - Ok(Some(("", "score magic"))) - } else { - Ok(Some(("", &message.content[..]))) + let cli = Cli::parse(); + match cli.command { + Command::Discord {} => { + // {{{ Poise options + let options = poise::FrameworkOptions { + commands: vec![ + commands::help(), + commands::score::score(), + commands::stats::stats(), + commands::chart::chart(), + ], + prefix_options: poise::PrefixFrameworkOptions { + stripped_dynamic_prefix: Some(|_ctx, message, _user_ctx| { + Box::pin(async { + if message.author.bot || Into::::into(message.author.id) == 1 { + Ok(None) + } else if message.content.starts_with("!") { + Ok(Some(message.content.split_at(1))) + } else if message.guild_id.is_none() { + if message.content.trim().len() == 0 { + Ok(Some(("", "score magic"))) + } else { + Ok(Some(("", &message.content[..]))) + } + } else { + Ok(None) + } + }) + }), + edit_tracker: Some(Arc::new(poise::EditTracker::for_timespan( + Duration::from_secs(3600), + ))), + ..Default::default() + }, + on_error: |error| Box::pin(on_error(error)), + ..Default::default() + }; + // }}} + // {{{ Start poise + let framework = poise::Framework::builder() + .setup(move |ctx, _ready, framework| { + Box::pin(async move { + println!("Logged in as {}", _ready.user.name); + poise::builtins::register_globally(ctx, &framework.options().commands) + .await?; + let ctx = UserContext::new().await?; + + if var("SHIMMERING_REGEN_SCORES").unwrap_or_default() == "1" { + timed!("generate_missing_scores", { + generate_missing_scores(&ctx).await?; + }); } - } else { - Ok(None) - } + + Ok(ctx) + }) }) - }), - edit_tracker: Some(Arc::new(poise::EditTracker::for_timespan( - Duration::from_secs(3600), - ))), - ..Default::default() - }, - on_error: |error| Box::pin(on_error(error)), - ..Default::default() - }; - // }}} - // {{{ Start poise - let framework = poise::Framework::builder() - .setup(move |ctx, _ready, framework| { - Box::pin(async move { - println!("Logged in as {}", _ready.user.name); - poise::builtins::register_globally(ctx, &framework.options().commands).await?; - let ctx = UserContext::new().await?; + .options(options) + .build(); - if var("SHIMMERING_REGEN_SCORES").unwrap_or_default() == "1" { - timed!("generate_missing_scores", { - generate_missing_scores(&ctx).await?; - }); - } + let token = var("SHIMMERING_DISCORD_TOKEN") + .expect("Missing `SHIMMERING_DISCORD_TOKEN` env var"); + let intents = serenity::GatewayIntents::non_privileged() + | serenity::GatewayIntents::MESSAGE_CONTENT; - Ok(ctx) - }) - }) - .options(options) - .build(); + let client = serenity::ClientBuilder::new(token, intents) + .framework(framework) + .await; - let token = - var("SHIMMERING_DISCORD_TOKEN").expect("Missing `SHIMMERING_DISCORD_TOKEN` env var"); - let intents = - serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::MESSAGE_CONTENT; - - let client = serenity::ClientBuilder::new(token, intents) - .framework(framework) - .await; - - client.unwrap().start().await.unwrap() - // }}} + client.unwrap().start().await.unwrap() + // }}} + } + Command::PrepareJackets {} => { + prepare_jackets().expect("Could not prepare jackets"); + } + } } diff --git a/src/recognition/fuzzy_song_name.rs b/src/recognition/fuzzy_song_name.rs index 32d0286..004b16a 100644 --- a/src/recognition/fuzzy_song_name.rs +++ b/src/recognition/fuzzy_song_name.rs @@ -10,6 +10,8 @@ //! databases extracted from the game, but this is still useful for having a //! "canonical" way to refer to some weirdly-named charts). +use anyhow::bail; + use crate::arcaea::chart::{Chart, Difficulty, Song, SongCache}; use crate::context::{Error, UserContext}; use crate::levenshtein::edit_distance_with; @@ -82,26 +84,29 @@ pub fn guess_chart_name<'a>( let song_title = &song.lowercase_title; distance_vec.clear(); + // Apply raw distance let base_distance = edit_distance_with(&text, &song_title, &mut levenshtein_vec); - if base_distance < 1.max(song.title.len() / 3) { + if base_distance <= song.title.len() / 3 { distance_vec.push(base_distance * 10 + 2); } + // Cut title to the length of the text, and then check let shortest_len = Ord::min(song_title.len(), text.len()); if let Some(sliced) = &song_title.get(..shortest_len) && (text.len() >= 6 || unsafe_heuristics) { let slice_distance = edit_distance_with(&text, sliced, &mut levenshtein_vec); - if slice_distance < 1 { - distance_vec.push(slice_distance * 10 + 3); + if slice_distance == 0 { + distance_vec.push(3); } } + // Shorthand-based matching if let Some(shorthand) = &chart.shorthand && unsafe_heuristics { let short_distance = edit_distance_with(&text, shorthand, &mut levenshtein_vec); - if short_distance < 1.max(shorthand.len() / 3) { + if short_distance <= shorthand.len() / 3 { distance_vec.push(short_distance * 10 + 1); } } @@ -113,12 +118,16 @@ pub fn guess_chart_name<'a>( }) .collect(); + close_enough.sort_by_key(|(song, _, _)| song.id); + close_enough.dedup_by_key(|(song, _, _)| song.id); + if close_enough.len() == 0 { if text.len() <= 1 { - Err(format!( + bail!( "Could not find match for chart name '{}' [{:?}]", - raw_text, difficulty - ))?; + raw_text, + difficulty + ); } else { text = &text[..text.len() - 1]; } @@ -129,7 +138,7 @@ pub fn guess_chart_name<'a>( close_enough.sort_by_key(|(_, _, distance)| *distance); break (close_enough[0].0, close_enough[0].1); } else { - return Err(format!("Name '{}' is too vague to choose a match", raw_text).into()); + bail!("Name '{}' is too vague to choose a match", raw_text); }; }; }; diff --git a/src/recognition/hyperglass.rs b/src/recognition/hyperglass.rs index 85edfe4..387dd73 100644 --- a/src/recognition/hyperglass.rs +++ b/src/recognition/hyperglass.rs @@ -21,6 +21,7 @@ //! aforementioned precomputed vectors are generated using almost the exact //! procedure described in steps 1-6, except the images are generated at //! startup using my very own bitmap rendering module (`crate::bitmap`). +use anyhow::{anyhow, bail}; use freetype::Face; use image::{DynamicImage, ImageBuffer, Luma}; use imageproc::{ @@ -58,7 +59,7 @@ impl ComponentVec { .bounds .get(component as usize - 1) .and_then(|o| o.as_ref()) - .ok_or_else(|| "Missing bounds for given connected component")?; + .ok_or_else(|| anyhow!("Missing bounds for given connected component"))?; for i in 0..(SPLIT_FACTOR * SPLIT_FACTOR) { let (iy, ix) = i.div_rem_euclid(&SPLIT_FACTOR); @@ -82,10 +83,7 @@ impl ComponentVec { let size = (x_end + 1 - x_start) * (y_end + 1 - y_start); if size == 0 { - return Err(format!( - "Got zero size for chunk [{x_start},{x_end}]x[{y_start},{y_end}]" - ) - .into()); + bail!("Got zero size for chunk [{x_start},{x_end}]x[{y_start},{y_end}]"); } chunks[i as usize] = count as f32 / size as f32; @@ -256,7 +254,7 @@ impl CharMeasurements { canvas.text(padding, &mut [face], style, &string)?; let buffer = ImageBuffer::from_raw(canvas.width, canvas.height(), canvas.buffer.to_vec()) - .ok_or_else(|| "Failed to turn buffer into canvas")?; + .ok_or_else(|| anyhow!("Failed to turn buffer into canvas"))?; let image = DynamicImage::ImageRgb8(buffer); debug_image_log(&image); @@ -270,14 +268,14 @@ impl CharMeasurements { .filter_map(|o| o.as_ref()) .map(|b| b.x_max - b.x_min) .max() - .ok_or_else(|| "No connected components found")?; + .ok_or_else(|| anyhow!("No connected components found"))?; let max_height = components .bounds .iter() .filter_map(|o| o.as_ref()) .map(|b| b.y_max - b.y_min) .max() - .ok_or_else(|| "No connected components found")?; + .ok_or_else(|| anyhow!("No connected components found"))?; // }}} let mut chars = Vec::with_capacity(string.len()); @@ -318,7 +316,7 @@ impl CharMeasurements { .filter_map(|o| o.as_ref()) .map(|b| b.y_max - b.y_min) .max() - .ok_or_else(|| "No connected components found")?; + .ok_or_else(|| anyhow!("No connected components found"))?; let max_width = self.max_width * max_height / self.max_height; for i in &components.bounds_by_position { @@ -334,7 +332,7 @@ impl CharMeasurements { d1.partial_cmp(d2).expect("NaN distance encountered") }) .map(|(i, _, d)| (d.sqrt(), i)) - .ok_or_else(|| "No chars in cache")?; + .ok_or_else(|| anyhow!("No chars in cache"))?; println!("char '{}', distance {}", best_match.1, best_match.0); if best_match.0 <= 0.75 { diff --git a/src/recognition/recognize.rs b/src/recognition/recognize.rs index 9109f31..bfecf0d 100644 --- a/src/recognition/recognize.rs +++ b/src/recognition/recognize.rs @@ -1,5 +1,6 @@ use std::fmt::Display; +use anyhow::{anyhow, bail}; use hypertesseract::{PageSegMode, Tesseract}; use image::imageops::FilterType; use image::{DynamicImage, GenericImageView}; @@ -170,7 +171,7 @@ impl ImageAnalyzer { }) { Ok(result) } else { - Err(format!("Score {result} is not vaild").into()) + Err(anyhow!("Score {result} is not vaild")) } } // }}} @@ -228,7 +229,7 @@ impl ImageAnalyzer { .zip(Difficulty::DIFFICULTY_STRINGS) .min_by_key(|(_, difficulty_string)| edit_distance(difficulty_string, &text)) .map(|(difficulty, _)| *difficulty) - .ok_or_else(|| format!("Unrecognised difficulty '{}'", text))?; + .ok_or_else(|| anyhow!("Unrecognised difficulty '{}'", text))?; Ok(difficulty) } @@ -272,12 +273,11 @@ impl ImageAnalyzer { )?; if conf < 20 && conf != 0 { - return Err(format!( + bail!( "Title text is not readable (confidence = {}, text = {}).", conf, text.trim() - ) - .into()); + ); } guess_chart_name(&text, &ctx.song_cache, Some(difficulty), false) @@ -319,10 +319,10 @@ impl ImageAnalyzer { let (distance, song_id) = ctx .jacket_cache .recognise(&*cropped) - .ok_or_else(|| "Could not recognise jacket")?; + .ok_or_else(|| anyhow!("Could not recognise jacket"))?; if distance > (IMAGE_VEC_DIM * 3) as f32 { - Err("No known jacket looks like this")?; + bail!("No known jacket looks like this"); } let (song, chart) = ctx.song_cache.lookup_by_difficulty(*song_id, difficulty)?; diff --git a/src/recognition/ui.rs b/src/recognition/ui.rs index 2d279b4..f649c15 100644 --- a/src/recognition/ui.rs +++ b/src/recognition/ui.rs @@ -1,5 +1,6 @@ use std::fs; +use anyhow::anyhow; use image::GenericImage; use crate::{assets::get_config_dir, bitmap::Rect, context::Error}; @@ -172,7 +173,7 @@ impl UIMeasurements { } } - Err(format!("Could no find rect for {rect:?} in image").into()) + Err(anyhow!("Could no find rect for {rect:?} in image")) } // }}} } diff --git a/src/user.rs b/src/user.rs index a32297f..1b46ba6 100644 --- a/src/user.rs +++ b/src/user.rs @@ -1,5 +1,6 @@ use std::str::FromStr; +use anyhow::anyhow; use poise::serenity_prelude::UserId; use rusqlite::Row; @@ -39,7 +40,7 @@ impl User { )? .query_map([&discord_id], |row| row.get("id"))? .next() - .ok_or_else(|| "Failed to create user")??; + .ok_or_else(|| anyhow!("Failed to create user"))??; Ok(Self { discord_id, @@ -57,7 +58,7 @@ impl User { .prepare_cached("SELECT * FROM users WHERE discord_id = ?")? .query_map([id], Self::from_row)? .next() - .ok_or_else(|| "You are not an user in my database, sowwy ^~^")??; + .ok_or_else(|| anyhow!("You are not an user in my database, sowwy ^~^"))??; Ok(user) } @@ -69,7 +70,7 @@ impl User { .prepare_cached("SELECT * FROM users WHERE id = ?")? .query_map([id], Self::from_row)? .next() - .ok_or_else(|| "You are not an user in my database, sowwy ^~^")??; + .ok_or_else(|| anyhow!("You are not an user in my database, sowwy ^~^"))??; Ok(user) }