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)