Provide a way to give pagination functions a total page count

youmubot-cf/src/lib.rs

@@ -1,6 +1,7 @@
 use std::{collections::HashMap, sync::Arc, time::Duration};
 
 use codeforces::Contest;
+use pagination::paginate_from_fn;
 use serenity::{
     builder::{CreateMessage, EditMessage},
     framework::standard::{
@@ -178,8 +179,8 @@ pub async fn ranks(ctx: &Context, m: &Message) -> CommandResult {
     let total_pages = (ranks.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE;
     let last_updated = ranks.iter().map(|(_, cfu)| cfu.last_update).min().unwrap();
 
-    paginate_reply_fn(
+    paginate_reply(
-        move |page, ctx, msg| {
+        paginate_from_fn(move |page, ctx, msg| {
             use Align::*;
             let ranks = ranks.clone();
             Box::pin(async move {
@@ -224,7 +225,8 @@ pub async fn ranks(ctx: &Context, m: &Message) -> CommandResult {
                 msg.edit(ctx, EditMessage::new().content(content)).await?;
                 Ok(true)
             })
-        },
+        })
+        .with_page_count(total_pages),
         ctx,
         m,
         std::time::Duration::from_secs(60),
@@ -315,8 +317,8 @@ pub(crate) async fn contest_rank_table(
     }
     let ranks = Arc::new(ranks);
 
-    paginate_reply_fn(
+    paginate_reply(
-        move |page, ctx, msg| {
+        paginate_from_fn(move |page, ctx, msg| {
             let contest = contest.clone();
             let problems = problems.clone();
             let ranks = ranks.clone();
@@ -390,7 +392,8 @@ pub(crate) async fn contest_rank_table(
                 msg.edit(ctx, EditMessage::new().content(content)).await?;
                 Ok(true)
             })
-        },
+        })
+        .with_page_count(total_pages),
         ctx,
         reply_to,
         Duration::from_secs(60),

youmubot-core/src/community/roles.rs

@@ -43,8 +43,8 @@ async fn list(ctx: &Context, m: &Message, _: Args) -> CommandResult {
             const ROLES_PER_PAGE: usize = 8;
             let pages = (roles.len() + ROLES_PER_PAGE - 1) / ROLES_PER_PAGE;
 
-            paginate_reply_fn(
+            paginate_reply(
-                |page, ctx, msg| {
+                paginate_from_fn(|page, ctx, msg| {
                     let roles = roles.clone();
                     Box::pin(async move {
                         let page = page as usize;
@@ -80,7 +80,8 @@ async fn list(ctx: &Context, m: &Message, _: Args) -> CommandResult {
                         msg.edit(ctx, EditMessage::new().content(content)).await?;
                         Ok(true)
                     })
-                },
+                })
+                .with_page_count(pages),
                 ctx,
                 m,
                 std::time::Duration::from_secs(60 * 10),

youmubot-core/src/fun/images.rs

@@ -65,8 +65,8 @@ async fn message_command(
         return Ok(());
     }
     let images = std::sync::Arc::new(images);
-    paginate_reply_fn(
+    paginate_reply(
-        move |page, ctx, msg: &mut Message| {
+        paginate_from_fn(|page, ctx, msg: &mut Message| {
             let images = images.clone();
             Box::pin(async move {
                 let page = page as usize;
@@ -87,7 +87,8 @@ async fn message_command(
                     .map_err(|e| e.into())
                 }
             })
-        },
+        })
+        .with_page_count(images.len()),
         ctx,
         msg,
         std::time::Duration::from_secs(120),

youmubot-osu/src/discord/hook.rs

@@ -1,7 +1,10 @@
 use std::str::FromStr;
+use std::sync::Arc;
 
 use lazy_static::lazy_static;
+use pagination::paginate_from_fn;
 use regex::Regex;
+use serenity::all::EditMessage;
 use serenity::{builder::CreateMessage, model::channel::Message, utils::MessageBuilder};
 
 use youmubot_prelude::*;

youmubot-osu/src/discord/server_rank.rs

@@ -101,12 +101,14 @@ pub async fn server_rank(ctx: &Context, m: &Message, mut args: Args) -> CommandR
         return Ok(());
     }
 
+    const ITEMS_PER_PAGE: usize = 10;
     let users = Arc::new(users);
     let last_update = last_update.unwrap();
-    paginate_reply_fn(
+    let total_len = users.len();
-        move |page: u8, ctx: &Context, m: &mut Message| {
+    let total_pages = (total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE;
+    paginate_reply(
+        paginate_from_fn(move |page: u8, ctx: &Context, m: &mut Message| {
             use Align::*;
-            const ITEMS_PER_PAGE: usize = 10;
             let users = users.clone();
             Box::pin(async move {
                 let start = (page as usize) * ITEMS_PER_PAGE;
@@ -114,7 +116,6 @@ pub async fn server_rank(ctx: &Context, m: &Message, mut args: Args) -> CommandR
                 if start >= end {
                     return Ok(false);
                 }
-                let total_len = users.len();
                 let users = &users[start..end];
                 let table = if matches!(mode, RankQuery::Mode(Mode::Std) | RankQuery::MapLength) {
                     const HEADERS: [&'static str; 5] =
@@ -167,14 +168,15 @@ pub async fn server_rank(ctx: &Context, m: &Message, mut args: Args) -> CommandR
                     .push_line(format!(
                         "Page **{}**/**{}**. Last updated: {}",
                         page + 1,
-                        (total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE,
+                        total_pages,
                         last_update.format("<t:%s:R>"),
                     ))
                     .build();
                 m.edit(ctx, EditMessage::new().content(content)).await?;
                 Ok(true)
             })
-        },
+        })
+        .with_page_count(total_pages),
         ctx,
         m,
         std::time::Duration::from_secs(60),
@@ -337,15 +339,17 @@ pub async fn show_leaderboard(ctx: &Context, msg: &Message, mut args: Args) -> C
     }
     let has_lazer_score = scores.iter().any(|(_, _, v)| v.score.is_none());
 
-    paginate_reply_fn(
-        move |page: u8, ctx: &Context, m: &mut Message| {
     const ITEMS_PER_PAGE: usize = 5;
+    let total_len = scores.len();
+    let total_pages = (total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE;
+
+    paginate_reply(
+        paginate_from_fn(move |page: u8, ctx: &Context, m: &mut Message| {
             let start = (page as usize) * ITEMS_PER_PAGE;
             let end = (start + ITEMS_PER_PAGE).min(scores.len());
             if start >= end {
                 return Box::pin(future::ready(Ok(false)));
             }
-            let total_len = scores.len();
             let scores = scores[start..end].to_vec();
             let bm = (bm.0.clone(), bm.1);
             Box::pin(async move {
@@ -392,7 +396,7 @@ pub async fn show_leaderboard(ctx: &Context, msg: &Message, mut args: Args) -> C
                     .push_line(format!(
                         "Page **{}**/**{}**. Not seeing your scores? Run `osu check` to update.",
                         page + 1,
-                        (total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE,
+                        total_pages,
                     ))
                     .push(
                         if let crate::models::ApprovalStatus::Ranked(_) = bm.0.approval {
@@ -408,7 +412,8 @@ pub async fn show_leaderboard(ctx: &Context, msg: &Message, mut args: Args) -> C
                 m.edit(&ctx, EditMessage::new().content(content)).await?;
                 Ok(true)
             })
-        },
+        })
+        .with_page_count(total_pages),
         ctx,
         msg,
         std::time::Duration::from_secs(60),

youmubot-prelude/src/lib.rs

@@ -18,7 +18,7 @@ pub use debugging_ok::OkPrint;
 pub use flags::Flags;
 pub use hook::Hook;
 pub use member_cache::MemberCache;
-pub use pagination::{paginate, paginate_fn, paginate_reply, paginate_reply_fn, Paginate};
+pub use pagination::{paginate, paginate_from_fn, paginate_reply, Paginate};
 
 pub mod announcer;
 pub mod args;

youmubot-prelude/src/pagination.rs

@@ -52,6 +52,70 @@ pub trait Paginate: Send + Sized {
     fn is_empty(&self) -> Option<bool> {
         self.len().map(|v| v == 0)
     }
+
+    /// Add a page count to the pagination.
+    fn with_page_count(self, page_count: usize) -> impl Paginate {
+        WithPageCount {
+            inner: self,
+            page_count,
+        }
+    }
+}
+
+pub fn paginate_from_fn(
+    pager: impl for<'m> FnMut(
+            u8,
+            &'m Context,
+            &'m mut Message,
+        ) -> std::pin::Pin<Box<dyn Future<Output = Result<bool>> + Send + 'm>>
+        + Send,
+) -> impl Paginate {
+    pager
+}
+
+struct WithPageCount<Inner> {
+    inner: Inner,
+    page_count: usize,
+}
+
+#[async_trait::async_trait]
+impl<Inner: Paginate> Paginate for WithPageCount<Inner> {
+    async fn render(&mut self, page: u8, ctx: &Context, m: &mut Message) -> Result<bool> {
+        if page as usize >= self.page_count {
+            return Ok(false);
+        }
+        self.inner.render(page, ctx, m).await
+    }
+    async fn prerender(&mut self, ctx: &Context, m: &mut Message) -> Result<()> {
+        self.inner.prerender(ctx, m).await
+    }
+
+    async fn handle_reaction(
+        &mut self,
+        page: u8,
+        ctx: &Context,
+        message: &mut Message,
+        reaction: &Reaction,
+    ) -> Result<Option<u8>> {
+        // handle normal reactions first, then fallback to the inner one
+        let new_page = handle_pagination_reaction(page, self, ctx, message, reaction).await?;
+
+        if new_page != page {
+            Ok(Some(new_page))
+        } else {
+            self.inner
+                .handle_reaction(page, ctx, message, reaction)
+                .await
+        }
+    }
+
+    fn len(&self) -> Option<usize> {
+        Some(self.page_count)
+    }
+
+    fn is_empty(&self) -> Option<bool> {
+        Some(self.page_count == 0)
+    }
 }
 
 #[async_trait::async_trait]
@@ -185,36 +249,6 @@ async fn paginate_with_first_message(
     res
 }
 
-/// Same as `paginate`, but for function inputs, especially anonymous functions.
-pub async fn paginate_fn(
-    pager: impl for<'m> FnMut(
-            u8,
-            &'m Context,
-            &'m mut Message,
-        ) -> std::pin::Pin<Box<dyn Future<Output = Result<bool>> + Send + 'm>>
-        + Send,
-    ctx: &Context,
-    channel: ChannelId,
-    timeout: std::time::Duration,
-) -> Result<()> {
-    paginate(pager, ctx, channel, timeout).await
-}
-
-/// Same as `paginate_reply`, but for function inputs, especially anonymous functions.
-pub async fn paginate_reply_fn(
-    pager: impl for<'m> FnMut(
-            u8,
-            &'m Context,
-            &'m mut Message,
-        ) -> std::pin::Pin<Box<dyn Future<Output = Result<bool>> + Send + 'm>>
-        + Send,
-    ctx: &Context,
-    reply_to: &Message,
-    timeout: std::time::Duration,
-) -> Result<()> {
-    paginate_reply(pager, ctx, reply_to, timeout).await
-}
-
 // Handle the reaction and return a new page number.
 pub async fn handle_pagination_reaction(
     page: u8,