Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: advance sort feature #106

Merged
merged 8 commits into from
Dec 2, 2024
Merged
51 changes: 46 additions & 5 deletions src/bot/meigen/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand Down Expand Up @@ -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)]
Expand All @@ -99,6 +125,19 @@ enum Command {
/// 指定した文字列を含む名言をリスト表示します
#[clap(long)]
content: Option<String>,

/// 指定した項目でソートします。
#[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,
},

/// 名言を削除します
Expand Down Expand Up @@ -149,13 +188,18 @@ impl<D: MeigenDatabase> BotService for MeigenBot<D> {
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?
}
Expand Down Expand Up @@ -210,14 +254,11 @@ impl<D: MeigenDatabase> MeigenBot<D> {
}

async fn search(&self, opt: FindOptions<'_>) -> Result<String> {
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);

Comment on lines -219 to -220
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

すまん

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#107 に免じてゆるされました

Ok(list(&res))
}

Expand Down
50 changes: 46 additions & 4 deletions src/db/mem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -337,12 +337,29 @@ impl MeigenDatabase for MemoryDB {
&& options.content.map_or(true, |c| x.content.contains(c))
})
.collect::<Vec<_>>();

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())
Expand Down Expand Up @@ -383,3 +400,28 @@ impl MeigenDatabase for MemoryDB {
Ok(true)
}
}

pub trait SortByKeyWithDirTraitExt<I> {
fn sort_by_key_with_dir<FnK, FnKR>(&mut self, key: FnK, dir: SortDirection)
where
FnK: Fn(&I) -> FnKR,
FnKR: Ord;
}

impl<I> SortByKeyWithDirTraitExt<I> for Vec<I> {
fn sort_by_key_with_dir<FnK, FnKR>(&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(),
}
})
}
}
64 changes: 45 additions & 19 deletions src/db/mongodb/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use {
meigen::{
self,
model::{Meigen, MeigenId},
MeigenDatabase,
MeigenDatabase, SortDirection, SortKey,
},
IsUpdated,
},
Expand Down Expand Up @@ -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
Expand Down
Loading