add support for custom suggestions in azalea-brigadier and cleanup a bit

12 files changed, 155 insertions, 117 deletions

diff --git a/azalea-brigadier/src/arguments/double_argument_type.rs b/azalea-brigadier/src/arguments/double_argument_type.rs
index 559d1cf5..2e50d291 100644
--- a/azalea-brigadier/src/arguments/double_argument_type.rs
+++ b/azalea-brigadier/src/arguments/double_argument_type.rs
@@ -17,25 +17,25 @@ impl ArgumentType for Double {
     fn parse(&self, reader: &mut StringReader) -> Result<Arc<dyn Any>, CommandSyntaxException> {
         let start = reader.cursor;
         let result = reader.read_double()?;
-        if let Some(minimum) = self.minimum {
-            if result < minimum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::DoubleTooSmall {
-                    found: result,
-                    min: minimum,
-                }
-                .create_with_context(reader));
+        if let Some(minimum) = self.minimum
+            && result < minimum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::DoubleTooSmall {
+                found: result,
+                min: minimum,
             }
+            .create_with_context(reader));
         }
-        if let Some(maximum) = self.maximum {
-            if result > maximum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::DoubleTooBig {
-                    found: result,
-                    max: maximum,
-                }
-                .create_with_context(reader));
+        if let Some(maximum) = self.maximum
+            && result > maximum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::DoubleTooBig {
+                found: result,
+                max: maximum,
             }
+            .create_with_context(reader));
         }
         Ok(Arc::new(result))
     }
diff --git a/azalea-brigadier/src/arguments/float_argument_type.rs b/azalea-brigadier/src/arguments/float_argument_type.rs
index 83e298ff..23dc88a5 100644
--- a/azalea-brigadier/src/arguments/float_argument_type.rs
+++ b/azalea-brigadier/src/arguments/float_argument_type.rs
@@ -17,25 +17,25 @@ impl ArgumentType for Float {
     fn parse(&self, reader: &mut StringReader) -> Result<Arc<dyn Any>, CommandSyntaxException> {
         let start = reader.cursor;
         let result = reader.read_float()?;
-        if let Some(minimum) = self.minimum {
-            if result < minimum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::FloatTooSmall {
-                    found: result,
-                    min: minimum,
-                }
-                .create_with_context(reader));
+        if let Some(minimum) = self.minimum
+            && result < minimum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::FloatTooSmall {
+                found: result,
+                min: minimum,
             }
+            .create_with_context(reader));
         }
-        if let Some(maximum) = self.maximum {
-            if result > maximum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::FloatTooBig {
-                    found: result,
-                    max: maximum,
-                }
-                .create_with_context(reader));
+        if let Some(maximum) = self.maximum
+            && result > maximum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::FloatTooBig {
+                found: result,
+                max: maximum,
             }
+            .create_with_context(reader));
         }
         Ok(Arc::new(result))
     }
diff --git a/azalea-brigadier/src/arguments/integer_argument_type.rs b/azalea-brigadier/src/arguments/integer_argument_type.rs
index 47d25e27..b993d200 100644
--- a/azalea-brigadier/src/arguments/integer_argument_type.rs
+++ b/azalea-brigadier/src/arguments/integer_argument_type.rs
@@ -17,25 +17,25 @@ impl ArgumentType for Integer {
     fn parse(&self, reader: &mut StringReader) -> Result<Arc<dyn Any>, CommandSyntaxException> {
         let start = reader.cursor;
         let result = reader.read_int()?;
-        if let Some(minimum) = self.minimum {
-            if result < minimum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::IntegerTooSmall {
-                    found: result,
-                    min: minimum,
-                }
-                .create_with_context(reader));
+        if let Some(minimum) = self.minimum
+            && result < minimum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::IntegerTooSmall {
+                found: result,
+                min: minimum,
             }
+            .create_with_context(reader));
         }
-        if let Some(maximum) = self.maximum {
-            if result > maximum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::IntegerTooBig {
-                    found: result,
-                    max: maximum,
-                }
-                .create_with_context(reader));
+        if let Some(maximum) = self.maximum
+            && result > maximum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::IntegerTooBig {
+                found: result,
+                max: maximum,
             }
+            .create_with_context(reader));
         }
         Ok(Arc::new(result))
     }
diff --git a/azalea-brigadier/src/arguments/long_argument_type.rs b/azalea-brigadier/src/arguments/long_argument_type.rs
index ba65479c..1e27cf9d 100644
--- a/azalea-brigadier/src/arguments/long_argument_type.rs
+++ b/azalea-brigadier/src/arguments/long_argument_type.rs
@@ -17,25 +17,25 @@ impl ArgumentType for Long {
     fn parse(&self, reader: &mut StringReader) -> Result<Arc<dyn Any>, CommandSyntaxException> {
         let start = reader.cursor;
         let result = reader.read_long()?;
-        if let Some(minimum) = self.minimum {
-            if result < minimum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::LongTooSmall {
-                    found: result,
-                    min: minimum,
-                }
-                .create_with_context(reader));
+        if let Some(minimum) = self.minimum
+            && result < minimum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::LongTooSmall {
+                found: result,
+                min: minimum,
             }
+            .create_with_context(reader));
         }
-        if let Some(maximum) = self.maximum {
-            if result > maximum {
-                reader.cursor = start;
-                return Err(BuiltInExceptions::LongTooBig {
-                    found: result,
-                    max: maximum,
-                }
-                .create_with_context(reader));
+        if let Some(maximum) = self.maximum
+            && result > maximum
+        {
+            reader.cursor = start;
+            return Err(BuiltInExceptions::LongTooBig {
+                found: result,
+                max: maximum,
             }
+            .create_with_context(reader));
         }
         Ok(Arc::new(result))
     }
diff --git a/azalea-brigadier/src/builder/argument_builder.rs b/azalea-brigadier/src/builder/argument_builder.rs
index 9ebe6400..bc7f0ac6 100644
--- a/azalea-brigadier/src/builder/argument_builder.rs
+++ b/azalea-brigadier/src/builder/argument_builder.rs
@@ -9,10 +9,20 @@ use crate::{
     tree::{Command, CommandNode},
 };
 
-#[derive(Debug, Clone)]
-pub enum ArgumentBuilderType {
+#[derive(Debug)]
+pub enum ArgumentBuilderType<S> {
     Literal(Literal),
-    Argument(Argument),
+    Argument(Argument<S>),
+}
+impl<S> Clone for ArgumentBuilderType<S> {
+    fn clone(&self) -> Self {
+        match self {
+            ArgumentBuilderType::Literal(literal) => ArgumentBuilderType::Literal(literal.clone()),
+            ArgumentBuilderType::Argument(argument) => {
+                ArgumentBuilderType::Argument(argument.clone())
+            }
+        }
+    }
 }
 
 /// A node that hasn't yet been built.
@@ -30,7 +40,7 @@ pub struct ArgumentBuilder<S> {
 
 /// A node that isn't yet built.
 impl<S> ArgumentBuilder<S> {
-    pub fn new(value: ArgumentBuilderType) -> Self {
+    pub fn new(value: ArgumentBuilderType<S>) -> Self {
         Self {
             arguments: CommandNode {
                 value,
diff --git a/azalea-brigadier/src/builder/literal_argument_builder.rs b/azalea-brigadier/src/builder/literal_argument_builder.rs
index 6627ffdc..ea5db2d3 100644
--- a/azalea-brigadier/src/builder/literal_argument_builder.rs
+++ b/azalea-brigadier/src/builder/literal_argument_builder.rs
@@ -12,7 +12,7 @@ impl Literal {
     }
 }
 
-impl From<Literal> for ArgumentBuilderType {
+impl<S> From<Literal> for ArgumentBuilderType<S> {
     fn from(literal: Literal) -> Self {
         Self::Literal(literal)
     }
diff --git a/azalea-brigadier/src/builder/required_argument_builder.rs b/azalea-brigadier/src/builder/required_argument_builder.rs
index 51f0acec..7c0f1015 100644
--- a/azalea-brigadier/src/builder/required_argument_builder.rs
+++ b/azalea-brigadier/src/builder/required_argument_builder.rs
@@ -1,25 +1,35 @@
-use std::{any::Any, fmt::Debug, sync::Arc};
+use std::{
+    any::Any,
+    fmt::{self, Debug},
+    sync::Arc,
+};
 
 use super::argument_builder::{ArgumentBuilder, ArgumentBuilderType};
 use crate::{
     arguments::ArgumentType,
+    context::CommandContext,
     exceptions::CommandSyntaxException,
     string_reader::StringReader,
-    suggestion::{Suggestions, SuggestionsBuilder},
+    suggestion::{SuggestionProvider, Suggestions, SuggestionsBuilder},
 };
 
 /// An argument node type. The `T` type parameter is the type of the argument,
 /// which can be anything.
-#[derive(Clone)]
-pub struct Argument {
+pub struct Argument<S> {
     pub name: String,
     parser: Arc<dyn ArgumentType + Send + Sync>,
+    custom_suggestions: Option<Arc<dyn SuggestionProvider<S> + Send + Sync>>,
 }
-impl Argument {
-    pub fn new(name: &str, parser: Arc<dyn ArgumentType + Send + Sync>) -> Self {
+impl<S> Argument<S> {
+    pub fn new(
+        name: &str,
+        parser: Arc<dyn ArgumentType + Send + Sync>,
+        custom_suggestions: Option<Arc<dyn SuggestionProvider<S> + Send + Sync>>,
+    ) -> Self {
         Self {
             name: name.to_string(),
             parser,
+            custom_suggestions,
         }
     }
 
@@ -27,11 +37,16 @@ impl Argument {
         self.parser.parse(reader)
     }
 
-    pub fn list_suggestions(&self, builder: SuggestionsBuilder) -> Suggestions {
-        // TODO: custom suggestions
-        // https://github.com/Mojang/brigadier/blob/master/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java#L71
-
-        self.parser.list_suggestions(builder)
+    pub fn list_suggestions(
+        &self,
+        context: CommandContext<S>,
+        builder: SuggestionsBuilder,
+    ) -> Suggestions {
+        if let Some(s) = &self.custom_suggestions {
+            s.get_suggestions(context, builder)
+        } else {
+            self.parser.list_suggestions(builder)
+        }
     }
 
     pub fn examples(&self) -> Vec<String> {
@@ -39,14 +54,14 @@ impl Argument {
     }
 }
 
-impl From<Argument> for ArgumentBuilderType {
-    fn from(argument: Argument) -> Self {
+impl<S> From<Argument<S>> for ArgumentBuilderType<S> {
+    fn from(argument: Argument<S>) -> Self {
         Self::Argument(argument)
     }
 }
 
-impl Debug for Argument {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+impl<S> Debug for Argument<S> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Argument")
             .field("name", &self.name)
             // .field("parser", &self.parser)
@@ -59,5 +74,15 @@ pub fn argument<S>(
     name: &str,
     parser: impl ArgumentType + Send + Sync + 'static,
 ) -> ArgumentBuilder<S> {
-    ArgumentBuilder::new(Argument::new(name, Arc::new(parser)).into())
+    ArgumentBuilder::new(Argument::new(name, Arc::new(parser), None).into())
+}
+
+impl<S> Clone for Argument<S> {
+    fn clone(&self) -> Self {
+        Self {
+            name: self.name.clone(),
+            parser: self.parser.clone(),
+            custom_suggestions: self.custom_suggestions.clone(),
+        }
+    }
 }
diff --git a/azalea-brigadier/src/exceptions/command_syntax_exception.rs b/azalea-brigadier/src/exceptions/command_syntax_exception.rs
index 657649b0..c9b8134f 100644
--- a/azalea-brigadier/src/exceptions/command_syntax_exception.rs
+++ b/azalea-brigadier/src/exceptions/command_syntax_exception.rs
@@ -54,23 +54,22 @@ impl CommandSyntaxException {
     }
 
     pub fn context(&self) -> Option<String> {
-        if let Some(input) = &self.input {
-            if let Some(cursor) = self.cursor {
-                let mut builder = String::new();
-                let cursor = cmp::min(input.len(), cursor);
+        if let Some(input) = &self.input
+            && let Some(cursor) = self.cursor
+        {
+            let mut builder = String::new();
+            let cursor = cmp::min(input.len(), cursor);
 
-                if cursor > CONTEXT_AMOUNT {
-                    builder.push_str("...");
-                }
+            if cursor > CONTEXT_AMOUNT {
+                builder.push_str("...");
+            }
 
-                builder.push_str(
-                    &input
-                        [(cmp::max(0, cursor as isize - CONTEXT_AMOUNT as isize) as usize)..cursor],
-                );
-                builder.push_str("<--[HERE]");
+            builder.push_str(
+                &input[(cmp::max(0, cursor as isize - CONTEXT_AMOUNT as isize) as usize)..cursor],
+            );
+            builder.push_str("<--[HERE]");
 
-                return Some(builder);
-            }
+            return Some(builder);
         }
         None
     }
diff --git a/azalea-brigadier/src/suggestion/mod.rs b/azalea-brigadier/src/suggestion/mod.rs
index a31b6837..a1c385cc 100644
--- a/azalea-brigadier/src/suggestion/mod.rs
+++ b/azalea-brigadier/src/suggestion/mod.rs
@@ -1,3 +1,4 @@
+mod suggestion_provider;
 mod suggestions;
 mod suggestions_builder;
 
@@ -12,6 +13,7 @@ use std::{
 use azalea_buf::AzaleaWrite;
 #[cfg(feature = "azalea-buf")]
 use azalea_chat::FormattedText;
+pub use suggestion_provider::SuggestionProvider;
 pub use suggestions::Suggestions;
 pub use suggestions_builder::SuggestionsBuilder;
 
diff --git a/azalea-brigadier/src/suggestion/suggestion_provider.rs b/azalea-brigadier/src/suggestion/suggestion_provider.rs
new file mode 100644
index 00000000..4c3a9f41
--- /dev/null
+++ b/azalea-brigadier/src/suggestion/suggestion_provider.rs
@@ -0,0 +1,10 @@
+use super::{Suggestions, SuggestionsBuilder};
+use crate::context::CommandContext;
+
+pub trait SuggestionProvider<S> {
+    fn get_suggestions(
+        &self,
+        context: CommandContext<S>,
+        builder: SuggestionsBuilder,
+    ) -> Suggestions;
+}
diff --git a/azalea-brigadier/src/tree/mod.rs b/azalea-brigadier/src/tree/mod.rs
index 690e5df3..8181f817 100644
--- a/azalea-brigadier/src/tree/mod.rs
+++ b/azalea-brigadier/src/tree/mod.rs
@@ -25,7 +25,7 @@ pub type Command<S> = Option<Arc<dyn Fn(&CommandContext<S>) -> i32 + Send + Sync
 /// An ArgumentBuilder that has been built.
 #[non_exhaustive]
 pub struct CommandNode<S> {
-    pub value: ArgumentBuilderType,
+    pub value: ArgumentBuilderType<S>,
 
     // this is a BTreeMap because children need to be ordered when getting command suggestions
     pub children: BTreeMap<String, Arc<RwLock<CommandNode<S>>>>,
@@ -66,7 +66,7 @@ impl<S> CommandNode<S> {
     }
     /// Gets the argument, or panics. You should use match if you're not certain
     /// about the type.
-    pub fn argument(&self) -> &Argument {
+    pub fn argument(&self) -> &Argument<S> {
         match self.value {
             ArgumentBuilderType::Argument(ref argument) => argument,
             _ => panic!("CommandNode::argument() called on non-argument node"),
@@ -214,9 +214,7 @@ impl<S> CommandNode<S> {
 
     pub fn list_suggestions(
         &self,
-        // context is here because that's how it is in mojang's brigadier, but we haven't
-        // implemented custom suggestions yet so this is unused rn
-        _context: CommandContext<S>,
+        context: CommandContext<S>,
         builder: SuggestionsBuilder,
     ) -> Suggestions {
         match &self.value {
@@ -231,7 +229,7 @@ impl<S> CommandNode<S> {
                     Suggestions::default()
                 }
             }
-            ArgumentBuilderType::Argument(argument) => argument.list_suggestions(builder),
+            ArgumentBuilderType::Argument(argument) => argument.list_suggestions(context, builder),
         }
     }
 }
@@ -239,7 +237,7 @@ impl<S> CommandNode<S> {
 impl<S> Debug for CommandNode<S> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CommandNode")
-            .field("value", &self.value)
+            // .field("value", &self.value)
             .field("children", &self.children)
             .field("command", &self.command.is_some())
             // .field("requirement", &self.requirement)
diff --git a/azalea-brigadier/tests/bevy_app_usage.rs b/azalea-brigadier/tests/bevy_app_usage.rs
index 24a2e369..e962d7d1 100644
--- a/azalea-brigadier/tests/bevy_app_usage.rs
+++ b/azalea-brigadier/tests/bevy_app_usage.rs
@@ -1,11 +1,6 @@
-use std::sync::Arc;
-
-use azalea_brigadier::{
-    arguments::integer_argument_type::integer,
-    builder::{literal_argument_builder::literal, required_argument_builder::argument},
-    command_dispatcher::CommandDispatcher,
-    context::CommandContext,
-};
+use std::{ops::Deref, sync::Arc};
+
+use azalea_brigadier::prelude::*;
 use bevy_app::App;
 use bevy_ecs::{prelude::*, system::RunSystemOnce};
 use parking_lot::Mutex;
@@ -145,8 +140,7 @@ impl DispatchStorage {
     ///
     /// Spawns a number of entities with the [`SpawnedEntity`] component.
     fn command_spawn_entity_num(context: &CommandContext<WorldAccessor>) -> i32 {
-        let num = context.argument("entities").unwrap();
-        let num = *num.downcast_ref::<i32>().unwrap();
+        let num = get_integer(context, "entities").unwrap();
 
         for _ in 0..num {
             context.source.lock().spawn(SpawnedEntity);
@@ -187,7 +181,7 @@ impl WorldAccessor {
 struct SpawnedEntity;
 
 /// Implemented for convenience.
-impl std::ops::Deref for WorldAccessor {
+impl Deref for WorldAccessor {
     type Target = Arc<Mutex<World>>;
     fn deref(&self) -> &Self::Target {
         &self.world