Merge pull request #106 from approvers/feature/advanced-sort-feature

feat: advance sort feature
approvers · Dec 2, 2024 · 7258731 · 7258731
2 parents 37f2917 + 7947253
commit 7258731
Show file tree

Hide file tree

Showing 3 changed files with 137 additions and 28 deletions.
diff --git 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<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,
     },
 
     /// 名言を削除します
@@ -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?
             }
@@ -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);
-
         Ok(list(&res))
     }
 

diff --git 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::<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())
@@ -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(),
+            }
+        })
+    }
+}
diff --git 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