From c225f92032231dcc244472157245af6ea53e486c Mon Sep 17 00:00:00 2001
From: Vankka <vankka.main@gmail.com>
Date: Sat, 20 Jul 2024 15:48:27 +0300
Subject: [PATCH] Fix recursive translation

---
 .../common/component/ComponentFactory.java    | 46 +++++++++++++++----
 .../translation/TranslationRegistry.java      | 24 ++++++----
 2 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/common/src/main/java/com/discordsrv/common/component/ComponentFactory.java b/common/src/main/java/com/discordsrv/common/component/ComponentFactory.java
index 7d661fcf..da13a9de 100644
--- a/common/src/main/java/com/discordsrv/common/component/ComponentFactory.java
+++ b/common/src/main/java/com/discordsrv/common/component/ComponentFactory.java
@@ -25,6 +25,8 @@ import com.discordsrv.common.DiscordSRV;
 import com.discordsrv.common.component.renderer.DiscordSRVMinecraftRenderer;
 import com.discordsrv.common.component.translation.Translation;
 import com.discordsrv.common.component.translation.TranslationRegistry;
+import com.discordsrv.common.logging.Logger;
+import com.discordsrv.common.logging.NamedLogger;
 import dev.vankka.enhancedlegacytext.EnhancedLegacyText;
 import dev.vankka.mcdiscordreserializer.discord.DiscordSerializer;
 import dev.vankka.mcdiscordreserializer.discord.DiscordSerializerOptions;
@@ -39,6 +41,9 @@ import net.kyori.adventure.text.serializer.plain.PlainTextComponentSerializer;
 import net.kyori.ansi.ColorLevel;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.HashSet;
+import java.util.Set;
+
 public class ComponentFactory implements MinecraftComponentFactory {
 
     public static final Class<?> UNRELOCATED_ADVENTURE_COMPONENT;
@@ -52,6 +57,8 @@ public class ComponentFactory implements MinecraftComponentFactory {
     }
 
     private final DiscordSRV discordSRV;
+    private final Logger logger;
+
     private final MinecraftSerializer minecraftSerializer;
     private final DiscordSerializer discordSerializer;
     private final PlainTextComponentSerializer plainSerializer;
@@ -62,6 +69,8 @@ public class ComponentFactory implements MinecraftComponentFactory {
 
     public ComponentFactory(DiscordSRV discordSRV) {
         this.discordSRV = discordSRV;
+        this.logger = new NamedLogger(discordSRV, "COMPONENT_FACTORY");
+
         this.minecraftSerializer = new MinecraftSerializer(
                 MinecraftSerializerOptions.defaults()
                         .addRenderer(new DiscordSRVMinecraftRenderer(discordSRV))
@@ -83,19 +92,40 @@ public class ComponentFactory implements MinecraftComponentFactory {
                 .build();
     }
 
+    private final ThreadLocal<Set<String>> translationHistory = new ThreadLocal<>();
     private String provideTranslation(TranslatableComponent component) {
-        Translation translation = translationRegistry.lookup(discordSRV.defaultLocale(), component.key());
+        Set<String> history = translationHistory.get();
+        if (history == null) {
+            history = new HashSet<>();
+        }
+
+        String key = component.key();
+        if (history.contains(key)) {
+            // Prevent infinite loop here
+            logger.debug("Preventing recursive translation: " + key);
+            return key;
+        }
+
+        Translation translation = translationRegistry.lookup(discordSRV.defaultLocale(), key);
         if (translation == null) {
             return null;
         }
 
-        return translation.translate(
-                component.arguments()
-                        .stream()
-                        // Prevent infinite loop here by using the default PlainTextSerializer
-                        .map(argument -> PlainTextComponentSerializer.plainText().serialize(argument.asComponent()))
-                        .toArray(Object[]::new)
-        );
+        try {
+            history.add(key);
+            translationHistory.set(history);
+
+            return translation.translate(
+                    component.arguments()
+                            .stream()
+                            .map(argument -> plainSerializer().serialize(argument.asComponent()))
+                            .toArray(Object[]::new)
+            );
+        } finally {
+            Set<String> newHistory = translationHistory.get();
+            newHistory.remove(key);
+            translationHistory.set(newHistory.isEmpty() ? null : newHistory);
+        }
     }
 
     @Override
diff --git a/common/src/main/java/com/discordsrv/common/component/translation/TranslationRegistry.java b/common/src/main/java/com/discordsrv/common/component/translation/TranslationRegistry.java
index b61d97fe..50c532ee 100644
--- a/common/src/main/java/com/discordsrv/common/component/translation/TranslationRegistry.java
+++ b/common/src/main/java/com/discordsrv/common/component/translation/TranslationRegistry.java
@@ -20,7 +20,6 @@ package com.discordsrv.common.component.translation;
 
 import org.jetbrains.annotations.Nullable;
 
-import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.Locale;
 import java.util.Map;
@@ -44,16 +43,23 @@ public class TranslationRegistry {
     }
 
     @Nullable
-    public Translation lookup(Locale locale, String key) {
+    private Translation getTranslationOrNull(Locale locale, String key) {
+        Map<String, Translation> localeMap;
         synchronized (translations) {
-            return translations.getOrDefault(
-                    locale, // Try the suggested locale first
-                    translations.getOrDefault(
-                            DEFAULT_LOCALE, // Then try the default locale
-                            Collections.emptyMap() // Then fail
-                    )
-            ).get(key);
+            localeMap = translations.get(locale);
         }
+
+        return localeMap != null ? localeMap.get(key) : null;
+    }
+
+    @Nullable
+    public Translation lookup(Locale locale, String key) {
+        Translation translation = getTranslationOrNull(locale, key);
+        if (translation != null) {
+            return translation;
+        }
+
+        return getTranslationOrNull(DEFAULT_LOCALE, key);
     }
 
     public void clear() {