diff --git a/src/bot/meigen/mod.rs b/src/bot/meigen/mod.rs index bd402db..bf80f38 100644 --- a/src/bot/meigen/mod.rs +++ b/src/bot/meigen/mod.rs @@ -4,17 +4,37 @@ use { }, anyhow::{Context as _, Result}, async_trait::async_trait, + clap::{ArgGroup, ValueEnum}, model::{Meigen, MeigenId}, }; pub mod model; +#[derive(ValueEnum, Clone, Debug, PartialEq, Eq, Default)] +pub enum SortKey { + #[default] + Id, + Love, + Length, +} + +#[derive(ValueEnum, Clone, Debug, PartialEq, Eq, Default)] +pub enum SortDirection { + #[clap(alias = "a")] + #[default] + Asc, + #[clap(alias = "d")] + Desc, +} + #[derive(Default)] pub struct FindOptions<'a> { pub author: Option<&'a str>, pub content: Option<&'a str>, pub offset: u32, pub limit: u8, + pub sort: SortKey, + pub dir: SortDirection, pub random: bool, } @@ -75,6 +95,12 @@ enum Command { Status, /// 名言をリスト表示します + #[clap(group( + ArgGroup::new("dir_conflict") + .args(&["dir"]) + .requires("dir") + .conflicts_with("reverse") + ))] // --dir and --reverse conflicts List { /// 表示する名言のオフセット #[clap(long)] @@ -99,6 +125,19 @@ enum Command { /// 指定した文字列を含む名言をリスト表示します #[clap(long)] content: Option, + + /// 指定した項目でソートします。 + #[clap(value_enum, long, default_value_t)] + sort: SortKey, + + /// ソートの順番を入れ替えます。 + #[clap(value_enum, long, default_value_t)] + dir: SortDirection, + + /// 降順にします。--dir desc のエイリアスです。 + #[clap(short = 'R', long, alias = "rev")] + #[clap(default_value_t = false)] + reverse: bool, }, /// 名言を削除します @@ -149,13 +188,18 @@ impl BotService for MeigenBot { random, author, content, + sort, + dir, + reverse, } => { self.search(FindOptions { author: author.as_deref(), content: content.as_deref(), - random, offset, limit, + sort, + dir: if reverse { SortDirection::Desc } else { dir }, + random, }) .await? } @@ -210,14 +254,11 @@ impl MeigenBot { } async fn search(&self, opt: FindOptions<'_>) -> Result { - let mut res = self.db.search(opt).await?; - + let res = self.db.search(opt).await?; if res.is_empty() { return Ok("条件に合致する名言が見つかりませんでした".into()); } - res.sort_unstable_by_key(|x| x.id); - Ok(list(&res)) } diff --git a/src/db/mem.rs b/src/db/mem.rs index 9f6d465..e2f3f22 100644 --- a/src/db/mem.rs +++ b/src/db/mem.rs @@ -6,14 +6,14 @@ use { meigen::{ self, model::{Meigen, MeigenId}, - MeigenDatabase, + MeigenDatabase, SortDirection, SortKey, }, IsUpdated, }, anyhow::{anyhow, Context as _, Result}, async_trait::async_trait, chrono::{DateTime, Duration, Utc}, - rand::seq::SliceRandom, + rand::{seq::SliceRandom, thread_rng}, serde::Serialize, std::{collections::HashMap, ops::DerefMut, sync::Arc}, tokio::sync::Mutex, @@ -337,12 +337,29 @@ impl MeigenDatabase for MemoryDB { && options.content.map_or(true, |c| x.content.contains(c)) }) .collect::>(); + if options.random { - meigens.shuffle(&mut rand::thread_rng()); + meigens.shuffle(&mut thread_rng()); + meigens.truncate(options.limit as usize); } + + match options.sort { + SortKey::Id => meigens.sort_by_key_with_dir(|meigen| meigen.id, options.dir), + SortKey::Love => { + meigens.sort_by_key_with_dir(|meigen| meigen.loved_user_id.len(), options.dir) + } + SortKey::Length => { + meigens.sort_by_key_with_dir(|meigen| meigen.content.len(), options.dir) + } + } + Ok(meigens .into_iter() - .skip(options.offset as usize) + .skip(if options.random { + 0 + } else { + options.offset as usize + }) .take(options.limit as usize) .cloned() .collect()) @@ -383,3 +400,28 @@ impl MeigenDatabase for MemoryDB { Ok(true) } } + +pub trait SortByKeyWithDirTraitExt { + fn sort_by_key_with_dir(&mut self, key: FnK, dir: SortDirection) + where + FnK: Fn(&I) -> FnKR, + FnKR: Ord; +} + +impl SortByKeyWithDirTraitExt for Vec { + fn sort_by_key_with_dir(&mut self, key: FnK, dir: SortDirection) + where + FnK: Fn(&I) -> FnKR, + FnKR: Ord, + { + self.sort_by(|left, right| { + let left_key = key(left); + let right_key = key(right); + + match dir { + SortDirection::Asc => left_key.cmp(&right_key), + SortDirection::Desc => left_key.cmp(&right_key).reverse(), + } + }) + } +} diff --git a/src/db/mongodb/mod.rs b/src/db/mongodb/mod.rs index c6c79a4..927a71e 100644 --- a/src/db/mongodb/mod.rs +++ b/src/db/mongodb/mod.rs @@ -9,7 +9,7 @@ use { meigen::{ self, model::{Meigen, MeigenId}, - MeigenDatabase, + MeigenDatabase, SortDirection, SortKey, }, IsUpdated, }, @@ -489,32 +489,58 @@ impl MeigenDatabase for MongoDb { content, offset, limit, + sort, + dir, random, } = options; - let mut pipeline = vec![ - { - let into_regex = |x| doc! { "$regex": format!(".*{}.*", regex::escape(x)) }; - let mut doc = Document::new(); - if let Some(author) = author { - doc.insert("author", into_regex(author)); - } - if let Some(content) = content { - doc.insert("content", into_regex(content)); - } - doc! { "$match": doc } // { $match: {} } is fine, it just matches to any document. - }, - doc! { "$sort": { "id": -1 } }, - doc! { "$skip": offset }, - ]; + let mut pipeline = vec![{ + let into_regex = |x| doc! { "$regex": format!(".*{}.*", regex::escape(x)) }; + let mut doc = Document::new(); + if let Some(author) = author { + doc.insert("author", into_regex(author)); + } + if let Some(content) = content { + doc.insert("content", into_regex(content)); + } + doc! { "$match": doc } // { $match: {} } is fine, it just matches to any document. + }]; if random { + // `Randomized` skips/limits before shuffling pipeline.extend([ + doc! { "$skip": offset }, doc! { "$sample": { "size": limit as u32 } }, // sample pipeline scrambles document order. - doc! { "$sort": { "id": -1 } }, + doc! { "$limit": limit as u32 }, ]); - } else { - pipeline.push(doc! { "$limit": limit as u32 }); + } + + let dir = match dir { + SortDirection::Asc => 1, + SortDirection::Desc => -1, + }; + + match sort { + SortKey::Id => pipeline.extend([doc! { "$sort": { "id": dir } }]), + SortKey::Love => pipeline.extend([ + doc! { + "$addFields": { + "loved_users": { + "$size": { "$ifNull": ["$loved_user_id", []] } + } + } + }, + doc! { "$sort": { "loved_users": dir } }, + ]), + SortKey::Length => pipeline.extend([ + doc! { "$addFields": { "length": { "$strLenCP": "$content" }}}, + doc! { "$sort": { "length": dir } }, + ]), + }; + + if !random { + // `Randomized` skips/limits before shuffling + pipeline.extend([doc! { "$skip": offset }, doc! { "$limit": limit as u32 }]); } self.inner