Allow planner to accept 3rd/4th args on select optimized ASTs from Java. Add a cross-stack test

eval/compiler/flat_expr_builder.cc

@@ -2233,12 +2233,17 @@ std::vector<ExecutionPathView> FlattenExpressionTable(
 }
 
 absl::Status CheckAstExtensions(
-    const std::vector<cel::ExtensionSpec>& extensions) {
+    const std::vector<cel::ExtensionSpec>& extensions,
+    bool enable_select_optimization) {
   for (const cel::ExtensionSpec& extension : extensions) {
     if (extension.id() == "cel_block" && extension.version().major() == 1) {
       // cel_block v1 is always supported.
       continue;
     }
+    if (extension.id() == "select_optimization" &&
+        extension.version().major() == 1 && enable_select_optimization) {
+      continue;
+    }
 
     // TODO(uncreated-issue/89): Add support for json field names.
     return absl::InvalidArgumentError(absl::StrCat(
@@ -2262,19 +2267,10 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
                                             : RuntimeIssue::Severity::kError;
   IssueCollector issue_collector(max_severity);
 
-  absl::StatusOr<std::vector<cel::ExtensionSpec>> runtime_extensions =
-      ExtractAndValidateRuntimeExtensions(*ast);
-
-  if (!runtime_extensions.ok()) {
-    CEL_RETURN_IF_ERROR(issue_collector.AddIssue(
-        RuntimeIssue::CreateError(runtime_extensions.status())));
-  }
-
-  auto status = CheckAstExtensions(*runtime_extensions);
-  if (!status.ok()) {
-    CEL_RETURN_IF_ERROR(
-        issue_collector.AddIssue(RuntimeIssue::CreateError(status)));
-  }
+  CEL_ASSIGN_OR_RETURN(std::vector<cel::ExtensionSpec> runtime_extensions,
+                       ExtractAndValidateRuntimeExtensions(*ast));
+  CEL_RETURN_IF_ERROR(
+      CheckAstExtensions(runtime_extensions, enable_select_optimization_));
 
   Resolver resolver(container_, function_registry_, type_registry_,
                     GetTypeProvider(),

eval/compiler/flat_expr_builder.h

@@ -81,6 +81,10 @@ class FlatExprBuilder {
 
   bool optional_types_enabled() const { return enable_optional_types_; }
 
+  // Called when select optimization is enabled to accept ASTs tagged with the
+  // `select_optimization` runtime extension.
+  void enable_select_optimization() { enable_select_optimization_ = true; }
+
  private:
   const cel::TypeProvider& GetTypeProvider() const;
 
@@ -90,6 +94,7 @@ class FlatExprBuilder {
   cel::RuntimeOptions options_;
   std::string container_;
   bool enable_optional_types_ = false;
+  bool enable_select_optimization_ = false;
   // TODO(uncreated-issue/45): evaluate whether we should use a shared_ptr here to
   // allow built expressions to keep the registries alive.
   const cel::FunctionRegistry& function_registry_;

eval/compiler/flat_expr_builder_test.cc

@@ -211,6 +211,51 @@ TEST(FlatExprBuilderTest, RuntimeExtensionsError) {
                        HasSubstr("unsupported CEL extension: ext1")));
 }
 
+TEST(FlatExprBuilderTest, SelectOptimizationRuntimeExtension) {
+  Expr expr;
+  expr.mutable_const_expr()->set_bool_value(true);
+  auto create = [&](CelExpressionBuilder& builder, absl::string_view id,
+                    int major, int minor) {
+    SourceInfo source_info;
+    auto* ext = source_info.add_extensions();
+    ext->set_id(id);
+    ext->mutable_version()->set_major(major);
+    ext->mutable_version()->set_minor(minor);
+    ext->add_affected_components(SourceInfo::Extension::COMPONENT_RUNTIME);
+    return builder.CreateExpression(&expr, &source_info).status();
+  };
+
+  CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv());
+  // Other program optimizers don't imply select optimization support.
+  builder.flat_expr_builder().AddProgramOptimizer(
+      [](auto&, const auto&) { return nullptr; });
+  EXPECT_THAT(
+      create(builder, "select_optimization", 1, 0),
+      StatusIs(
+          absl::StatusCode::kInvalidArgument,
+          HasSubstr("unsupported CEL extension: select_optimization@1.0")));
+
+  builder.flat_expr_builder().enable_select_optimization();
+  EXPECT_THAT(create(builder, "select_optimization", 1, 0), IsOk());
+  // Minor versions are informational (see SourceInfo.Extension.Version).
+  EXPECT_THAT(create(builder, "select_optimization", 1, 1), IsOk());
+  EXPECT_THAT(create(builder, "select_optimization", 2, 0),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+  EXPECT_THAT(create(builder, "ext1", 1, 0),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+
+  // The legacy API also plans the 3-arg form under default fail_on_warnings.
+  ASSERT_OK_AND_ASSIGN(ParsedExpr parsed,
+                       parser::Parse("f(x, [[1, 'field', 5]], 0)"));
+  parsed.mutable_expr()->mutable_call_expr()->set_function("cel.@attribute");
+  expr = parsed.expr();
+  InterpreterOptions options;
+  options.enable_select_optimization = true;
+  EXPECT_THAT(
+      create(*CreateCelExpressionBuilder(options), "select_optimization", 1, 0),
+      IsOk());
+}
+
 TEST(FlatExprBuilderTest, ConstValueUnset) {
   Expr expr;
   SourceInfo source_info;

eval/public/cel_expr_builder_factory.cc

@@ -122,19 +122,23 @@ std::unique_ptr<CelExpressionBuilder> CreateCelExpressionBuilder(
     // Add overloads for select optimization signature.
     // These are never bound, only used to prevent the builder from failing on
     // the overloads check.
-    absl::Status status =
-        builder->GetRegistry()->RegisterLazyFunction(CelFunctionDescriptor(
-            kCelAttribute, false, {cel::Kind::kAny, cel::Kind::kList}));
-    if (!status.ok()) {
-      ABSL_LOG(ERROR) << "Failed to register " << kCelAttribute << ": "
-                      << status;
-    }
-    status = builder->GetRegistry()->RegisterLazyFunction(CelFunctionDescriptor(
-        kCelHasField, false, {cel::Kind::kAny, cel::Kind::kList}));
-    if (!status.ok()) {
-      ABSL_LOG(ERROR) << "Failed to register " << kCelHasField << ": "
-                      << status;
+    for (const auto& desc : {
+             CelFunctionDescriptor(kCelAttribute, false,
+                                   {cel::Kind::kAny, cel::Kind::kList}),
+             CelFunctionDescriptor(
+                 kCelAttribute, false,
+                 {cel::Kind::kAny, cel::Kind::kList, cel::Kind::kAny}),
+             CelFunctionDescriptor(kCelHasField, false,
+                                   {cel::Kind::kAny, cel::Kind::kList}),
+         }) {
+      if (absl::Status status =
+              builder->GetRegistry()->RegisterLazyFunction(desc);
+          !status.ok()) {
+        ABSL_LOG(ERROR) << "Failed to register " << desc.name() << ": "
+                        << status;
+      }
     }
+    flat_expr_builder.enable_select_optimization();
     // Add runtime implementation.
     flat_expr_builder.AddProgramOptimizer(
         CreateSelectOptimizationProgramOptimizer());

extensions/select_optimization.cc

@@ -166,7 +166,8 @@ std::optional<SelectInstruction> GetSelectInstruction(
 }
 
 absl::StatusOr<SelectQualifier> SelectQualifierFromList(const ListExpr& list) {
-  if (list.elements().size() != 2) {
+  // Type codes may follow [number, name]; they're unused here.
+  if (list.elements().size() < 2 || list.elements().size() > 4) {
     return absl::InvalidArgumentError("Invalid cel.attribute select list");
   }
 
@@ -774,29 +775,22 @@ absl::Status SelectOptimizer::OnPostVisit(PlannerContext& context,
   if (fn != kCelHasField && fn != kCelAttribute) {
     return absl::OkStatus();
   }
+  const bool presence_test = fn == kCelHasField;
 
-  if (node.call_expr().args().size() < 2 ||
-      node.call_expr().args().size() > 3) {
+  // cel.@attribute may take a 3rd arg that only types the result; it's unused.
+  const size_t max_args = presence_test ? 2 : 3;
+  const size_t num_args = node.call_expr().args().size();
+  if (num_args < 2 || num_args > max_args) {
     return absl::InvalidArgumentError("Invalid cel.attribute call");
   }
 
-  if (node.call_expr().args().size() == 3) {
-    return absl::UnimplementedError("Optionals not yet supported");
-  }
-
   CEL_ASSIGN_OR_RETURN(std::vector<SelectQualifier> instructions,
                        SelectInstructionsFromCall(node.call_expr()));
 
   if (instructions.empty()) {
     return absl::InvalidArgumentError("Invalid cel.attribute no select steps.");
   }
 
-  bool presence_test = false;
-
-  if (fn == kCelHasField) {
-    presence_test = true;
-  }
-
   const Expr& operand = node.call_expr().args()[0];
   absl::string_view identifier;
   if (operand.has_ident_expr()) {
@@ -892,9 +886,13 @@ absl::Status EnableSelectOptimization(
   // the overloads check.
   CEL_RETURN_IF_ERROR(builder.function_registry().RegisterLazyFunction(
       FunctionDescriptor(kCelAttribute, false, {Kind::kAny, Kind::kList})));
+  CEL_RETURN_IF_ERROR(
+      builder.function_registry().RegisterLazyFunction(FunctionDescriptor(
+          kCelAttribute, false, {Kind::kAny, Kind::kList, Kind::kAny})));
 
   CEL_RETURN_IF_ERROR(builder.function_registry().RegisterLazyFunction(
       FunctionDescriptor(kCelHasField, false, {Kind::kAny, Kind::kList})));
+  flat_expr_builder->enable_select_optimization();
   // Add runtime implementation.
   flat_expr_builder->AddProgramOptimizer(
       CreateSelectOptimizationProgramOptimizer(options));

extensions/select_optimization.h

@@ -38,7 +38,8 @@ struct SelectOptimizationOptions {
 };
 
 // Enable select optimization on the given RuntimeBuilder, replacing long
-// select chains with a single operation.
+// select chains with a single operation. Also enables planning ASTs tagged with
+// the `select_optimization` runtime extension.
 //
 // This assumes that the type information at check time agrees with the
 // configured types at runtime.

extensions/select_optimization_test.cc

@@ -1753,6 +1753,39 @@ class SelectOptimizationUnexpectedAstTest : public SelectOptimizationTest {
     return element;
   }
 
+  Expr NextInt(int64_t value) {
+    Expr result = NextExpr();
+    result.mutable_const_expr().set_int_value(value);
+    return result;
+  }
+
+  Expr NextString(absl::string_view value) {
+    Expr result = NextExpr();
+    result.mutable_const_expr().set_string_value(value);
+    return result;
+  }
+
+  template <typename... Exprs>
+  Expr NextList(Exprs... elements) {
+    Expr result = NextExpr();
+    (result.mutable_list_expr().add_elements().set_expr(std::move(elements)),
+     ...);
+    return result;
+  }
+
+  // Plans `fn(ident, args...)` with only the select program optimizer.
+  template <typename... Exprs>
+  absl::Status PlanSelectCall(absl::string_view fn, Exprs... args) {
+    auto ast = std::make_unique<Ast>(NextExpr(), SourceInfo());
+    auto& call = ast->mutable_root_expr().mutable_call_expr();
+    call.set_function(fn);
+    call.add_args().mutable_ident_expr().set_name("ident");
+    (call.mutable_args().push_back(std::move(args)), ...);
+    FlatExprBuilder builder(env_, runtime_options_);
+    builder.AddProgramOptimizer(CreateSelectOptimizationProgramOptimizer());
+    return builder.CreateExpressionImpl(std::move(ast), nullptr).status();
+  }
+
  protected:
   int64_t next_id_;
 };
@@ -1907,37 +1940,46 @@ TEST_F(SelectOptimizationUnexpectedAstTest,
               StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
-TEST_F(SelectOptimizationUnexpectedAstTest, OptionalNotYetSupported) {
-  std::unique_ptr<Ast> ast = std::make_unique<Ast>(NextExpr(), SourceInfo());
-
-  ast->mutable_root_expr().mutable_call_expr().set_function(kCelAttribute);
-  auto& call_args = ast->mutable_root_expr().mutable_call_expr().mutable_args();
-  call_args.emplace_back(NextExpr()).mutable_ident_expr().set_name("ident");
-
-  auto& list_expr = call_args.emplace_back(NextExpr()).mutable_list_expr();
-  auto& fields = list_expr.mutable_elements()
-                     .emplace_back(NextListExprElement())
-                     .mutable_expr()
-                     .mutable_list_expr()
-                     .mutable_elements();
-
-  fields.emplace_back(NextListExprElement())
-      .mutable_expr()
-      .mutable_const_expr()
-      .set_int64_value(1);
-  fields.emplace_back(NextListExprElement())
-      .mutable_expr()
-      .mutable_const_expr()
-      .set_string_value("field");
-
-  call_args.emplace_back(NextExpr()).mutable_const_expr().set_int64_value(0);
-
-  FlatExprBuilder builder(env_, runtime_options_);
+TEST_F(SelectOptimizationUnexpectedAstTest, TypedTupleAttributeAccepted) {
+  using FieldDescriptor = google::protobuf::FieldDescriptor;
+  // [field_num, field_name, type_code] hops, a map leaf
+  // [field_num, field_name, -1, [key_type, value_type]], and a 3rd argument
+  // that only binds the result type.
+  EXPECT_THAT(
+      PlanSelectCall(
+          kCelAttribute,
+          NextList(NextList(NextInt(1), NextString("field"),
+                            NextInt(FieldDescriptor::TYPE_MESSAGE)),
+                   NextList(NextInt(2), NextString("map_field"), NextInt(-1),
+                            NextList(NextInt(FieldDescriptor::TYPE_STRING),
+                                     NextInt(FieldDescriptor::TYPE_INT32)))),
+          NextInt(0)),
+      IsOk());
+}
 
-  builder.AddProgramOptimizer(CreateSelectOptimizationProgramOptimizer());
+TEST_F(SelectOptimizationUnexpectedAstTest, TooManyArguments) {
+  EXPECT_THAT(
+      PlanSelectCall(kCelHasField,
+                     NextList(NextList(NextInt(1), NextString("field"))),
+                     NextInt(0)),
+      StatusIs(absl::StatusCode::kInvalidArgument,
+               HasSubstr("Invalid cel.attribute call")));
+  EXPECT_THAT(
+      PlanSelectCall(kCelAttribute,
+                     NextList(NextList(NextInt(1), NextString("field"))),
+                     NextInt(0), NextInt(0)),
+      StatusIs(absl::StatusCode::kInvalidArgument,
+               HasSubstr("Invalid cel.attribute call")));
+}
 
-  EXPECT_THAT(builder.CreateExpressionImpl(std::move(ast), nullptr),
-              StatusIs(absl::StatusCode::kUnimplemented));
+TEST_F(SelectOptimizationUnexpectedAstTest,
+       MalformedSelectPathTooManyElements) {
+  EXPECT_THAT(
+      PlanSelectCall(kCelAttribute,
+                     NextList(NextList(NextInt(1), NextString("field"),
+                                       NextInt(5), NextInt(0), NextInt(0)))),
+      StatusIs(absl::StatusCode::kInvalidArgument,
+               HasSubstr("Invalid cel.attribute select list")));
 }
 
 }  // namespace