function_registry.h

// Copyright 2023 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.

#ifndef THIRD_PARTY_CEL_CPP_RUNTIME_FUNCTION_REGISTRY_H_
#define THIRD_PARTY_CEL_CPP_RUNTIME_FUNCTION_REGISTRY_H_

#include <cstddef>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/node_hash_map.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "common/function_descriptor.h"
#include "common/kind.h"
#include "runtime/function.h"
#include "runtime/function_overload_reference.h"
#include "runtime/function_provider.h"

namespace cel {

// FunctionRegistry manages binding builtin or custom CEL functions to
// implementations.
//
// The registry is consulted during program planning to tie overload candidates
// to the CEL function in the AST getting planned.
//
// The registry takes ownership of the cel::Function objects -- the registry
// must outlive any program planned using it.
//
// This class is move-only.
class FunctionRegistry {
 public:
  // Represents a single overload for a lazily provided function.
  struct LazyOverload {
    const cel::FunctionDescriptor& descriptor;
    const cel::runtime_internal::FunctionProvider& provider;
  };

  FunctionRegistry() = default;

  // Move-only
  FunctionRegistry(FunctionRegistry&&) = default;
  FunctionRegistry& operator=(FunctionRegistry&&) = default;

  // Register a function implementation for the given descriptor.
  // Function registration should be performed prior to CelExpression creation.
  absl::Status Register(const cel::FunctionDescriptor& descriptor,
                        std::unique_ptr<cel::Function> implementation);

  // Register a lazily provided function.
  // Internally, the registry binds a FunctionProvider that provides an overload
  // at evaluation time by resolving against the overloads provided by an
  // implementation of cel::ActivationInterface.
  absl::Status RegisterLazyFunction(const cel::FunctionDescriptor& descriptor);

  // Find subset of cel::Function implementations that match overload conditions
  // As types may not be available during expression compilation,
  // further narrowing of this subset will happen at evaluation stage.
  //
  // name - the name of CEL function (as distinct from overload ID);
  // receiver_style - indicates whether function has receiver style;
  // types - argument types. If  type is not known during compilation,
  // cel::Kind::kAny should be passed.
  //
  // Results refer to underlying registry entries by reference. Results are
  // invalid after the registry is deleted.
  std::vector<cel::FunctionOverloadReference> FindStaticOverloads(
      absl::string_view name, bool receiver_style,
      absl::Span<const cel::Kind> types) const;

  std::vector<cel::FunctionOverloadReference> FindStaticOverloadsByArity(
      absl::string_view name, bool receiver_style, size_t arity) const;

  // Find subset of cel::Function providers that match overload conditions.
  // As types may not be available during expression compilation,
  // further narrowing of this subset will happen at evaluation stage.
  //
  // name - the name of CEL function (as distinct from overload ID);
  // receiver_style - indicates whether function has receiver style;
  // types - argument types. If  type is not known during compilation,
  // cel::Kind::kAny should be passed.
  //
  // Results refer to underlying registry entries by reference. Results are
  // invalid after the registry is deleted.
  std::vector<LazyOverload> FindLazyOverloads(
      absl::string_view name, bool receiver_style,
      absl::Span<const cel::Kind> types) const;

  std::vector<LazyOverload> FindLazyOverloadsByArity(absl::string_view name,
                                                     bool receiver_style,
                                                     size_t arity) const;

  // Retrieve list of registered function descriptors. This includes both
  // static and lazy functions.
  absl::node_hash_map<std::string, std::vector<const cel::FunctionDescriptor*>>
  ListFunctions() const;

 private:
  struct StaticFunctionEntry {
    StaticFunctionEntry(const cel::FunctionDescriptor& descriptor,
                        std::unique_ptr<cel::Function> impl)
        : descriptor(std::make_unique<cel::FunctionDescriptor>(descriptor)),
          implementation(std::move(impl)) {}

    // Extra indirection needed to preserve pointer stability for the
    // descriptors.
    std::unique_ptr<cel::FunctionDescriptor> descriptor;
    std::unique_ptr<cel::Function> implementation;
  };

  struct LazyFunctionEntry {
    LazyFunctionEntry(
        const cel::FunctionDescriptor& descriptor,
        std::unique_ptr<cel::runtime_internal::FunctionProvider> provider)
        : descriptor(std::make_unique<cel::FunctionDescriptor>(descriptor)),
          function_provider(std::move(provider)) {}

    // Extra indirection needed to preserve pointer stability for the
    // descriptors.
    std::unique_ptr<cel::FunctionDescriptor> descriptor;
    std::unique_ptr<cel::runtime_internal::FunctionProvider> function_provider;
  };

  struct RegistryEntry {
    std::vector<StaticFunctionEntry> static_overloads;
    std::vector<LazyFunctionEntry> lazy_overloads;
  };

  // Returns whether the descriptor is registered either as a lazy function or
  // as a static function.
  bool DescriptorRegistered(const cel::FunctionDescriptor& descriptor) const;

  // Returns true if after adding this function, the rule "a non-strict
  // function should have only a single overload" will be preserved.
  bool ValidateNonStrictOverload(
      const cel::FunctionDescriptor& descriptor) const;

  // indexed by function name (not type checker overload id).
  absl::flat_hash_map<std::string, RegistryEntry> functions_;
};

}  // namespace cel

#endif  // THIRD_PARTY_CEL_CPP_RUNTIME_FUNCTION_REGISTRY_H_