Support overload lookup and filtering by signature as an alternative to overload_id

PiperOrigin-RevId: 928833427

Author: dmitriplotnikov

checker/BUILD

@@ -229,6 +229,7 @@ cc_library(
     hdrs = ["type_checker_subset_factory.h"],
     deps = [
         ":type_checker_builder",
+        "//common:decl",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",

checker/internal/type_checker_builder_impl.cc

@@ -163,8 +163,8 @@ std::optional<FunctionDecl> FilterDecl(FunctionDecl decl,
   FunctionDecl filtered;
   std::string name = decl.release_name();
   std::vector<OverloadDecl> overloads = decl.release_overloads();
-  for (const auto& ovl : overloads) {
-    if (subset.should_include_overload(name, ovl.id())) {
+  for (auto& ovl : overloads) {
+    if (subset.should_include_overload(name, ovl)) {
       absl::Status s = filtered.AddOverload(std::move(ovl));
       if (!s.ok()) {
         // Should not be possible to construct the original decl in a way that

checker/type_checker_builder.h

@@ -51,7 +51,7 @@ struct CheckerLibrary {
 // Represents a declaration to only use a subset of a library.
 struct TypeCheckerSubset {
   using FunctionPredicate = absl::AnyInvocable<bool(
-      absl::string_view function, absl::string_view overload_id) const>;
+      absl::string_view function, const OverloadDecl& overload) const>;
 
   // The id of the library to subset. Only one subset can be applied per
   // library id.

checker/type_checker_builder_factory_test.cc

@@ -235,8 +235,8 @@ TEST(TypeCheckerBuilderTest, AddLibraryIncludeSubset) {
   ASSERT_THAT(
       builder->AddLibrarySubset(
           {"testlib",
-           [](absl::string_view /*function*/, absl::string_view overload_id) {
-             return (overload_id == "add_int" || overload_id == "sub_int");
+           [](absl::string_view /*function*/, const OverloadDecl& overload) {
+             return (overload.id() == "add_int" || overload.id() == "sub_int");
            }}),
       IsOk());
   ASSERT_OK_AND_ASSIGN(auto checker, builder->Build());
@@ -274,9 +274,8 @@ TEST(TypeCheckerBuilderTest, AddLibraryExcludeSubset) {
   ASSERT_THAT(
       builder->AddLibrarySubset(
           {"testlib",
-           [](absl::string_view /*function*/, absl::string_view overload_id) {
-             return (overload_id != "add_int" && overload_id != "sub_int");
-             ;
+           [](absl::string_view /*function*/, const OverloadDecl& overload) {
+             return (overload.id() != "add_int" && overload.id() != "sub_int");
            }}),
       IsOk());
   ASSERT_OK_AND_ASSIGN(auto checker, builder->Build());
@@ -313,7 +312,7 @@ TEST(TypeCheckerBuilderTest, AddLibrarySubsetRemoveAllOvl) {
   ASSERT_THAT(builder->AddLibrary(SubsetTestlib()), IsOk());
   ASSERT_THAT(builder->AddLibrarySubset({"testlib",
                                          [](absl::string_view function,
-                                            absl::string_view /*overload_id*/) {
+                                            const OverloadDecl& /*overload*/) {
                                            return function != "add";
                                          }}),
               IsOk());
@@ -352,12 +351,12 @@ TEST(TypeCheckerBuilderTest, AddLibraryOneSubsetPerLibraryId) {
   ASSERT_THAT(
       builder->AddLibrarySubset(
           {"testlib", [](absl::string_view function,
-                         absl::string_view /*overload_id*/) { return true; }}),
+                         const OverloadDecl& /*overload*/) { return true; }}),
       IsOk());
   EXPECT_THAT(
       builder->AddLibrarySubset(
           {"testlib", [](absl::string_view function,
-                         absl::string_view /*overload_id*/) { return true; }}),
+                         const OverloadDecl& /*overload*/) { return true; }}),
       StatusIs(absl::StatusCode::kAlreadyExists));
 }
 
@@ -369,7 +368,7 @@ TEST(TypeCheckerBuilderTest, AddLibrarySubsetLibraryIdRequireds) {
   ASSERT_THAT(builder->AddLibrary(SubsetTestlib()), IsOk());
   EXPECT_THAT(builder->AddLibrarySubset({"",
                                          [](absl::string_view function,
-                                            absl::string_view /*overload_id*/) {
+                                            const OverloadDecl& /*overload*/) {
                                            return function == "add";
                                          }}),
               StatusIs(absl::StatusCode::kInvalidArgument));

checker/type_checker_subset_factory.cc

@@ -21,14 +21,16 @@
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "checker/type_checker_builder.h"
+#include "common/decl.h"
 
 namespace cel {
 
 TypeCheckerSubset::FunctionPredicate IncludeOverloadsByIdPredicate(
     absl::flat_hash_set<std::string> overload_ids) {
   return [overload_ids = std::move(overload_ids)](
-             absl::string_view /*function*/, absl::string_view overload_id) {
-    return overload_ids.contains(overload_id);
+             absl::string_view /*function*/, const OverloadDecl& overload) {
+    return overload_ids.contains(overload.id()) ||
+           overload_ids.contains(overload.signature());
   };
 }
 
@@ -41,8 +43,9 @@ TypeCheckerSubset::FunctionPredicate IncludeOverloadsByIdPredicate(
 TypeCheckerSubset::FunctionPredicate ExcludeOverloadsByIdPredicate(
     absl::flat_hash_set<std::string> overload_ids) {
   return [overload_ids = std::move(overload_ids)](
-             absl::string_view /*function*/, absl::string_view overload_id) {
-    return !overload_ids.contains(overload_id);
+             absl::string_view /*function*/, const OverloadDecl& overload) {
+    return !overload_ids.contains(overload.id()) &&
+           !overload_ids.contains(overload.signature());
   };
 }
 

checker/type_checker_subset_factory_test.cc

@@ -43,6 +43,8 @@ TEST(TypeCheckerSubsetFactoryTest, IncludeOverloadsByIdPredicate) {
       StandardOverloadIds::kEquals,
       StandardOverloadIds::kNotEquals,
       StandardOverloadIds::kNotStrictlyFalse,
+      "matches(string,string)",
+      "string.matches(string)",
   };
   ASSERT_THAT(builder->AddLibrary(StandardCompilerLibrary()), IsOk());
   ASSERT_THAT(builder->GetCheckerBuilder().AddLibrarySubset({
@@ -65,15 +67,19 @@ TEST(TypeCheckerSubsetFactoryTest, IncludeOverloadsByIdPredicate) {
 
   EXPECT_TRUE(r.IsValid());
 
+  // Allowed by signature.
+  ASSERT_OK_AND_ASSIGN(r, compiler->Compile("r'foo.*'.matches('foobar')"));
+  EXPECT_TRUE(r.IsValid());
+
+  ASSERT_OK_AND_ASSIGN(r, compiler->Compile("matches(r'foo.*', 'foobar')"));
+  EXPECT_TRUE(r.IsValid());
+
   // Not in allowlist.
   ASSERT_OK_AND_ASSIGN(r, compiler->Compile("1 + 2 < 3"));
   EXPECT_FALSE(r.IsValid());
 
   ASSERT_OK_AND_ASSIGN(r, compiler->Compile("'abc' + 'def'"));
   EXPECT_FALSE(r.IsValid());
-
-  ASSERT_OK_AND_ASSIGN(r, compiler->Compile("r'foo.*'.matches('foobar')"));
-  EXPECT_FALSE(r.IsValid());
 }
 
 TEST(TypeCheckerSubsetFactoryTest, ExcludeOverloadsByIdPredicate) {
@@ -83,6 +89,8 @@ TEST(TypeCheckerSubsetFactoryTest, ExcludeOverloadsByIdPredicate) {
   absl::string_view exclude_list[] = {
       StandardOverloadIds::kMatches,
       StandardOverloadIds::kMatchesMember,
+      "size(string)",
+      "string.size()",
   };
   ASSERT_THAT(builder->AddLibrary(StandardCompilerLibrary()), IsOk());
   ASSERT_THAT(builder->GetCheckerBuilder().AddLibrarySubset({
@@ -105,18 +113,35 @@ TEST(TypeCheckerSubsetFactoryTest, ExcludeOverloadsByIdPredicate) {
 
   EXPECT_TRUE(r.IsValid());
 
-  // Not in allowlist.
+  // Allowed.
   ASSERT_OK_AND_ASSIGN(r, compiler->Compile("1 + 2 < 3"));
   EXPECT_TRUE(r.IsValid());
 
   ASSERT_OK_AND_ASSIGN(r, compiler->Compile("'abc' + 'def'"));
   EXPECT_TRUE(r.IsValid());
 
+  // Excluded by ID.
   ASSERT_OK_AND_ASSIGN(r, compiler->Compile("r'foo.*'.matches('foobar')"));
   EXPECT_FALSE(r.IsValid());
 
   ASSERT_OK_AND_ASSIGN(r, compiler->Compile("matches(r'foo.*', 'foobar')"));
   EXPECT_FALSE(r.IsValid());
+
+  // Excluded by signature (top-level function).
+  ASSERT_OK_AND_ASSIGN(r, compiler->Compile("size('abc')"));
+  EXPECT_FALSE(r.IsValid());
+
+  // Allowed (other overloads of size).
+  ASSERT_OK_AND_ASSIGN(r, compiler->Compile("size([1, 2, 3])"));
+  EXPECT_TRUE(r.IsValid());
+
+  // Excluded by signature (member function).
+  ASSERT_OK_AND_ASSIGN(r, compiler->Compile("'abc'.size()"));
+  EXPECT_FALSE(r.IsValid());
+
+  // Allowed (other overloads of size member).
+  ASSERT_OK_AND_ASSIGN(r, compiler->Compile("[1, 2, 3].size()"));
+  EXPECT_TRUE(r.IsValid());
 }
 
 }  // namespace

common/BUILD

@@ -151,9 +151,8 @@ cc_library(
         "//internal:status_macros",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/hash",
-        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

common/decl.cc

@@ -20,8 +20,8 @@
 #include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
-#include "absl/log/absl_check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
@@ -109,43 +109,47 @@ bool SignaturesOverlap(const OverloadDecl& lhs, const OverloadDecl& rhs) {
 template <typename Overload>
 void AddOverloadInternal(std::string_view function_name,
                          std::vector<OverloadDecl>& insertion_order,
-                         OverloadDeclHashSet& overloads, Overload&& overload,
-                         absl::Status& status) {
+                         absl::flat_hash_map<std::string, size_t>& by_id,
+                         absl::flat_hash_map<std::string, size_t>& by_signature,
+                         Overload&& overload, absl::Status& status) {
   if (!status.ok()) {
     return;
   }
 
-  if (overload.id().empty()) {
-    OverloadDecl overload_decl = overload;
-    absl::StatusOr<std::string> overload_id =
-        common_internal::MakeOverloadSignature(
-            function_name, overload_decl.args(), overload_decl.member());
-    if (!overload_id.ok()) {
-      status = overload_id.status();
-      return;
-    }
-    overload_decl.set_id(*overload_id);
-    AddOverloadInternal(function_name, insertion_order, overloads,
-                        std::move(overload_decl), status);
+  absl::StatusOr<std::string> signature =
+      common_internal::MakeOverloadSignature(function_name, overload.args(),
+                                             overload.member());
+  if (!signature.ok()) {
+    status = signature.status();
     return;
   }
 
-  if (auto it = overloads.find(overload.id()); it != overloads.end()) {
+  OverloadDecl mutable_overload = std::forward<Overload>(overload);
+  mutable_overload.set_signature(*signature);
+
+  if (mutable_overload.id().empty()) {
+    mutable_overload.set_id(mutable_overload.signature());
+  }
+
+  if (auto it = by_id.find(mutable_overload.id()); it != by_id.end()) {
     status = absl::AlreadyExistsError(
-        absl::StrCat("overload already exists: ", overload.id()));
+        absl::StrCat("overload exists: ", mutable_overload.id()));
     return;
   }
-  for (const auto& existing : overloads) {
-    if (SignaturesOverlap(overload, existing)) {
+
+  for (const auto& existing : insertion_order) {
+    if (SignaturesOverlap(mutable_overload, existing)) {
       status = absl::InvalidArgumentError(
           absl::StrCat("overload signature collision: ", existing.id(),
-                       " collides with ", overload.id()));
+                       " collides with ", mutable_overload.id()));
       return;
     }
   }
-  const auto inserted = overloads.insert(std::forward<Overload>(overload));
-  ABSL_DCHECK(inserted.second);
-  insertion_order.push_back(*inserted.first);
+
+  size_t index = insertion_order.size();
+  by_id[mutable_overload.id()] = index;
+  by_signature[mutable_overload.signature()] = index;
+  insertion_order.push_back(std::move(mutable_overload));
 }
 
 void CollectTypeParams(absl::flat_hash_set<std::string>& type_params,
@@ -195,14 +199,25 @@ absl::flat_hash_set<std::string> OverloadDecl::GetTypeParams() const {
 
 void FunctionDecl::AddOverloadImpl(const OverloadDecl& overload,
                                    absl::Status& status) {
-  AddOverloadInternal(name_, overloads_.insertion_order, overloads_.set,
-                      overload, status);
+  AddOverloadInternal(name_, overloads_.insertion_order, overloads_.by_id,
+                      overloads_.by_signature, overload, status);
 }
 
 void FunctionDecl::AddOverloadImpl(OverloadDecl&& overload,
                                    absl::Status& status) {
-  AddOverloadInternal(name_, overloads_.insertion_order, overloads_.set,
-                      std::move(overload), status);
+  AddOverloadInternal(name_, overloads_.insertion_order, overloads_.by_id,
+                      overloads_.by_signature, std::move(overload), status);
+}
+
+const OverloadDecl* FunctionDecl::FindOverloadById(absl::string_view id) const {
+  if (auto it = overloads_.by_id.find(id); it != overloads_.by_id.end()) {
+    return &overloads_.insertion_order[it->second];
+  }
+  if (auto it = overloads_.by_signature.find(id);
+      it != overloads_.by_signature.end()) {
+    return &overloads_.insertion_order[it->second];
+  }
+  return nullptr;
 }
 
 }  // namespace cel