diff --git a/runtime/function_registry.cc b/runtime/function_registry.cc index b1f08e00a..fa6d75bae 100644 --- a/runtime/function_registry.cc +++ b/runtime/function_registry.cc @@ -145,6 +145,19 @@ FunctionRegistry::FindStaticOverloadsByArity(absl::string_view name, return matched_funcs; } + size_t count = 0; + for (const auto& overload : overloads->second.static_overloads) { + if (overload.descriptor->receiver_style() == receiver_style && + overload.descriptor->types().size() == arity) { + ++count; + } + } + // per go/dram-cy we are OK trading CPU for RAM as long as it's net positive + // in SWE terms + if (count > 1) { + matched_funcs.reserve(count); + } + for (const auto& overload : overloads->second.static_overloads) { if (overload.descriptor->receiver_style() == receiver_style && overload.descriptor->types().size() == arity) { diff --git a/runtime/function_registry_test.cc b/runtime/function_registry_test.cc index 40445f43a..6d9786daa 100644 --- a/runtime/function_registry_test.cc +++ b/runtime/function_registry_test.cc @@ -59,7 +59,7 @@ TEST(FunctionRegistryTest, InsertAndRetrieveLazyFunction) { cel::FunctionDescriptor lazy_function_desc{"LazyFunction", false, {}}; FunctionRegistry registry; Activation activation; - ASSERT_OK(registry.RegisterLazyFunction(lazy_function_desc)); + ASSERT_THAT(registry.RegisterLazyFunction(lazy_function_desc)); const auto descriptors = registry.FindLazyOverloads("LazyFunction", false, {}); @@ -72,7 +72,7 @@ TEST(FunctionRegistryTest, InsertAndRetrieveLazyFunction) { TEST(FunctionRegistryTest, LazyAndStaticFunctionShareDescriptorSpace) { FunctionRegistry registry; cel::FunctionDescriptor desc = ConstIntFunction::MakeDescriptor(); - ASSERT_OK(registry.RegisterLazyFunction(desc)); + ASSERT_THAT(registry.RegisterLazyFunction(desc)); absl::Status status = registry.Register(ConstIntFunction::MakeDescriptor(), std::make_unique()); @@ -82,7 +82,7 @@ TEST(FunctionRegistryTest, LazyAndStaticFunctionShareDescriptorSpace) { TEST(FunctionRegistryTest, FindStaticOverloadsReturns) { FunctionRegistry registry; cel::FunctionDescriptor desc = ConstIntFunction::MakeDescriptor(); - ASSERT_OK(registry.Register(desc, std::make_unique())); + ASSERT_THAT(registry.Register(desc, std::make_unique())); std::vector overloads = registry.FindStaticOverloads(desc.name(), false, {}); @@ -95,13 +95,32 @@ TEST(FunctionRegistryTest, FindStaticOverloadsReturns) { << "Expected single ConstFunction()"; } +TEST(FunctionRegistryTest, FindStaticOverloadsByArityExactCapacity) { + FunctionRegistry registry; + ASSERT_THAT(registry.Register({"Func", false, {Kind::kInt64}}, + std::make_unique()), + IsOk()); + ASSERT_THAT(registry.Register({"Func", false, {Kind::kDouble}}, + std::make_unique()), + IsOk()); + ASSERT_THAT(registry.Register({"Func", false, {Kind::kBool}}, + std::make_unique()), + IsOk()); + + std::vector overloads = + registry.FindStaticOverloadsByArity("Func", false, 1); + EXPECT_EQ(overloads.size(), 3); + EXPECT_EQ(overloads.size(), overloads.capacity()); +} + TEST(FunctionRegistryTest, ListFunctions) { cel::FunctionDescriptor lazy_function_desc{"LazyFunction", false, {}}; FunctionRegistry registry; - ASSERT_OK(registry.RegisterLazyFunction(lazy_function_desc)); + ASSERT_THAT(registry.RegisterLazyFunction(lazy_function_desc)); EXPECT_OK(registry.Register(ConstIntFunction::MakeDescriptor(), - std::make_unique())); + std::make_unique()), + IsOk()); auto registered_functions = registry.ListFunctions(); @@ -217,7 +236,7 @@ TEST_P(NonStrictRegistrationFailTest, /*receiver_style=*/false, {Kind::kAny}, /*is_strict=*/true); if (existing_function_is_lazy) { - ASSERT_OK(registry.RegisterLazyFunction(descriptor)); + ASSERT_THAT(registry.RegisterLazyFunction(descriptor)); } else { ASSERT_OK( registry.Register(descriptor, std::make_unique())); @@ -246,7 +265,7 @@ TEST_P(NonStrictRegistrationFailTest, /*receiver_style=*/false, {Kind::kAny}, /*is_strict=*/false); if (existing_function_is_lazy) { - ASSERT_OK(registry.RegisterLazyFunction(descriptor)); + ASSERT_THAT(registry.RegisterLazyFunction(descriptor)); } else { ASSERT_OK( registry.Register(descriptor, std::make_unique())); @@ -274,7 +293,7 @@ TEST_P(NonStrictRegistrationFailTest, CanRegisterStrictFunctionsWithoutLimit) { /*receiver_style=*/false, {Kind::kAny}, /*is_strict=*/true); if (existing_function_is_lazy) { - ASSERT_OK(registry.RegisterLazyFunction(descriptor)); + ASSERT_THAT(registry.RegisterLazyFunction(descriptor)); } else { ASSERT_OK( registry.Register(descriptor, std::make_unique()));