From 339e54c0b2ccf31029857429cc8767d7ed174a07 Mon Sep 17 00:00:00 2001
From: Jesse Costello-Good <jcg@google.com>
Date: Fri, 20 Oct 2023 09:38:00 -0700
Subject: [PATCH] BFS with stream successors; avoid collecting to Array just to
 add to queue.

PiperOrigin-RevId: 575242725
---
 .../template/soy/internal/util/TreeStreams.java    | 14 +++++++++++++-
 .../google/template/soy/jssrc/dsl/CodeChunks.java  |  7 ++-----
 .../google/template/soy/jssrc/dsl/Expression.java  |  6 +++---
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/java/src/com/google/template/soy/internal/util/TreeStreams.java b/java/src/com/google/template/soy/internal/util/TreeStreams.java
index 47ad38b0bd..bfa88ebcfd 100644
--- a/java/src/com/google/template/soy/internal/util/TreeStreams.java
+++ b/java/src/com/google/template/soy/internal/util/TreeStreams.java
@@ -27,6 +27,7 @@
 import java.util.List;
 import java.util.Spliterator;
 import java.util.Spliterators.AbstractSpliterator;
+import java.util.function.BiConsumer;
 import java.util.function.BiPredicate;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -73,6 +74,17 @@ public boolean tryAdvance(Consumer<? super T> action) {
    */
   public static <T> Stream<? extends T> breadthFirst(
       T root, Function<T, Iterable<? extends T>> successors) {
+    return breadthFirstInternal(
+        root, (queue, next) -> Iterables.addAll(queue, successors.apply(next)));
+  }
+
+  public static <T> Stream<? extends T> breadthFirstWithStream(
+      T root, Function<T, Stream<? extends T>> successors) {
+    return breadthFirstInternal(root, (queue, next) -> successors.apply(next).forEach(queue::add));
+  }
+
+  private static <T> Stream<? extends T> breadthFirstInternal(
+      T root, BiConsumer<Deque<T>, T> pusher) {
     Deque<T> queue = new ArrayDeque<>();
     queue.add(root);
     return StreamSupport.stream(
@@ -87,7 +99,7 @@ public boolean tryAdvance(Consumer<? super T> action) {
             if (next == null) {
               return false;
             }
-            Iterables.addAll(queue, successors.apply(next));
+            pusher.accept(queue, next);
             action.accept(next);
             return true;
           }
diff --git a/java/src/com/google/template/soy/jssrc/dsl/CodeChunks.java b/java/src/com/google/template/soy/jssrc/dsl/CodeChunks.java
index 1c1e3ed3ea..e2478413de 100644
--- a/java/src/com/google/template/soy/jssrc/dsl/CodeChunks.java
+++ b/java/src/com/google/template/soy/jssrc/dsl/CodeChunks.java
@@ -120,14 +120,11 @@ public static Stream<CodeChunk> flatten(Stream<CodeChunk> chunk) {
   }
 
   public static Stream<? extends CodeChunk> breadthFirst(CodeChunk root) {
-    return TreeStreams.breadthFirst(root, c -> c.childrenStream().collect(Collectors.toList()));
+    return TreeStreams.breadthFirstWithStream(root, CodeChunk::childrenStream);
   }
 
   public static Stream<? extends CodeChunk> breadthFirst(List<? extends CodeChunk> roots) {
     return roots.stream()
-        .flatMap(
-            root ->
-                TreeStreams.<CodeChunk>breadthFirst(
-                    root, c -> c.childrenStream().collect(Collectors.toList())));
+        .flatMap(root -> TreeStreams.breadthFirstWithStream(root, CodeChunk::childrenStream));
   }
 }
diff --git a/java/src/com/google/template/soy/jssrc/dsl/Expression.java b/java/src/com/google/template/soy/jssrc/dsl/Expression.java
index a4da8f8bed..7e9502c5fd 100644
--- a/java/src/com/google/template/soy/jssrc/dsl/Expression.java
+++ b/java/src/com/google/template/soy/jssrc/dsl/Expression.java
@@ -365,13 +365,13 @@ final void doFormatInitialStatements(FormattingContext ctx) {
   }
 
   private Stream<Statement> initialStatementsStream() {
-    return TreeStreams.<CodeChunk>breadthFirst(
+    return TreeStreams.<CodeChunk>breadthFirstWithStream(
             this,
             c -> {
               if (c instanceof Expression && !(c instanceof InitialStatementsScope)) {
-                return c.childrenStream().collect(toImmutableList());
+                return c.childrenStream();
               }
-              return ImmutableList.of();
+              return Stream.of();
             })
         .filter(HasInitialStatements.class::isInstance)
         .map(HasInitialStatements.class::cast)