rewrite the entire thing ig

Author: kezz

api/src/main/java/net/kyori/adventure/text/flattener/ComponentFlattener.java

@@ -90,6 +90,14 @@ public interface ComponentFlattener extends Buildable<ComponentFlattener, Compon
    * @since 4.7.0
    */
   interface Builder extends AbstractBuilder<ComponentFlattener>, Buildable.Builder<ComponentFlattener> {
+    /**
+     * A constant value for {@link #maximumComplexity(int)} that indicates that the maximum complexity is unlimited.
+     *
+     * @see #maximumComplexity(int)
+     * @since 4.22.0
+     */
+    int UNLIMITED_COMPLEXITY = -1;
+
     /**
      * Register a type of component to be handled.
      *
@@ -125,12 +133,16 @@ interface Builder extends AbstractBuilder<ComponentFlattener>, Buildable.Builder
     @NotNull Builder unknownMapper(final @Nullable Function<Component, String> converter);
 
     /**
-     * Sets the maximum depth of the flattening.
+     * Sets the maximum complexity of the flattening.
+     *
+     * <p>Complexity is defined as the number of recursive flatten operations that take place.
+     * This includes both depth and breadth as an additional count.</p>
      *
-     * @param maximumDepth the maximum depth
+     * @param maximumComplexity the maximum complexity, defaulting to {@link #UNLIMITED_COMPLEXITY}
      * @return this builder
+     * @see #UNLIMITED_COMPLEXITY
      * @since 4.22.0
      */
-    @NotNull Builder maximumDepth(final int maximumDepth);
+    @NotNull Builder maximumComplexity(final int maximumComplexity);
   }
 }

api/src/main/java/net/kyori/adventure/text/flattener/ComponentFlattenerImpl.java

@@ -60,27 +60,32 @@ final class ComponentFlattenerImpl implements ComponentFlattener {
     .mapper(TextComponent.class, TextComponent::content)
     .build();
 
+  private static final int MAX_DEPTH = 512;
+
   private final InheritanceAwareMap<Component, Handler> flatteners;
   private final Function<Component, String> unknownHandler;
-  private final int maximumDepth;
+  private final int maximumComplexity;
 
-  ComponentFlattenerImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler, final int maximumDepth) {
+  ComponentFlattenerImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler, final int maximumComplexity) {
     this.flatteners = flatteners;
     this.unknownHandler = unknownHandler;
-    this.maximumDepth = maximumDepth;
+    this.maximumComplexity = maximumComplexity;
   }
 
   @Override
   public void flatten(final @NotNull Component input, final @NotNull FlattenerListener listener) {
-    this.flatten0(input, listener, 0);
+    this.flatten0(input, listener, new State(), 0);
   }
 
-  private void flatten0(final @NotNull Component input, final @NotNull FlattenerListener listener, final int depth) {
+  private void flatten0(final @NotNull Component input, final @NotNull FlattenerListener listener, final @NotNull State state, final int depth) {
     requireNonNull(input, "input");
     requireNonNull(listener, "listener");
     if (input == Component.empty()) return;
-    if (depth > this.maximumDepth) {
-      throw new IllegalStateException("Exceeded maximum depth of " + this.maximumDepth + " while attempting to flatten components!");
+    if (this.maximumComplexity != Builder.UNLIMITED_COMPLEXITY && state.complexity > this.maximumComplexity) {
+      throw new IllegalStateException("Exceeded maximum complexity of " + this.maximumComplexity + " while attempting to flatten components!");
+    }
+    if (depth >= MAX_DEPTH) {
+      throw new IllegalStateException("Exceeded maximum depth of " + MAX_DEPTH + " while attempting to flatten components!");
     }
 
     final @Nullable Handler flattener = this.flattener(input);
@@ -89,12 +94,14 @@ private void flatten0(final @NotNull Component input, final @NotNull FlattenerLi
     listener.pushStyle(inputStyle);
     try {
       if (flattener != null) {
-        flattener.handle(this, input, listener, depth + 1);
+        state.complexity++;
+        flattener.handle(this, input, listener, state, depth + 1);
       }
 
       if (!input.children().isEmpty() && listener.shouldContinue()) {
         for (final Component child : input.children()) {
-          this.flatten0(child, listener, depth + 1);
+          state.complexity++;
+          this.flatten0(child, listener, state, depth + 1);
         }
       }
     } finally {
@@ -106,57 +113,59 @@ private void flatten0(final @NotNull Component input, final @NotNull FlattenerLi
     final Handler flattener = this.flatteners.get(test.getClass());
 
     if (flattener == null && this.unknownHandler != null) {
-      return (self, component, listener, depth) -> listener.component(this.unknownHandler.apply(component));
+      return (self, component, listener, state, depth) -> listener.component(this.unknownHandler.apply(component));
     } else {
       return flattener;
     }
   }
 
   @Override
   public ComponentFlattener.@NotNull Builder toBuilder() {
-    return new BuilderImpl(this.flatteners, this.unknownHandler, this.maximumDepth);
+    return new BuilderImpl(this.flatteners, this.unknownHandler, this.maximumComplexity);
+  }
+
+  static final class State {
+    int complexity = 0;
   }
 
   // A function that allows nesting other flatten operations
   @FunctionalInterface
   interface Handler {
-    void handle(final ComponentFlattenerImpl self, final Component input, final FlattenerListener listener, final int depth);
+    void handle(final ComponentFlattenerImpl self, final Component input, final FlattenerListener listener, final State state, final int depth);
   }
 
   static final class BuilderImpl implements Builder {
-    private static final int DEFAULT_MAX_DEPTH = 512;
-
     private final InheritanceAwareMap.Builder<Component, Handler> flatteners;
     private @Nullable Function<Component, String> unknownHandler;
-    private int maximumDepth;
+    private int maximumComplexity;
 
     BuilderImpl() {
       this.flatteners = InheritanceAwareMap.<Component, Handler>builder().strict(true);
-      this.maximumDepth = DEFAULT_MAX_DEPTH;
+      this.maximumComplexity = UNLIMITED_COMPLEXITY;
     }
 
-    BuilderImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler, final int maximumDepth) {
+    BuilderImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler, final int maximumComplexity) {
       this.flatteners = InheritanceAwareMap.builder(flatteners).strict(true);
       this.unknownHandler = unknownHandler;
-      this.maximumDepth = maximumDepth;
+      this.maximumComplexity = maximumComplexity;
     }
 
     @Override
     public @NotNull ComponentFlattener build() {
-      return new ComponentFlattenerImpl(this.flatteners.build(), this.unknownHandler, this.maximumDepth);
+      return new ComponentFlattenerImpl(this.flatteners.build(), this.unknownHandler, this.maximumComplexity);
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public <T extends Component> ComponentFlattener.@NotNull Builder mapper(final @NotNull Class<T> type, final @NotNull Function<T, String> converter) {
-      this.flatteners.put(type, (self, component, listener, depth) -> listener.component(converter.apply((T) component)));
+      this.flatteners.put(type, (self, component, listener, state, depth) -> listener.component(converter.apply((T) component)));
       return this;
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public <T extends Component> ComponentFlattener.@NotNull Builder complexMapper(final @NotNull Class<T> type, final @NotNull BiConsumer<T, Consumer<Component>> converter) {
-      this.flatteners.put(type, (self, component, listener, depth) -> converter.accept((T) component, c -> self.flatten0(c, listener, depth)));
+      this.flatteners.put(type, (self, component, listener, state, depth) -> converter.accept((T) component, c -> self.flatten0(c, listener, state, depth)));
       return this;
     }
 
@@ -167,9 +176,11 @@ static final class BuilderImpl implements Builder {
     }
 
     @Override
-    public @NotNull Builder maximumDepth(final int maximumDepth) {
-      if (maximumDepth <= 0) throw new IllegalArgumentException("maxDepth must be greater than 0, was " + maximumDepth);
-      this.maximumDepth = maximumDepth;
+    public @NotNull Builder maximumComplexity(final int maximumComplexity) {
+      if (maximumComplexity != UNLIMITED_COMPLEXITY && maximumComplexity <= 0) {
+        throw new IllegalArgumentException("maximumComplexity must be greater than 0, was " + maximumComplexity);
+      }
+      this.maximumComplexity = maximumComplexity;
       return this;
     }
   }

api/src/test/java/net/kyori/adventure/text/flattener/ComponentFlattenerTest.java

@@ -27,14 +27,18 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Locale;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import net.kyori.adventure.text.BlockNBTComponent;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.NBTComponent;
 import net.kyori.adventure.text.TranslatableComponent;
+import net.kyori.adventure.text.TranslationArgument;
 import net.kyori.adventure.text.format.NamedTextColor;
 import net.kyori.adventure.text.format.Style;
 import net.kyori.adventure.text.format.TextDecoration;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
@@ -289,6 +293,48 @@ void testEarlyExit() {
       .assertContents("Hello", "How are you?", "Not great");
   }
 
+  private static final Pattern PAPERS_WEIRD_LOCALIZATION_PATTERN = Pattern.compile("%(?:(\\d+)\\$)?s");
+  public static final ComponentFlattener PAPERS_WEIRD_FLATTENER = ComponentFlattener.basic().toBuilder()
+    .complexMapper(TranslatableComponent.class, (translatable, consumer) -> {
+      final String key = translatable.key();
+      final Matcher matcher = PAPERS_WEIRD_LOCALIZATION_PATTERN.matcher(key);
+      final List<TranslationArgument> args = translatable.arguments();
+      int argPosition = 0;
+      int lastIdx = 0;
+      while (matcher.find()) {
+        // append prior
+        if (lastIdx < matcher.start()) {
+          consumer.accept(Component.text(key.substring(lastIdx, matcher.start())));
+        }
+        lastIdx = matcher.end();
+
+        final @Nullable String argIdx = matcher.group(1);
+        // calculate argument position
+        if (argIdx != null) {
+          try {
+            final int idx = Integer.parseInt(argIdx) - 1;
+            if (idx < args.size()) {
+              consumer.accept(args.get(idx).asComponent());
+            }
+          } catch (final NumberFormatException ex) {
+            // ignore, drop the format placeholder
+          }
+        } else {
+          final int idx = argPosition++;
+          if (idx < args.size()) {
+            consumer.accept(args.get(idx).asComponent());
+          }
+        }
+      }
+
+      // append tail
+      if (lastIdx < key.length()) {
+        consumer.accept(Component.text(key.substring(lastIdx)));
+      }
+    })
+    .maximumComplexity(32)
+    .build();
+
   public static Component createNestedComponent(final int depth, final String finalText) {
     Component component = Component.text(finalText);
 
@@ -303,6 +349,6 @@ public static Component createNestedComponent(final int depth, final String fina
   void testGiantComponent() {
     final Component component = createNestedComponent(34, "only 34?!");
     final StringBuilder sb = new StringBuilder();
-    ComponentFlattener.basic().flatten(component, sb::append);
+    assertThrows(IllegalStateException.class, () -> PAPERS_WEIRD_FLATTENER.flatten(component, sb::append));
   }
 }