Avoid hotpath cast to OptimizedFunctionOverload in DynamicDispatchOverload

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

@@ -106,17 +106,13 @@ public static CelResolvedOverload of(
       CelFunctionOverload definition,
       boolean isStrict,
       List<Class<?>> parameterTypes) {
-    OptimizedFunctionOverload optimizedDef =
-        (definition instanceof OptimizedFunctionOverload)
-            ? (OptimizedFunctionOverload) definition
-            : definition::apply;
     return new AutoValue_CelResolvedOverload(
         functionName,
         overloadId,
         ImmutableList.copyOf(parameterTypes),
         isStrict,
         definition,
-        optimizedDef);
+        OptimizedFunctionOverload.of(definition));
   }
 
   /**

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

@@ -189,7 +189,7 @@ private OverloadEntry mergeDynamicDispatchesOrThrow(
         DynamicDispatchOverload mergedOverload =
             new DynamicDispatchOverload(
                 overloadId,
-                ImmutableSet.<CelFunctionBinding>builder()
+                ImmutableSet.<FunctionBindingImpl>builder()
                     .addAll(existingOverload.getOverloadBindings())
                     .addAll(incomingOverload.getOverloadBindings())
                     .build());

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

@@ -33,6 +33,8 @@ final class FunctionBindingImpl implements InternalCelFunctionBinding {
 
   private final CelFunctionOverload definition;
 
+  private final OptimizedFunctionOverload optimizedDefinition;
+
   private final boolean isStrict;
 
   @Override
@@ -55,6 +57,10 @@ public CelFunctionOverload getDefinition() {
     return definition;
   }
 
+  OptimizedFunctionOverload getOptimizedDefinition() {
+    return optimizedDefinition;
+  }
+
   @Override
   public boolean isStrict() {
     return isStrict;
@@ -70,6 +76,7 @@ public boolean isStrict() {
     this.overloadId = overloadId;
     this.argTypes = argTypes;
     this.definition = definition;
+    this.optimizedDefinition = OptimizedFunctionOverload.of(definition);
     this.isStrict = isStrict;
   }
 
@@ -83,12 +90,15 @@ public boolean isStrict() {
 
   static ImmutableSet<CelFunctionBinding> groupOverloadsToFunction(
       String functionName, ImmutableSet<CelFunctionBinding> overloadBindings) {
-    ImmutableSet.Builder<CelFunctionBinding> builder = ImmutableSet.builder();
+    ImmutableSet.Builder<FunctionBindingImpl> functionBindingsBuilder = ImmutableSet.builder();
     for (CelFunctionBinding b : overloadBindings) {
-      builder.add(
+      functionBindingsBuilder.add(
           new FunctionBindingImpl(
               functionName, b.getOverloadId(), b.getArgTypes(), b.getDefinition(), b.isStrict()));
     }
+    ImmutableSet<FunctionBindingImpl> functionBindings = functionBindingsBuilder.build();
+    ImmutableSet.Builder<CelFunctionBinding> builder =
+        ImmutableSet.<CelFunctionBinding>builder().addAll(functionBindings);
 
     // If there is already a binding with the same name as the function, we treat it as a
     // "Singleton" binding and do not create a dynamic dispatch wrapper for it.
@@ -107,7 +117,7 @@ static ImmutableSet<CelFunctionBinding> groupOverloadsToFunction(
                 singleBinding.getDefinition(),
                 singleBinding.isStrict()));
       } else if (overloadBindings.size() > 1) {
-        builder.add(new DynamicDispatchBinding(functionName, overloadBindings));
+        builder.add(new DynamicDispatchBinding(functionName, functionBindings));
       }
     }
 
@@ -146,7 +156,7 @@ public boolean isStrict() {
     }
 
     private DynamicDispatchBinding(
-        String functionName, ImmutableSet<CelFunctionBinding> overloadBindings) {
+        String functionName, ImmutableSet<FunctionBindingImpl> overloadBindings) {
       this.isStrict = overloadBindings.stream().allMatch(CelFunctionBinding::isStrict);
       this.dynamicDispatchOverload = new DynamicDispatchOverload(functionName, overloadBindings);
     }
@@ -155,11 +165,11 @@ private DynamicDispatchBinding(
   @Immutable
   static final class DynamicDispatchOverload implements OptimizedFunctionOverload {
     private final String functionName;
-    private final ImmutableSet<CelFunctionBinding> overloadBindings;
+    private final ImmutableSet<FunctionBindingImpl> overloadBindings;
 
     @Override
     public Object apply(Object[] args) throws CelEvaluationException {
-      for (CelFunctionBinding overload : overloadBindings) {
+      for (FunctionBindingImpl overload : overloadBindings) {
         if (CelFunctionOverload.canHandle(args, overload.getArgTypes(), overload.isStrict())) {
           return overload.getDefinition().apply(args);
         }
@@ -174,12 +184,9 @@ public Object apply(Object[] args) throws CelEvaluationException {
 
     @Override
     public Object apply(Object arg) throws CelEvaluationException {
-      for (CelFunctionBinding overload : overloadBindings) {
+      for (FunctionBindingImpl overload : overloadBindings) {
         if (CelFunctionOverload.canHandle(arg, overload.getArgTypes(), overload.isStrict())) {
-          // Unary function binding setups always instantiate an OptimizedFunctionOverload, avoiding
-          // an instanceof check on the hotpath.
-          OptimizedFunctionOverload def = (OptimizedFunctionOverload) overload.getDefinition();
-          return def.apply(arg);
+          return overload.getOptimizedDefinition().apply(arg);
         }
       }
       throw new CelOverloadNotFoundException(
@@ -191,13 +198,10 @@ public Object apply(Object arg) throws CelEvaluationException {
 
     @Override
     public Object apply(Object arg1, Object arg2) throws CelEvaluationException {
-      for (CelFunctionBinding overload : overloadBindings) {
+      for (FunctionBindingImpl overload : overloadBindings) {
         if (CelFunctionOverload.canHandle(
             arg1, arg2, overload.getArgTypes(), overload.isStrict())) {
-          // Binary function binding setups always instantiate an OptimizedFunctionOverload,
-          // avoiding an instanceof check on the hotpath.
-          OptimizedFunctionOverload def = (OptimizedFunctionOverload) overload.getDefinition();
-          return def.apply(arg1, arg2);
+          return overload.getOptimizedDefinition().apply(arg1, arg2);
         }
       }
       throw new CelOverloadNotFoundException(
@@ -207,12 +211,12 @@ public Object apply(Object arg1, Object arg2) throws CelEvaluationException {
               .collect(toImmutableList()));
     }
 
-    ImmutableSet<CelFunctionBinding> getOverloadBindings() {
+    ImmutableSet<FunctionBindingImpl> getOverloadBindings() {
       return overloadBindings;
     }
 
     DynamicDispatchOverload(
-        String functionName, ImmutableSet<CelFunctionBinding> overloadBindings) {
+        String functionName, ImmutableSet<FunctionBindingImpl> overloadBindings) {
       this.functionName = functionName;
       this.overloadBindings = overloadBindings;
     }

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

@@ -32,4 +32,10 @@ default Object apply(Object arg) throws CelEvaluationException {
   default Object apply(Object arg1, Object arg2) throws CelEvaluationException {
     return apply(new Object[] {arg1, arg2});
   }
+
+  static OptimizedFunctionOverload of(CelFunctionOverload overload) {
+    return (overload instanceof OptimizedFunctionOverload)
+        ? (OptimizedFunctionOverload) overload
+        : overload::apply;
+  }
 }

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

@@ -31,7 +31,7 @@
 public final class FunctionBindingImplTest {
 
   @Test
-  public void dynamicDispatch_unaryOptimizedOverload_invokesOptimizedUnaryApply() throws Exception {
+  public void dynamicDispatch_unary_invokesOptimizedApplyAndAdaptsNonOptimized() throws Exception {
     OptimizedFunctionOverload mockOverload =
         new OptimizedFunctionOverload() {
           @Override
@@ -48,7 +48,8 @@ public Object apply(Object[] args) {
         CelFunctionBinding.from(
             "custom_opt_unary_long", ImmutableList.of(Long.class), mockOverload);
     CelFunctionBinding b2 =
-        CelFunctionBinding.from("custom_opt_unary_str", String.class, (String arg) -> arg + "!");
+        CelFunctionBinding.from(
+            "custom_unary_str", ImmutableList.of(String.class), args -> args[0] + "!");
     ImmutableSet<CelFunctionBinding> bindings =
         FunctionBindingImpl.groupOverloadsToFunction("custom_opt_unary", ImmutableSet.of(b1, b2));
     OptimizedFunctionOverload overload =
@@ -61,8 +62,7 @@ public Object apply(Object[] args) {
   }
 
   @Test
-  public void dynamicDispatch_binaryOptimizedOverload_invokesOptimizedBinaryApply()
-      throws Exception {
+  public void dynamicDispatch_binary_invokesOptimizedApplyAndAdaptsNonOptimized() throws Exception {
     OptimizedFunctionOverload mockOverload =
         new OptimizedFunctionOverload() {
           @Override
@@ -81,7 +81,9 @@ public Object apply(Object[] args) {
             "custom_opt_binary_long", ImmutableList.of(Long.class, Long.class), mockOverload);
     CelFunctionBinding b2 =
         CelFunctionBinding.from(
-            "custom_opt_bin_str", String.class, String.class, (String a, String b) -> a + b);
+            "custom_bin_str",
+            ImmutableList.of(String.class, String.class),
+            args -> (String) args[0] + args[1]);
     ImmutableSet<CelFunctionBinding> bindings =
         FunctionBindingImpl.groupOverloadsToFunction("custom_opt_binary", ImmutableSet.of(b1, b2));
     OptimizedFunctionOverload overload =