Swap Command and related structs to avoid String allocs (#318)

serenity-rs · Nov 13, 2024 · 6d2ae48 · 6d2ae48
1 parent 3a2ab5e
commit 6d2ae48
Show file tree

Hide file tree

Showing 16 changed files with 143 additions and 115 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -20,7 +20,7 @@ jobs:
 
         include:
           - name: MSRV
-            toolchain: 1.74.0
+            toolchain: 1.80.0
             # don't do doctests because they rely on new features for brevity
             # copy known Cargo.lock to avoid random dependency MSRV bumps to mess up our test
             command: cp .github/Cargo-msrv.lock Cargo.lock && cargo test --all-features --lib --tests

diff --git a/Cargo.toml b/Cargo.toml
@@ -3,7 +3,7 @@ authors = ["kangalio <[email protected]>"]
 edition = "2021"
 name = "poise"
 version = "0.6.1"
-rust-version = "1.74.0"
+rust-version = "1.80.0"
 description = "A Discord bot framework for serenity"
 license = "MIT"
 repository = "https://github.com/serenity-rs/poise/"

diff --git a/README.md b/README.md
@@ -3,7 +3,7 @@
 [![Docs](https://img.shields.io/badge/docs-online-informational)](https://docs.rs/poise/)
 [![Docs (git)](https://img.shields.io/badge/docs%20%28git%29-online-informational)](https://serenity-rs.github.io/poise/)
 [![License: MIT](https://img.shields.io/badge/license-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
-[![Rust: 1.74+](https://img.shields.io/badge/rust-1.74+-93450a)](https://blog.rust-lang.org/2023/11/16/Rust-1.74.0.html)
+[![Rust: 1.80+](https://img.shields.io/badge/rust-1.80+-93450a)](https://blog.rust-lang.org/2024/07/25/Rust-1.80.0.html)
 
 # Poise
 Poise is an opinionated Discord bot framework with a few distinctive features:

diff --git a/examples/fluent_localization/main.rs b/examples/fluent_localization/main.rs
@@ -4,7 +4,7 @@ use poise::serenity_prelude as serenity;
 use translation::tr;
 
 pub struct Data {
-    translations: translation::Translations,
+    translations: &'static translation::Translations,
 }
 
 type Error = Box<dyn std::error::Error + Send + Sync>;
@@ -62,7 +62,12 @@ async fn main() {
 
     let mut commands = vec![welcome(), info(), register()];
     let translations = translation::read_ftl().expect("failed to read translation files");
-    translation::apply_translations(&translations, &mut commands);
+
+    // We leak the translations so we can easily copy around `&'static str`s, to the downside
+    // that the OS will reclaim the memory at the end of `main` instead of the Drop implementation.
+    let translations: &'static translation::Translations = Box::leak(Box::new(translations));
+
+    translation::apply_translations(translations, &mut commands);
 
     let token = std::env::var("TOKEN").unwrap();
     let intents = serenity::GatewayIntents::non_privileged();

diff --git a/examples/fluent_localization/translation.rs b/examples/fluent_localization/translation.rs
@@ -1,5 +1,7 @@
 //! Wraps the fluent API and provides easy to use functions and macros for translation
 
+use std::borrow::Cow;
+
 use crate::{Context, Data, Error};
 
 type FluentBundle = fluent::bundle::FluentBundle<
@@ -30,27 +32,27 @@ pub(crate) use tr;
 
 /// Given a language file and message identifier, returns the translation
 pub fn format(
-    bundle: &FluentBundle,
+    bundle: &'static FluentBundle,
     id: &str,
     attr: Option<&str>,
     args: Option<&fluent::FluentArgs<'_>>,
-) -> Option<String> {
+) -> Option<Cow<'static, str>> {
     let message = bundle.get_message(id)?;
     let pattern = match attr {
         Some(attribute) => message.get_attribute(attribute)?.value(),
         None => message.value()?,
     };
     let formatted = bundle.format_pattern(pattern, args, &mut vec![]);
-    Some(formatted.into_owned())
+    Some(formatted)
 }
 
 /// Retrieves the appropriate language file depending on user locale and calls [`format`]
 pub fn get(
     ctx: Context,
-    id: &str,
+    id: &'static str,
     attr: Option<&str>,
     args: Option<&fluent::FluentArgs<'_>>,
-) -> String {
+) -> Cow<'static, str> {
     let translations = &ctx.data().translations;
     ctx.locale()
         // Try to get the language-specific translation
@@ -60,7 +62,7 @@ pub fn get(
         // If this message ID is not present in any translation files whatsoever
         .unwrap_or_else(|| {
             tracing::warn!("unknown fluent message identifier `{}`", id);
-            id.to_string()
+            Cow::Borrowed(id)
         })
 }
 
@@ -97,7 +99,7 @@ pub fn read_ftl() -> Result<Translations, Error> {
 
 /// Given a set of language files, fills in command strings and their localizations accordingly
 pub fn apply_translations(
-    translations: &Translations,
+    translations: &'static Translations,
     commands: &mut [poise::Command<Data, Error>],
 ) {
     for command in &mut *commands {
@@ -108,21 +110,24 @@ pub fn apply_translations(
                 Some(x) => x,
                 None => continue, // no localization entry => skip localization
             };
-            command
-                .name_localizations
-                .insert(locale.clone(), localized_command_name);
-            command.description_localizations.insert(
+
+            let locale = Cow::Borrowed(locale.as_str());
+            let name_localizations = command.name_localizations.to_mut();
+            let description_localizations = command.description_localizations.to_mut();
+
+            name_localizations.push((locale.clone(), localized_command_name));
+            description_localizations.push((
                 locale.clone(),
                 format(bundle, &command.name, Some("description"), None).unwrap(),
-            );
+            ));
 
             for parameter in &mut command.parameters {
                 // Insert localized parameter name and description
-                parameter.name_localizations.insert(
+                parameter.name_localizations.to_mut().push((
                     locale.clone(),
                     format(bundle, &command.name, Some(&parameter.name), None).unwrap(),
-                );
-                parameter.description_localizations.insert(
+                ));
+                parameter.description_localizations.to_mut().push((
                     locale.clone(),
                     format(
                         bundle,
@@ -131,14 +136,14 @@ pub fn apply_translations(
                         None,
                     )
                     .unwrap(),
-                );
+                ));
 
                 // If this is a choice parameter, insert its localized variants
-                for choice in &mut parameter.choices {
-                    choice.localizations.insert(
+                for choice in parameter.choices.to_mut().iter_mut() {
+                    choice.localizations.to_mut().push((
                         locale.clone(),
                         format(bundle, &choice.name, None, None).unwrap(),
-                    );
+                    ));
                 }
             }
         }
@@ -170,7 +175,7 @@ pub fn apply_translations(
             );
 
             // If this is a choice parameter, set the choice names to en-US
-            for choice in &mut parameter.choices {
+            for choice in parameter.choices.to_mut().iter_mut() {
                 choice.name = format(bundle, &choice.name, None, None).unwrap();
             }
         }

diff --git a/macros/src/choice_parameter.rs b/macros/src/choice_parameter.rs
@@ -67,14 +67,16 @@ pub fn choice_parameter(input: syn::DeriveInput) -> Result<TokenStream, darling:
     let indices = 0..variant_idents.len();
     Ok(quote::quote! {
         impl poise::ChoiceParameter for #enum_ident {
-            fn list() -> Vec<poise::CommandParameterChoice> {
-                vec![ #( poise::CommandParameterChoice {
+            fn list() -> std::borrow::Cow<'static, [poise::CommandParameterChoice]> {
+                use ::std::borrow::Cow;
+
+                Cow::Borrowed(&[ #( poise::CommandParameterChoice {
                     __non_exhaustive: (),
-                    name: #names.to_string(),
-                    localizations: std::collections::HashMap::from([
-                        #( (#locales.to_string(), #localized_names.to_string()) ),*
+                    name: Cow::Borrowed(#names),
+                    localizations: Cow::Borrowed(&[
+                        #( (Cow::Borrowed(#locales), Cow::Borrowed(#localized_names)) ),*
                     ]),
-                }, )* ]
+                }, )* ])
             }
 
             fn from_index(index: usize) -> Option<Self> {

diff --git a/macros/src/command/mod.rs b/macros/src/command/mod.rs
@@ -2,7 +2,7 @@ mod prefix;
 mod slash;
 
 use crate::util::{
-    iter_tuple_2_to_hash_map, wrap_option, wrap_option_and_map, wrap_option_to_string,
+    iter_tuple_2_to_vec_map, wrap_option, wrap_option_and_map, wrap_option_to_string,
 };
 use proc_macro::TokenStream;
 use syn::spanned::Spanned as _;
@@ -343,9 +343,9 @@ fn generate_command(mut inv: Invocation) -> Result<proc_macro2::TokenStream, dar
         None => quote::quote! { Box::new(()) },
     };
 
-    let name_localizations = iter_tuple_2_to_hash_map(inv.args.name_localized.into_iter());
+    let name_localizations = iter_tuple_2_to_vec_map(inv.args.name_localized.into_iter());
     let description_localizations =
-        iter_tuple_2_to_hash_map(inv.args.description_localized.into_iter());
+        iter_tuple_2_to_vec_map(inv.args.description_localized.into_iter());
 
     let function_ident =
         std::mem::replace(&mut inv.function.sig.ident, syn::parse_quote! { inner });
@@ -359,6 +359,8 @@ fn generate_command(mut inv: Invocation) -> Result<proc_macro2::TokenStream, dar
             <#ctx_type_with_static as poise::_GetGenerics>::U,
             <#ctx_type_with_static as poise::_GetGenerics>::E,
         > {
+            use ::std::borrow::Cow;
+
             #function
 
             ::poise::Command {
@@ -368,11 +370,11 @@ fn generate_command(mut inv: Invocation) -> Result<proc_macro2::TokenStream, dar
 
                 subcommands: vec![ #( #subcommands() ),* ],
                 subcommand_required: #subcommand_required,
-                name: #command_name.to_string(),
+                name: Cow::Borrowed(#command_name),
                 name_localizations: #name_localizations,
-                qualified_name: String::from(#command_name), // properly filled in later by Framework
-                identifying_name: String::from(#identifying_name),
-                source_code_name: String::from(#function_name),
+                qualified_name: Cow::Borrowed(#command_name), // properly filled in later by Framework
+                identifying_name: Cow::Borrowed(#identifying_name),
+                source_code_name: Cow::Borrowed(#function_name),
                 category: #category,
                 description: #description,
                 description_localizations: #description_localizations,
@@ -396,7 +398,7 @@ fn generate_command(mut inv: Invocation) -> Result<proc_macro2::TokenStream, dar
                 parameters: vec![ #( #parameters ),* ],
                 custom_data: #custom_data,
 
-                aliases: vec![ #( #aliases.to_string(), )* ],
+                aliases: Cow::Borrowed(&[ #( Cow::Borrowed(#aliases), )* ]),
                 invoke_on_edit: #invoke_on_edit,
                 track_deletion: #track_deletion,
                 broadcast_typing: #broadcast_typing,

diff --git a/macros/src/command/slash.rs b/macros/src/command/slash.rs
@@ -1,7 +1,6 @@
 use super::Invocation;
 use crate::util::{
-    extract_type_parameter, iter_tuple_2_to_hash_map, tuple_2_iter_deref, wrap_option_to_string,
-    List,
+    extract_type_parameter, iter_tuple_2_to_vec_map, tuple_2_iter_deref, wrap_option_to_string,
 };
 use quote::format_ident;
 use syn::spanned::Spanned as _;
@@ -42,9 +41,9 @@ pub fn generate_parameters(inv: &Invocation) -> Result<Vec<proc_macro2::TokenStr
 
         let param_name = &param.name;
         let name_localizations =
-            iter_tuple_2_to_hash_map(tuple_2_iter_deref(&param.args.name_localized));
+            iter_tuple_2_to_vec_map(tuple_2_iter_deref(&param.args.name_localized));
         let desc_localizations =
-            iter_tuple_2_to_hash_map(tuple_2_iter_deref(&param.args.description_localized));
+            iter_tuple_2_to_vec_map(tuple_2_iter_deref(&param.args.description_localized));
 
         let autocomplete_callback = match &param.args.autocomplete {
             Some(autocomplete_fn) => {
@@ -91,37 +90,34 @@ pub fn generate_parameters(inv: &Invocation) -> Result<Vec<proc_macro2::TokenStr
         };
         // TODO: theoretically a problem that we don't store choices for non slash commands
         // TODO: move this to poise::CommandParameter::choices (is there a reason not to?)
-        let choices = match inv.args.slash_command {
-            true => {
-                if let Some(List(choices)) = &param.args.choices {
-                    let choices = choices
-                        .iter()
-                        .map(lit_to_string)
-                        .collect::<Result<Vec<_>, _>>()?;
-
-                    quote::quote! { vec![#( ::poise::CommandParameterChoice {
-                        name: String::from(#choices),
-                        localizations: Default::default(),
-                        __non_exhaustive: (),
-                    } ),*] }
-                } else {
-                    quote::quote! { poise::slash_argument_choices!(#type_) }
-                }
+        let choices = if inv.args.slash_command {
+            if let Some(choices) = &param.args.choices {
+                let choices_iter = choices.0.iter();
+                let choices: Vec<_> = choices_iter.map(lit_to_string).collect::<Result<_, _>>()?;
+
+                quote::quote! { Cow::Borrowed(&[#( ::poise::CommandParameterChoice {
+                    name: Cow::Borrowed(#choices),
+                    localizations: Cow::Borrowed(&[]),
+                    __non_exhaustive: (),
+                } ),*]) }
+            } else {
+                quote::quote! { poise::slash_argument_choices!(#type_) }
             }
-            false => quote::quote! { vec![] },
+        } else {
+            quote::quote! { Cow::Borrowed(&[]) }
         };
 
         let channel_types = match &param.args.channel_types {
             Some(crate::util::List(channel_types)) => quote::quote! { Some(
-                vec![ #( poise::serenity_prelude::ChannelType::#channel_types ),* ]
+                Cow::Borrowed(&[ #( poise::serenity_prelude::ChannelType::#channel_types ),* ])
             ) },
             None => quote::quote! { None },
         };
 
         parameter_structs.push((
             quote::quote! {
                 ::poise::CommandParameter {
-                    name: #param_name.to_string(),
+                    name: ::std::borrow::Cow::Borrowed(#param_name),
                     name_localizations: #name_localizations,
                     description: #description,
                     description_localizations: #desc_localizations,

diff --git a/macros/src/util.rs b/macros/src/util.rs
@@ -38,8 +38,8 @@ pub fn wrap_option_and_map<T: quote::ToTokens>(
 }
 
 pub fn wrap_option_to_string<T: quote::ToTokens>(literal: Option<T>) -> syn::Expr {
-    let to_string_path = quote::quote!(::std::string::ToString::to_string);
-    wrap_option_and_map(literal, to_string_path)
+    let cowstr_path = quote::quote!(Cow::Borrowed);
+    wrap_option_and_map(literal, cowstr_path)
 }
 
 /// Syn Fold to make all lifetimes 'static. Used to access trait items of a type without having its
@@ -99,13 +99,13 @@ where
         .map(|Tuple2(t, v)| Tuple2(t.deref(), v.deref()))
 }
 
-pub fn iter_tuple_2_to_hash_map<I, T>(v: I) -> proc_macro2::TokenStream
+pub fn iter_tuple_2_to_vec_map<I, T>(v: I) -> proc_macro2::TokenStream
 where
     I: ExactSizeIterator<Item = Tuple2<T>>,
     T: quote::ToTokens,
 {
     if v.len() == 0 {
-        return quote::quote!(std::collections::HashMap::new());
+        return quote::quote!(Cow::Borrowed(&[]));
     }
 
     let (keys, values) = v
@@ -114,8 +114,8 @@ where
         .unzip::<_, _, Vec<_>, Vec<_>>();
 
     quote::quote! {
-        std::collections::HashMap::from([
-            #( (#keys.to_string(), #values.to_string()) ),*
+        Cow::Borrowed(&[
+            #( (Cow::Borrowed(#keys), Cow::Borrowed(#values)) ),*
         ])
     }
 }
diff --git a/src/builtins/mod.rs b/src/builtins/mod.rs
@@ -235,7 +235,7 @@ pub async fn autocomplete_command<'a, U, E>(
         .take(25);
 
     let choices = filtered_commands
-        .map(|cmd| serenity::AutocompleteChoice::from(&cmd.name))
+        .map(|cmd| serenity::AutocompleteChoice::from(cmd.name.as_ref()))
         .collect();
 
     serenity::CreateAutocompleteResponse::new().set_choices(choices)

diff --git a/src/choice_parameter.rs b/src/choice_parameter.rs
@@ -1,13 +1,13 @@
 //! Contains the [`ChoiceParameter`] trait and the blanket [`crate::SlashArgument`] and
 //! [`crate::PopArgument`] impl
 
-use crate::serenity_prelude as serenity;
+use crate::{serenity_prelude as serenity, CowVec};
 
 /// This trait is implemented by [`crate::macros::ChoiceParameter`]. See its docs for more
 /// information
 pub trait ChoiceParameter: Sized {
     /// Returns all possible choices for this parameter, in the order they will appear in Discord.
-    fn list() -> Vec<crate::CommandParameterChoice>;
+    fn list() -> CowVec<crate::CommandParameterChoice>;
 
     /// Returns an instance of [`Self`] corresponding to the given index into [`Self::list()`]
     fn from_index(index: usize) -> Option<Self>;
@@ -50,7 +50,7 @@ impl<T: ChoiceParameter> crate::SlashArgument for T {
         builder.kind(serenity::CommandOptionType::Integer)
     }
 
-    fn choices() -> Vec<crate::CommandParameterChoice> {
+    fn choices() -> CowVec<crate::CommandParameterChoice> {
         Self::list()
     }
 }