Fix custom suggestions not being triggered when transformed (#24)

Author: alexstaeding

brigadier/src/jvmMain/kotlin/org/anvilpowered/kbrig/brigadier/BrigadierConverter.kt

@@ -23,6 +23,7 @@ import org.anvilpowered.kbrig.context.CommandContext
 import org.anvilpowered.kbrig.context.StringRange
 import org.anvilpowered.kbrig.coroutine.coroutineToFuture
 import org.anvilpowered.kbrig.suggestion.Suggestion
+import org.anvilpowered.kbrig.suggestion.SuggestionProvider
 import org.anvilpowered.kbrig.suggestion.Suggestions
 import org.anvilpowered.kbrig.suggestion.SuggestionsBuilder
 import org.anvilpowered.kbrig.tree.ArgumentCommandNode
@@ -42,6 +43,7 @@ import com.mojang.brigadier.arguments.StringArgumentType as BrigadierStringArgum
 import com.mojang.brigadier.context.CommandContext as BrigadierCommandContext
 import com.mojang.brigadier.context.StringRange as BrigadierStringRange
 import com.mojang.brigadier.suggestion.Suggestion as BrigadierSuggestion
+import com.mojang.brigadier.suggestion.SuggestionProvider as BrigadierSuggestionProvider
 import com.mojang.brigadier.suggestion.Suggestions as BrigadierSuggestions
 import com.mojang.brigadier.suggestion.SuggestionsBuilder as BrigadierSuggestionsBuilder
 import com.mojang.brigadier.tree.ArgumentCommandNode as BrigadierArgumentCommandNode
@@ -71,7 +73,7 @@ fun <S, T> ArgumentCommandNode<S, T>.toBrigadier(): BrigadierArgumentCommandNode
         redirect?.toBrigadier(),
         null,
         forks,
-        null,
+        customSuggestions?.toBrigadier(),
     ).withChildrenFrom(this)
 }
 
@@ -140,6 +142,15 @@ private fun <T> ArgumentType<T>.toDefaultBrigadierArgument(): BrigadierArgumentT
     }
 }
 
+private fun <S> SuggestionProvider<S>.toBrigadier(): BrigadierSuggestionProvider<S> {
+    val original = this
+    return BrigadierSuggestionProvider { context, builder ->
+        coroutineToFuture {
+            original.getSuggestions(context.toKBrig(), builder.toKBrig()).toBrigadier()
+        }
+    }
+}
+
 private fun Suggestions.toBrigadier(): BrigadierSuggestions = BrigadierSuggestions(range.toBrigadier(), list.map { it.toBrigadier() })
 private fun Suggestion.toBrigadier(): BrigadierSuggestion = BrigadierSuggestion(range.toBrigadier(), text)
 private fun StringRange.toBrigadier(): BrigadierStringRange = BrigadierStringRange(start, end)

core/src/commonMain/kotlin/org/anvilpowered/kbrig/tree/ArgumentCommandNode.kt

@@ -23,7 +23,7 @@ class ArgumentCommandNode<S, T>(
     redirect: CommandNode<S>? = null,
     forks: Boolean = false,
     children: Map<String, CommandNode<S>>,
-    private val customSuggestions: SuggestionProvider<S>? = null,
+    val customSuggestions: SuggestionProvider<S>? = null,
 ) : CommandNode<S>(name, command, requirement, redirect, forks, children) {
 
     override val usageText: String

core/src/commonMain/kotlin/org/anvilpowered/kbrig/tree/SourceConverter.kt

@@ -11,6 +11,7 @@ package org.anvilpowered.kbrig.tree
 
 import org.anvilpowered.kbrig.Command
 import org.anvilpowered.kbrig.context.CommandContext
+import org.anvilpowered.kbrig.suggestion.SuggestionProvider
 import kotlin.jvm.JvmName
 
 /**
@@ -25,6 +26,7 @@ fun <S, R, T> ArgumentCommandNode<S, T>.mapSource(mapper: (R) -> S): ArgumentCom
         redirect?.mapSource(mapper),
         forks,
         children.mapValues { (_, child) -> child.mapSource(mapper) },
+        customSuggestions?.mapSource(mapper),
     )
 
 /**
@@ -72,6 +74,9 @@ private fun <S, R> Command<S>.mapSource(mapper: (R) -> S): Command<R> =
  */
 private fun <S, R> ((S) -> Boolean).mapSource(mapper: (R) -> S): (R) -> Boolean = { this(mapper(it)) }
 
+private fun <S, R> SuggestionProvider<S>.mapSource(mapper: (R) -> S): SuggestionProvider<R> =
+    SuggestionProvider { context, builder -> getSuggestions(context.mapToOriginalSource(mapper), builder) }
+
 /**
  * Converts the given [CommandContext] with source type [R] to an [CommandContext] with source type [S].
  */