Add async state and in-flight tracking primitives.

PiperOrigin-RevId: 975452308

Author: l46kok

runtime/planner/BUILD.bazel

@@ -35,3 +35,9 @@ java_library(
     visibility = ["//:internal"],
     exports = ["//runtime/src/main/java/dev/cel/runtime/planner:async_completion_coordinator"],
 )
+
+java_library(
+    name = "async_call_state_tracker",
+    visibility = ["//:internal"],
+    exports = ["//runtime/src/main/java/dev/cel/runtime/planner:async_call_state_tracker"],
+)

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

@@ -19,6 +19,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import org.jspecify.annotations.Nullable;
@@ -35,6 +36,7 @@ public final class AccumulatedUnknowns {
   private static final int MAX_UNKNOWN_ATTRIBUTE_SIZE = 500_000;
   private final Set<Long> exprIds;
   private final Set<CelAttribute> attributes;
+  private final Set<Long> callIds;
 
   Set<Long> exprIds() {
     return exprIds;
@@ -44,6 +46,14 @@ Set<CelAttribute> attributes() {
     return attributes;
   }
 
+  public Set<Long> callIds() {
+    return Collections.unmodifiableSet(callIds);
+  }
+
+  public boolean hasCallIds() {
+    return !callIds.isEmpty();
+  }
+
   /**
    * Evaluates if the right hand side is an accumulated unknown, and if so, merges it into the
    * accumulator.
@@ -62,6 +72,7 @@ public AccumulatedUnknowns merge(AccumulatedUnknowns arg) {
     enforceMaxAttributeSize(this.attributes, arg.attributes);
     this.exprIds.addAll(arg.exprIds);
     this.attributes.addAll(arg.attributes);
+    this.callIds.addAll(arg.callIds);
     return this;
   }
 
@@ -75,7 +86,20 @@ static AccumulatedUnknowns create(Collection<Long> ids) {
 
   public static AccumulatedUnknowns create(
       Collection<Long> exprIds, Collection<CelAttribute> attributes) {
-    return new AccumulatedUnknowns(new HashSet<>(exprIds), new HashSet<>(attributes));
+    return new AccumulatedUnknowns(
+        new HashSet<>(exprIds), new HashSet<>(attributes), new HashSet<>());
+  }
+
+  /**
+   * Creates an accumulated unknown for a pending asynchronous call, recording {@code exprId} so the
+   * unknown retains its origin when adapted into a {@link CelUnknownSet}.
+   */
+  public static AccumulatedUnknowns createForAsyncCall(long exprId, long callId) {
+    HashSet<Long> exprIds = new HashSet<>();
+    exprIds.add(exprId);
+    HashSet<Long> callIds = new HashSet<>();
+    callIds.add(callId);
+    return new AccumulatedUnknowns(exprIds, new HashSet<>(), callIds);
   }
 
   private static void enforceMaxAttributeSize(
@@ -88,8 +112,9 @@ private static void enforceMaxAttributeSize(
     }
   }
 
-  private AccumulatedUnknowns(Set<Long> exprIds, Set<CelAttribute> attributes) {
+  private AccumulatedUnknowns(Set<Long> exprIds, Set<CelAttribute> attributes, Set<Long> callIds) {
     this.exprIds = exprIds;
     this.attributes = attributes;
+    this.callIds = callIds;
   }
 }

runtime/src/main/java/dev/cel/runtime/planner/AsyncCallKey.java

@@ -0,0 +1,247 @@
+// Copyright 2026 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 at
+//
+//      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.planner;
+
+import static java.util.Objects.requireNonNull;
+
+import dev.cel.runtime.RuntimeEquality;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * Unique cache key for an asynchronous function invocation at a given AST expression node.
+ *
+ * <p>Equality delegates to {@link RuntimeEquality}, except that {@link Double#NaN} and {@link
+ * Float#NaN} argument values compare equal so that re-evaluating a node with a NaN argument hits
+ * its existing call record. Map keys do not get this override because lookup goes through {@link
+ * RuntimeEquality#findInMap}.
+ */
+final class AsyncCallKey {
+  private final long exprId;
+  private final String functionName;
+  private final String overloadId;
+  private final Object[] args;
+  private final RuntimeEquality runtimeEquality;
+  private final int hashCode;
+
+  static AsyncCallKey create(
+      long exprId,
+      String functionName,
+      String overloadId,
+      Object[] args,
+      RuntimeEquality runtimeEquality) {
+    return new AsyncCallKey(exprId, functionName, overloadId, args, runtimeEquality);
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof AsyncCallKey)) {
+      return false;
+    }
+    AsyncCallKey other = (AsyncCallKey) o;
+    if (exprId != other.exprId
+        || !functionName.equals(other.functionName)
+        || !overloadId.equals(other.overloadId)
+        || args.length != other.args.length) {
+      return false;
+    }
+    for (int i = 0; i < args.length; i++) {
+      if (!argEquals(args[i], other.args[i], runtimeEquality)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  @Override
+  public int hashCode() {
+    return hashCode;
+  }
+
+  private static boolean argEquals(Object a, Object b, RuntimeEquality runtimeEquality) {
+    if (a == b) {
+      return true;
+    }
+    if (a == null || b == null) {
+      return false;
+    }
+    if (a instanceof Number && b instanceof Number) {
+      double da = ((Number) a).doubleValue();
+      double db = ((Number) b).doubleValue();
+      // CEL defines NaN != NaN; override so a node re-evaluated with NaN hits its existing record.
+      if (Double.isNaN(da) && Double.isNaN(db)) {
+        return true;
+      }
+      if (da == 0.0d && db == 0.0d) {
+        return celEquals(normalizeSignedZero(a, da), normalizeSignedZero(b, db), runtimeEquality);
+      }
+      return celEquals(a, b, runtimeEquality);
+    }
+    if (a instanceof List && b instanceof List) {
+      List<?> listA = (List<?>) a;
+      List<?> listB = (List<?>) b;
+      if (listA.size() != listB.size()) {
+        return false;
+      }
+      Iterator<?> iterA = listA.iterator();
+      Iterator<?> iterB = listB.iterator();
+      while (iterA.hasNext() && iterB.hasNext()) {
+        if (!argEquals(iterA.next(), iterB.next(), runtimeEquality)) {
+          return false;
+        }
+      }
+      return true;
+    }
+    if (a instanceof Map && b instanceof Map) {
+      Map<?, ?> mapA = (Map<?, ?>) a;
+      Map<?, ?> mapB = (Map<?, ?>) b;
+      if (mapA.size() != mapB.size()) {
+        return false;
+      }
+      for (Map.Entry<?, ?> entry : mapA.entrySet()) {
+        Optional<Object> valB = findInMap(mapB, entry.getKey(), runtimeEquality);
+        if (valB.isPresent()) {
+          if (!argEquals(entry.getValue(), valB.get(), runtimeEquality)) {
+            return false;
+          }
+        } else {
+          if (!mapB.containsKey(entry.getKey())
+              || entry.getValue() != null
+              || mapB.get(entry.getKey()) != null) {
+            return false;
+          }
+        }
+      }
+      return true;
+    }
+    if (a instanceof byte[] && b instanceof byte[]) {
+      return Arrays.equals((byte[]) a, (byte[]) b);
+    }
+    if (a instanceof Object[] && b instanceof Object[]) {
+      return Arrays.deepEquals((Object[]) a, (Object[]) b);
+    }
+    return celEquals(a, b, runtimeEquality);
+  }
+
+  /**
+   * Applies CEL heterogeneous equality, treating incomparable argument pairs (which throw unchecked
+   * exceptions from {@link RuntimeEquality#objectEquals}) as unequal.
+   */
+  private static boolean celEquals(Object a, Object b, RuntimeEquality runtimeEquality) {
+    try {
+      return runtimeEquality.objectEquals(a, b);
+    } catch (RuntimeException e) {
+      return false;
+    }
+  }
+
+  /**
+   * Normalizes {@code -0.0} to {@code 0.0} before comparison, because {@link
+   * RuntimeEquality#objectEquals} returns {@code false} for cross-type pairs such as {@code (0L,
+   * -0.0d)}.
+   */
+  private static Object normalizeSignedZero(Object value, double asDouble) {
+    if (asDouble == 0.0d && (value instanceof Double || value instanceof Float)) {
+      return 0.0d;
+    }
+    return value;
+  }
+
+  private static Optional<Object> findInMap(
+      Map<?, ?> map, Object key, RuntimeEquality runtimeEquality) {
+    try {
+      return runtimeEquality.findInMap(map, key);
+    } catch (RuntimeException e) {
+      return Optional.empty();
+    }
+  }
+
+  private static int computeHashCode(
+      long exprId,
+      String functionName,
+      String overloadId,
+      Object[] args,
+      RuntimeEquality runtimeEquality) {
+    int result = (int) (exprId ^ (exprId >>> 32));
+    result = 31 * result + functionName.hashCode();
+    result = 31 * result + overloadId.hashCode();
+    for (Object arg : args) {
+      result = 31 * result + hashArg(arg, runtimeEquality);
+    }
+    return result;
+  }
+
+  /**
+   * Hashes a single argument consistently with {@link #argEquals}.
+   *
+   * <p>Does not delegate to {@link RuntimeEquality#hashCode} because that method hashes {@code
+   * -0.0} and {@code 0.0} differently, which would break the {@link Object#hashCode} contract for
+   * keys containing signed zero.
+   */
+  private static int hashArg(Object arg, RuntimeEquality runtimeEquality) {
+    if (arg == null) {
+      return 0;
+    }
+    if (arg instanceof Number) {
+      double d = ((Number) arg).doubleValue();
+      // Normalize -0.0d to +0.0d so that values CEL considers equal hash identically.
+      if (d == 0.0d) {
+        d = 0.0d;
+      }
+      return Double.hashCode(d);
+    }
+    if (arg instanceof Iterable) {
+      int h = 1;
+      for (Object elem : (Iterable<?>) arg) {
+        h = h * 31 + hashArg(elem, runtimeEquality);
+      }
+      return h;
+    }
+    if (arg instanceof Map) {
+      int h = 0;
+      for (Map.Entry<?, ?> entry : ((Map<?, ?>) arg).entrySet()) {
+        h += hashArg(entry.getKey(), runtimeEquality) ^ hashArg(entry.getValue(), runtimeEquality);
+      }
+      return h;
+    }
+    if (arg instanceof byte[]) {
+      return Arrays.hashCode((byte[]) arg);
+    }
+    if (arg instanceof Object[]) {
+      return Arrays.deepHashCode((Object[]) arg);
+    }
+    return runtimeEquality.hashCode(arg);
+  }
+
+  private AsyncCallKey(
+      long exprId,
+      String functionName,
+      String overloadId,
+      Object[] args,
+      RuntimeEquality runtimeEquality) {
+    this.exprId = exprId;
+    this.functionName = requireNonNull(functionName);
+    this.overloadId = requireNonNull(overloadId);
+    this.args = args.clone();
+    this.runtimeEquality = requireNonNull(runtimeEquality);
+    this.hashCode = computeHashCode(exprId, functionName, overloadId, this.args, runtimeEquality);
+  }
+}