feature(api): Rewrite ComponentFlattener to remove recursion and add a configurable maximum nesting depth

Author: kezz

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

@@ -31,13 +31,21 @@
 import net.kyori.adventure.util.Buildable;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.Range;
 
 /**
  * A 'flattener' to convert a component tree to a linear string for display.
  *
  * @since 4.7.0
  */
 public interface ComponentFlattener extends Buildable<ComponentFlattener, ComponentFlattener.Builder> {
+  /**
+   * A constant representing a flattener with no limit on nested flatten calls.
+   *
+   * @since 4.22.0
+   */
+  int NO_NESTING_LIMIT = -1;
+
   /**
    * Create a new builder for a flattener.
    *
@@ -123,5 +131,16 @@ interface Builder extends AbstractBuilder<ComponentFlattener>, Buildable.Builder
      * @since 4.7.0
      */
     @NotNull Builder unknownMapper(final @Nullable Function<Component, String> converter);
+
+    /**
+     * Sets the limit of nested flatten calls.
+     *
+     * <p>The default value is {@link #NO_NESTING_LIMIT}, which means there is no limit on nesting.</p>
+     *
+     * @param limit the new limit (must be a positive integer, or {@link #NO_NESTING_LIMIT})
+     * @return this builder
+     * @since 4.22.0
+     */
+    @NotNull Builder nestingLimit(final @Range(from = 1, to = Integer.MAX_VALUE) int limit);
   }
 }

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

@@ -23,6 +23,9 @@
  */
 package net.kyori.adventure.text.flattener;
 
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.List;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -36,6 +39,7 @@
 import net.kyori.adventure.util.InheritanceAwareMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.Range;
 
 import static java.util.Objects.requireNonNull;
 
@@ -64,94 +68,136 @@ final class ComponentFlattenerImpl implements ComponentFlattener {
 
   private final InheritanceAwareMap<Component, Handler> flatteners;
   private final Function<Component, String> unknownHandler;
+  private final int maxNestedDepth;
 
-  ComponentFlattenerImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler) {
+  ComponentFlattenerImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler, final int maxNestedDepth) {
     this.flatteners = flatteners;
     this.unknownHandler = unknownHandler;
+    this.maxNestedDepth = maxNestedDepth;
+  }
+
+  private static final class StackEntry {
+    final Component component;
+    final int depth;
+
+    StackEntry(final Component component, final int depth) {
+      this.component = component;
+      this.depth = depth;
+    }
   }
 
   @Override
   public void flatten(final @NotNull Component input, final @NotNull FlattenerListener listener) {
-    this.flatten0(input, listener, 0);
+    this.flatten0(input, listener, 0, 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 int depth, final int nestedDepth) {
     requireNonNull(input, "input");
     requireNonNull(listener, "listener");
     if (input == Component.empty()) return;
-    if (depth > MAX_DEPTH) {
-      throw new IllegalStateException("Exceeded maximum depth of " + MAX_DEPTH + " while attempting to flatten components!");
+
+    if (this.maxNestedDepth != ComponentFlattener.NO_NESTING_LIMIT && nestedDepth > this.maxNestedDepth) {
+      throw new IllegalStateException("Exceeded maximum nesting depth of " + this.maxNestedDepth + " while attempting to flatten components!");
     }
 
-    final @Nullable Handler flattener = this.flattener(input);
-    final Style inputStyle = input.style();
+    final Deque<StackEntry> componentStack = new ArrayDeque<>();
+    final Deque<Style> styleStack = new ArrayDeque<>();
+
+    // Push the starting component.
+    componentStack.push(new StackEntry(input, depth));
+
+    while (!componentStack.isEmpty()) {
+      final StackEntry entry = componentStack.pop();
+      final int currentDepth = entry.depth;
+
+      if (currentDepth > MAX_DEPTH) {
+        throw new IllegalStateException("Exceeded maximum depth of " + MAX_DEPTH + " while attempting to flatten components!");
+      }
+
+      final Component component = entry.component;
+      final @Nullable Handler flattener = this.flattener(component);
+      final Style componentStyle = component.style();
+
+      // Push the style to both the listener and the stack (so we can pop later).
+      listener.pushStyle(componentStyle);
+      styleStack.push(componentStyle);
 
-    listener.pushStyle(inputStyle);
-    try {
       if (flattener != null) {
-        flattener.handle(this, input, listener, depth + 1);
+        flattener.handle(this, component, listener, currentDepth, nestedDepth);
       }
 
-      if (!input.children().isEmpty() && listener.shouldContinue()) {
-        for (final Component child : input.children()) {
-          this.flatten0(child, listener, depth + 1);
+      if (!component.children().isEmpty() && listener.shouldContinue()) {
+        // Push any children onto the stack in reverse order so they are popped in the right order.
+        final List<Component> children = component.children();
+        for (int i = children.size() - 1; i >= 0; i--) {
+          componentStack.push(new StackEntry(children.get(i), currentDepth + 1));
         }
+      } else {
+        // If there are no children, we pop the latest style to go back "up" the tree.
+        final Style style = styleStack.pop();
+        listener.popStyle(style);
       }
-    } finally {
-      listener.popStyle(inputStyle);
+    }
+
+    // Pop any remaining styles at the end.
+    while (!styleStack.isEmpty()) {
+      final Style style = styleStack.pop();
+      listener.popStyle(style);
     }
   }
 
   private <T extends Component> @Nullable Handler flattener(final T test) {
     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, depth, nestedDepth) -> listener.component(this.unknownHandler.apply(component));
     } else {
       return flattener;
     }
   }
 
   @Override
   public ComponentFlattener.@NotNull Builder toBuilder() {
-    return new BuilderImpl(this.flatteners, this.unknownHandler);
+    return new BuilderImpl(this.flatteners, this.unknownHandler, this.maxNestedDepth);
   }
 
   // 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 int depth, final int nestedDepth);
   }
 
   static final class BuilderImpl implements Builder {
     private final InheritanceAwareMap.Builder<Component, Handler> flatteners;
     private @Nullable Function<Component, String> unknownHandler;
+    private int maxNestedDepth = ComponentFlattener.NO_NESTING_LIMIT;
 
     BuilderImpl() {
       this.flatteners = InheritanceAwareMap.<Component, Handler>builder().strict(true);
     }
 
-    BuilderImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler) {
+    BuilderImpl(final InheritanceAwareMap<Component, Handler> flatteners, final @Nullable Function<Component, String> unknownHandler, final int maxNestedDepth) {
       this.flatteners = InheritanceAwareMap.builder(flatteners).strict(true);
       this.unknownHandler = unknownHandler;
+      this.maxNestedDepth = maxNestedDepth;
     }
 
     @Override
     public @NotNull ComponentFlattener build() {
-      return new ComponentFlattenerImpl(this.flatteners.build(), this.unknownHandler);
+      return new ComponentFlattenerImpl(this.flatteners.build(), this.unknownHandler, this.maxNestedDepth);
     }
 
     @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, depth, nestedDepth) -> 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, depth, nestedDepth) -> converter.accept((T) component, c -> self.flatten0(c, listener, depth, nestedDepth + 1)));
       return this;
     }
 
@@ -160,5 +206,13 @@ static final class BuilderImpl implements Builder {
       this.unknownHandler = converter;
       return this;
     }
+
+    @Override
+    public @NotNull Builder nestingLimit(final @Range(from = 1, to = Integer.MAX_VALUE) int limit) {
+      // noinspection ConstantValue (the Range annotation tells IDEA this will never happen, but API users could ignore that)
+      if (limit != ComponentFlattener.NO_NESTING_LIMIT && limit < 1) throw new IllegalArgumentException("limit must be positive or ComponentFlattener.NO_NESTING_LIMIT");
+      this.maxNestedDepth = limit;
+      return this;
+    }
   }
 }

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

@@ -27,20 +27,25 @@
 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;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertIterableEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 class ComponentFlattenerTest {
   static class TrackingFlattener implements FlattenerListener {
@@ -288,4 +293,64 @@ void testEarlyExit() {
       .assertPushesAndPops(3)
       .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)));
+      }
+    })
+    .nestingLimit(4)
+    .build();
+
+  public static Component createNestedComponent(final int depth, final String finalText) {
+    Component component = Component.text(finalText);
+
+    for (int i = 0; i < depth; i++) {
+      component = Component.translatable("%1$s%1$s%1$s", component);
+    }
+
+    return component;
+  }
+
+  @Test
+  void testGiantComponent() {
+    final Component component = createNestedComponent(34, "only 34?!");
+    final StringBuilder sb = new StringBuilder();
+    final Exception exception = assertThrows(IllegalStateException.class, () -> PAPERS_WEIRD_FLATTENER.flatten(component, sb::append));
+    assertTrue(exception.getMessage().startsWith("Exceeded maximum nesting depth"));
+  }
 }