Fix filter/map macro to be linear in time and space complexity

PiperOrigin-RevId: 781245388

Author: l46kok

runtime/src/main/java/dev/cel/runtime/BUILD.bazel

@@ -272,6 +272,7 @@ java_library(
     exports = [":base"],
     deps = [
         ":base",
+        ":concatenated_list_view",
         ":dispatcher",
         ":evaluation_exception",
         ":evaluation_exception_builder",
@@ -306,6 +307,7 @@ cel_android_library(
     visibility = ["//visibility:private"],
     deps = [
         ":base_android",
+        ":concatenated_list_view",
         ":dispatcher_android",
         ":evaluation_exception",
         ":evaluation_exception_builder",
@@ -398,6 +400,7 @@ cel_android_library(
     tags = [
     ],
     deps = [
+        ":concatenated_list_view",
         "//common:error_codes",
         "//common:options",
         "//common:runtime_exception",
@@ -419,6 +422,7 @@ java_library(
     tags = [
     ],
     deps = [
+        ":concatenated_list_view",
         "//common:error_codes",
         "//common:options",
         "//common:runtime_exception",
@@ -1094,3 +1098,10 @@ cel_android_library(
         "@maven_android//:com_google_guava_guava",
     ],
 )
+
+java_library(
+    name = "concatenated_list_view",
+    srcs = ["ConcatenatedListView.java"],
+    # used_by_android
+    visibility = ["//visibility:private"],
+)

runtime/src/main/java/dev/cel/runtime/ConcatenatedListView.java

@@ -0,0 +1,102 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License aj
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.runtime;
+
+import java.util.AbstractList;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+/**
+ * A custom list view implementation that allows O(1) concatenation of two lists. Its primary
+ * purpose is to facilitate efficient accumulation of lists for later materialization. (ex:
+ * comprehensions that dispatch `add_list` to concat N lists together).
+ *
+ * <p>This does not support any of the standard list operations from {@link java.util.List}.
+ */
+final class ConcatenatedListView<E> extends AbstractList<E> {
+  private final List<List<? extends E>> sourceLists;
+  private int totalSize = 0;
+
+  ConcatenatedListView() {
+    this.sourceLists = new ArrayList<>();
+  }
+
+  ConcatenatedListView(Collection<? extends E> collection) {
+    this();
+    addAll(collection);
+  }
+
+  @Override
+  public boolean addAll(Collection<? extends E> collection) {
+    if (!(collection instanceof List)) {
+      // size() is O(1) iff it's a list
+      throw new IllegalStateException("addAll must be called with lists, not collections");
+    }
+
+    sourceLists.add((List<? extends E>) collection);
+    totalSize += collection.size();
+    return true;
+  }
+
+  @Override
+  public E get(int index) {
+    throw new UnsupportedOperationException("get method not supported.");
+  }
+
+  @Override
+  public int size() {
+    return totalSize;
+  }
+
+  @Override
+  public Iterator<E> iterator() {
+    return new ConcatenatingIterator();
+  }
+
+  /** Custom iterator to provide a flat view of all concatenated collections */
+  private class ConcatenatingIterator implements Iterator<E> {
+    private int index = 0;
+    private Iterator<? extends E> iterator = null;
+
+    @Override
+    public boolean hasNext() {
+      while (iterator == null || !iterator.hasNext()) {
+        if (index < sourceLists.size()) {
+          iterator = sourceLists.get(index).iterator();
+          index++;
+        } else {
+          return false;
+        }
+      }
+      return true;
+    }
+
+    @Override
+    public E next() {
+      if (!hasNext()) {
+        throw new NoSuchElementException();
+      }
+      return iterator.next();
+    }
+
+    @Override
+    public void remove() {
+      throw new UnsupportedOperationException("remove method not supported");
+    }
+  }
+}

runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java

@@ -173,6 +173,7 @@ public Object evalTrackingUnknowns(
         throws CelEvaluationException {
       ExecutionFrame frame = newExecutionFrame(resolver, functionResolver, listener);
       IntermediateResult internalResult = evalInternal(frame, ast.getExpr());
+
       return internalResult.value();
     }
 
@@ -924,6 +925,11 @@ private IntermediateResult evalComprehension(
         accuValue = IntermediateResult.create(new LazyExpression(compre.accuInit()));
       } else {
         accuValue = evalNonstrictly(frame, compre.accuInit());
+        if (accuValue.value() instanceof List) {
+          ConcatenatedListView<Object> lv =
+              new ConcatenatedListView<>((List<Object>) accuValue.value());
+          accuValue = IntermediateResult.create(lv);
+        }
       }
       int i = 0;
       for (Object elem : iterRange) {
@@ -950,6 +956,12 @@ private IntermediateResult evalComprehension(
         frame.popScope();
       }
 
+      // Materialize view back into an immutable list to facilitate O(1) lookups.
+      if (accuValue.value() instanceof ConcatenatedListView) {
+        accuValue =
+            IntermediateResult.create(ImmutableList.copyOf((Collection<?>) accuValue.value()));
+      }
+
       frame.pushScope(Collections.singletonMap(accuVar, accuValue));
       IntermediateResult result;
       try {

runtime/src/main/java/dev/cel/runtime/RuntimeHelpers.java

@@ -91,7 +91,12 @@ public static boolean matches(String string, String regexp, CelOptions celOption
 
   /** Concatenates two lists into a new list. */
   public static <E> List<E> concat(List<E> first, List<E> second) {
-    // TODO: return a view instead of an actual copy.
+    if (first instanceof ConcatenatedListView) {
+      // Comprehensions use a more efficient list view for performing O(1) concatenation
+      first.addAll(second);
+      return first;
+    }
+
     List<E> result = new ArrayList<>(first.size() + second.size());
     result.addAll(first);
     result.addAll(second);

runtime/src/test/java/dev/cel/runtime/CelRuntimeTest.java

@@ -14,12 +14,14 @@
 
 package dev.cel.runtime;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
 
 import com.google.api.expr.v1alpha1.Constant;
 import com.google.api.expr.v1alpha1.Expr;
 import com.google.api.expr.v1alpha1.Type.PrimitiveType;
+import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.Any;
@@ -41,6 +43,7 @@
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import dev.cel.common.types.CelV1AlphaTypes;
+import dev.cel.common.types.ListType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.StructTypeReference;
 import dev.cel.compiler.CelCompiler;
@@ -51,7 +54,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.LongStream;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -685,4 +690,34 @@ public void standardEnvironmentDisabledForRuntime_throws() throws Exception {
         .hasMessageThat()
         .contains("No matching overload for function 'size'. Overload candidates: size_string");
   }
+
+  @Test
+  @TestParameters("{run: 400000}")
+  @TestParameters("{run: 800000}")
+  @TestParameters("{run: 1200000}")
+  @TestParameters("{run: 1600000}")
+  @TestParameters("{run: 2000000}")
+  @SuppressWarnings("unchecked") // Test only
+  public void smokeTest(long run) throws Exception {
+    CelCompiler celCompiler =
+        CelCompilerFactory.standardCelCompilerBuilder()
+            .addVar("list", ListType.create(SimpleType.INT))
+            .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
+            .build();
+    CelRuntime celRuntime = CelRuntimeFactory.standardCelRuntimeBuilder().build();
+    ImmutableList<Long> tooLongList =
+        LongStream.range(0L, run + 1).boxed().collect(toImmutableList());
+    CelAbstractSyntaxTree ast = celCompiler.compile("list.map(x, x)").getAst();
+
+    // Wake cache up
+    List<Long> result =
+        (List<Long>) celRuntime.createProgram(ast).eval(ImmutableMap.of("list", tooLongList));
+    Stopwatch sw = Stopwatch.createStarted();
+    result = (List<Long>) celRuntime.createProgram(ast).eval(ImmutableMap.of("list", tooLongList));
+    sw.stop();
+
+    System.err.println("Elapsed: " + sw.elapsed(TimeUnit.MILLISECONDS));
+
+    assertThat(result).isEqualTo(tooLongList);
+  }
 }