BFS with stream successors; avoid collecting to Array just to add to …

…queue. PiperOrigin-RevId: 575886387
google · Oct 23, 2023 · 1ca3f4f · 1ca3f4f
1 parent f92f148
commit 1ca3f4f
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 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
@@ -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
@@ -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
@@ -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)