osu: Make c/last/lb sensitive to replied-to messages

youmubot-osu/src/discord/hook.rs

@@ -12,13 +12,13 @@ use youmubot_prelude::*;
 use super::embeds::beatmap_embed;
 
 lazy_static! {
-    static ref OLD_LINK_REGEX: Regex = Regex::new(
+    pub(crate) static ref OLD_LINK_REGEX: Regex = Regex::new(
         r"(?:https?://)?osu\.ppy\.sh/(?P<link_type>s|b)/(?P<id>\d+)(?:[\&\?]m=(?P<mode>\d))?(?:\+(?P<mods>[A-Z]+))?"
     ).unwrap();
-    static ref NEW_LINK_REGEX: Regex = Regex::new(
+    pub(crate) static ref NEW_LINK_REGEX: Regex = Regex::new(
         r"(?:https?://)?osu\.ppy\.sh/beatmapsets/(?P<set_id>\d+)/?(?:\#(?P<mode>osu|taiko|fruits|mania)(?:/(?P<beatmap_id>\d+)|/?))?(?:\+(?P<mods>[A-Z]+))?"
     ).unwrap();
-    static ref SHORT_LINK_REGEX: Regex = Regex::new(
+    pub(crate) static ref SHORT_LINK_REGEX: Regex = Regex::new(
         r"(?:^|\s|\W)(?P<main>/b/(?P<id>\d+)(?:/(?P<mode>osu|taiko|fruits|mania))?(?:\+(?P<mods>[A-Z]+))?)"
     ).unwrap();
 }

youmubot-osu/src/discord/mod.rs

@@ -3,7 +3,7 @@ use crate::{
     discord::display::ScoreListStyle,
     discord::oppai_cache::{BeatmapCache, BeatmapInfo},
     models::{Beatmap, Mode, Mods, User},
-    request::UserID,
+    request::{BeatmapRequestKind, UserID},
     Client as OsuHttpClient,
 };
 use serenity::{
@@ -31,6 +31,7 @@ mod server_rank;
 use db::OsuUser;
 use db::{OsuLastBeatmap, OsuSavedUsers, OsuUserBests};
 use embeds::{beatmap_embed, score_embed, user_embed};
+use hook::SHORT_LINK_REGEX;
 pub use hook::{dot_osu_hook, hook};
 use server_rank::{SERVER_RANK_COMMAND, UPDATE_LEADERBOARD_COMMAND};
 
@@ -389,6 +390,61 @@ impl FromStr for OptBeatmapset {
     }
 }
 
+/// Load the mentioned beatmap from the given message.
+pub(crate) async fn load_beatmap(
+    ctx: &Context,
+    msg: &Message,
+) -> Option<(BeatmapWithMode, Option<Mods>)> {
+    let data = ctx.data.read().await;
+
+    if let Some(replied) = &msg.referenced_message {
+        // Try to look for a mention of the replied message.
+        let beatmap_id = SHORT_LINK_REGEX.captures(&replied.content).or_else(|| {
+            msg.embeds.iter().find_map(|e| {
+                e.description
+                    .as_ref()
+                    .and_then(|v| SHORT_LINK_REGEX.captures(&v))
+                    .or_else(|| {
+                        e.fields
+                            .iter()
+                            .find_map(|f| SHORT_LINK_REGEX.captures(&f.value))
+                    })
+            })
+        });
+        if let Some(caps) = beatmap_id {
+            let id: u64 = caps.name("id").unwrap().as_str().parse().unwrap();
+            let mode = caps
+                .name("mode")
+                .and_then(|m| Mode::parse_from_new_site(m.as_str()));
+            let mods = caps
+                .name("mods")
+                .and_then(|m| m.as_str().parse::<Mods>().ok());
+            let osu = data.get::<OsuClient>().unwrap();
+            let bms = osu
+                .beatmaps(BeatmapRequestKind::Beatmap(id), |f| f.maybe_mode(mode))
+                .await
+                .ok()
+                .and_then(|v| v.into_iter().next());
+            if let Some(beatmap) = bms {
+                let bm_mode = beatmap.mode.clone();
+                let bm = BeatmapWithMode(beatmap, mode.unwrap_or(bm_mode));
+                // Store the beatmap in history
+                cache::save_beatmap(&*data, msg.channel_id, &bm)
+                    .await
+                    .pls_ok();
+
+                return Some((bm, mods));
+            }
+        }
+    }
+
+    let b = cache::get_beatmap(&*data, msg.channel_id)
+        .await
+        .ok()
+        .flatten();
+    b.map(|b| (b, None))
+}
+
 #[command]
 #[description = "Show information from the last queried beatmap."]
 #[usage = "[--set/-s/--beatmapset] / [mods = no mod]"]
@@ -396,12 +452,12 @@ impl FromStr for OptBeatmapset {
 #[max_args(2)]
 pub async fn last(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let data = ctx.data.read().await;
-    let b = cache::get_beatmap(&*data, msg.channel_id).await?;
+    let b = load_beatmap(ctx, msg).await;
     let beatmapset = args.find::<OptBeatmapset>().is_ok();
 
     match b {
-        Some(BeatmapWithMode(b, m)) => {
-            let mods = args.find::<Mods>().unwrap_or(Mods::NOMOD);
+        Some((BeatmapWithMode(b, m), mods_def)) => {
+            let mods = args.find::<Mods>().ok().or(mods_def).unwrap_or(Mods::NOMOD);
             if beatmapset {
                 let beatmap_cache = data.get::<BeatmapMetaCache>().unwrap();
                 let beatmapset = beatmap_cache.get_beatmapset(b.beatmapset_id).await?;
@@ -447,15 +503,15 @@ pub async fn last(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult
 #[max_args(3)]
 pub async fn check(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let data = ctx.data.read().await;
-    let mods = args.find::<Mods>().unwrap_or(Mods::NOMOD);
-    let bm = cache::get_beatmap(&*data, msg.channel_id).await?;
+    let bm = load_beatmap(ctx, msg).await;
 
     match bm {
         None => {
             msg.reply(&ctx, "No beatmap queried on this channel.")
                 .await?;
         }
-        Some(bm) => {
+        Some((bm, mods_def)) => {
+            let mods = args.find::<Mods>().ok().or(mods_def).unwrap_or(Mods::NOMOD);
             let b = &bm.0;
             let m = bm.1;
             let style = args

youmubot-osu/src/discord/server_rank.rs

@@ -1,5 +1,4 @@
 use super::{
-    cache::get_beatmap,
     db::{OsuSavedUsers, OsuUserBests},
     ModeArg, OsuClient,
 };
@@ -179,7 +178,6 @@ impl std::str::FromStr for OrderBy {
 pub async fn update_leaderboard(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
     let sort_order = args.single::<OrderBy>().unwrap_or_default();
     let style = args.single::<ScoreListStyle>().unwrap_or_default();
-    let mods = args.find::<Mods>().unwrap_or(Mods::NOMOD);
 
     let guild = m.guild_id.unwrap();
     let data = ctx.data.read().await;
@@ -191,8 +189,11 @@ pub async fn update_leaderboard(ctx: &Context, m: &Message, mut args: Args) -> C
         }
         Some(v) => v,
     };
-    let bm = match get_beatmap(&*data, m.channel_id).await? {
-        Some(bm) => bm,
+    let (bm, mods) = match super::load_beatmap(ctx, m).await {
+        Some((bm, mods_def)) => {
+            let mods = args.find::<Mods>().ok().or(mods_def).unwrap_or(Mods::NOMOD);
+            (bm, mods)
+        }
         None => {
             m.reply(&ctx, "No beatmap queried on this channel.").await?;
             return Ok(());

youmubot-osu/src/request.rs

@@ -108,6 +108,14 @@ pub mod builders {
             self
         }
 
+        pub fn maybe_mode(&mut self, mode: Option<Mode>) -> &mut Self {
+            if let Some(m) = mode {
+                self.mode(m, true)
+            } else {
+                self
+            }
+        }
+
         pub fn mode(&mut self, mode: Mode, converted: bool) -> &mut Self {
             self.mode = Some((mode, converted));
             self