flat_expr_builder_extensions.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.
//
// API definitions for planner extensions.
//
// These are provided to indirect build dependencies for optional features and
// require detailed understanding of how the flat expression builder works and
// its assumptions.
//
// These interfaces should not be implemented directly by CEL users.
#ifndef THIRD_PARTY_CEL_CPP_EVAL_COMPILER_FLAT_EXPR_BUILDER_EXTENSIONS_H_
#define THIRD_PARTY_CEL_CPP_EVAL_COMPILER_FLAT_EXPR_BUILDER_EXTENSIONS_H_

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <utility>
#include <variant>
#include <vector>

#include "absl/base/attributes.h"
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "absl/log/absl_check.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/types/variant.h"
#include "base/ast.h"
#include "base/type_provider.h"
#include "common/expr.h"
#include "common/type.h"
#include "common/type_reflector.h"
#include "eval/compiler/resolver.h"
#include "eval/eval/evaluator_core.h"
#include "eval/eval/expression_step_logic.h"
#include "runtime/internal/issue_collector.h"
#include "runtime/internal/runtime_env.h"
#include "runtime/runtime_options.h"
#include "google/protobuf/arena.h"
#include "google/protobuf/descriptor.h"
#include "google/protobuf/message.h"

namespace google::api::expr::runtime {

// Class representing a CEL program being built.
//
// Maintains tree structure and mapping from the AST representation to
// subexpressions. Maintains an insertion point for new steps and
// subexpressions.
//
// This class is thread-hostile and not intended for direct access outside of
// the Expression builder. Extensions should interact with this through the
// the PlannerContext member functions.
class ProgramBuilder {
 public:
  class Subexpression;

 private:
  using SubprogramMap =
      absl::flat_hash_map<const cel::Expr*,
                          std::unique_ptr<ProgramBuilder::Subexpression>>;

 public:
  // Represents a subexpression.
  //
  // Steps apply operations on the stack machine for the C++ runtime.
  // For most expression types, this maps to a post order traversal -- for all
  // nodes, evaluate dependencies (pushing their results to stack) then evaluate
  // self.
  //
  // Must be tied to a ProgramBuilder to coordinate relationships.
  class Subexpression {
   private:
    using Element = absl::variant<ExpressionStep, Subexpression* absl_nonnull>;

    using TreePlan = std::vector<Element>;
    using FlattenedPlan = ExecutionPath;

   public:
    ~Subexpression() = default;

    // Not copyable or movable.
    Subexpression(const Subexpression&) = delete;
    Subexpression& operator=(const Subexpression&) = delete;
    Subexpression(Subexpression&&) = delete;
    Subexpression& operator=(Subexpression&&) = delete;

    // Add a program step at the current end of the subexpression.
    bool AddStep(ExpressionStep step) {
      if (IsFlattened()) {
        flattened_elements().push_back(std::move(step));
        return true;
      }

      elements().push_back({std::move(step)});
      return true;
    }

    void AddSubexpression(Subexpression* absl_nonnull expr) {
      ABSL_DCHECK(absl::holds_alternative<TreePlan>(program_));
      ABSL_DCHECK(owner_ == expr->owner_);
      elements().push_back(expr);
    }

    ExpressionStep* absl_nullable GetIfExpressionStep(size_t index) {
      if (index >= elements().size()) {
        return nullptr;
      }
      return std::get_if<ExpressionStep>(&elements()[index]);
    }

    // Accessor for elements (either simple steps or subexpressions).
    //
    // Value is undefined if in the expression has already been flattened.
    std::vector<Element>& elements() {
      ABSL_DCHECK(absl::holds_alternative<TreePlan>(program_));
      return absl::get<TreePlan>(program_);
    }

    const std::vector<Element>& elements() const {
      ABSL_DCHECK(absl::holds_alternative<TreePlan>(program_));
      return absl::get<TreePlan>(program_);
    }

    // Accessor for program steps.
    //
    // Value is undefined if in the expression has not yet been flattened.
    ExecutionPath& flattened_elements() {
      ABSL_DCHECK(IsFlattened());
      return absl::get<FlattenedPlan>(program_);
    }

    const ExecutionPath& flattened_elements() const {
      ABSL_DCHECK(IsFlattened());
      return absl::get<FlattenedPlan>(program_);
    }

    // Compute the current number of program steps in this subexpression and
    // its dependencies.
    size_t ComputeSize() const;

    // Calculate the number of steps from the end of base to before target,
    // (including negative offsets).
    int CalculateOffset(int base, int target) const;

    // Extract a child subexpression.
    //
    // The expression is removed from the elements array.
    //
    // Returns nullptr if child is not an element of this subexpression.
    Subexpression* absl_nullable ExtractChild(Subexpression* child);

    // Flatten the subexpression.
    //
    // This removes the structure tracking for subexpressions, but makes the
    // subprogram evaluable on the runtime's stack machine.
    void Flatten();

    bool IsFlattened() const {
      return absl::holds_alternative<FlattenedPlan>(program_);
    }

    // Extract a flattened subexpression into the given vector. Transferring
    // ownership of the given steps.
    //
    // Returns false if the subexpression is not currently flattened.
    bool ExtractTo(ExecutionPath& out);

    // Returns true if this subexpression is trivial (reduces to a single
    // constant or identifier step) and can be inlined at reference sites.
    bool IsInlinable() const;

    // Creates an inlined ExpressionStep for this subexpression with the
    // given expr_id, or returns std::nullopt if !IsInlinable().
    std::optional<ExpressionStep> MaybeInline(int64_t expr_id) const;

   private:
    Subexpression(const cel::Expr* self, ProgramBuilder* owner);

    friend class ProgramBuilder;

    // Some extensions expect the program plan to be contiguous mid-planning.
    //
    // This adds complexity, but supports swapping to a flat representation as
    // needed.
    absl::variant<TreePlan, FlattenedPlan> program_;

    const cel::Expr* self_;
    const cel::Expr* absl_nullable parent_;
    ProgramBuilder* owner_;
  };

  ProgramBuilder();

  // Flatten the main subexpression and return its value.
  //
  // This transfers ownership of the program, returning the builder to starting
  // state. (See FlattenSubexpressions).
  ExecutionPath FlattenMain();

  // Flatten extracted subprograms.
  //
  // This transfers ownership of the subprograms, returning the extracted
  // programs table to starting state.
  std::vector<ExecutionPath> FlattenSubexpressions();

  // Return the current number of subexpressions in the program builder.
  size_t ExtractedSubexpressionCount() const {
    return extracted_subexpressions_.size();
  }

  // Returns the current subexpression where steps and new subexpressions are
  // added.
  //
  // May return null if the builder is not currently planning an expression.
  Subexpression* absl_nullable current() { return current_; }

  // Enter a subexpression context.
  //
  // Adds a subexpression at the current insertion point and move insertion
  // to the subexpression.
  //
  // Returns the new current() value.
  //
  // May return nullptr if the expression is already indexed in the program
  // builder.
  Subexpression* absl_nullable EnterSubexpression(const cel::Expr* expr,
                                                  size_t size_hint = 0);

  // Exit a subexpression context.
  //
  // Sets insertion point to parent.
  //
  // Returns the new current() value or nullptr if called out of order.
  Subexpression* absl_nullable ExitSubexpression(const cel::Expr* expr);

  // Return the subexpression mapped to the given expression.
  //
  // Returns nullptr if the mapping doesn't exist either due to the
  // program being overwritten or not encountering the expression.
  Subexpression* absl_nullable GetSubexpression(const cel::Expr* expr);

  // Return the extracted subexpression mapped to the given index.
  //
  // Returns nullptr if the mapping doesn't exist
  Subexpression* absl_nullable GetExtractedSubexpression(size_t index) {
    if (index >= extracted_subexpressions_.size()) {
      return nullptr;
    }

    return extracted_subexpressions_[index];
  }

  const Subexpression* absl_nullable GetExtractedSubexpression(
      size_t index) const {
    if (index >= extracted_subexpressions_.size()) {
      return nullptr;
    }

    return extracted_subexpressions_[index];
  }

  // Return index to the extracted subexpression.
  //
  // Returns -1 if the subexpression is not found.
  int ExtractSubexpression(const cel::Expr* expr);

  // Add a program step to the current subexpression.
  // If successful, returns the step pointer.
  //
  // Note: If successful, the pointer should remain valid until a further call
  // to AddStep, AddSubexpression, or Flatten.
  ExpressionStep* absl_nullable AddStep(ExpressionStep step);

  void Reset();

 private:
  static ExecutionPath FlattenSubexpression(Subexpression* absl_nonnull expr);

  Subexpression* absl_nullable MakeSubexpression(const cel::Expr* expr);

  Subexpression* absl_nullable root_;
  std::vector<Subexpression* absl_nonnull> extracted_subexpressions_;
  Subexpression* absl_nullable current_;
  SubprogramMap subprogram_map_;
};

// Class representing FlatExpr internals exposed to extensions.
class PlannerContext {
 public:
  PlannerContext(
      std::shared_ptr<const cel::runtime_internal::RuntimeEnv> environment,
      const Resolver& resolver, const cel::RuntimeOptions& options,
      const cel::TypeReflector& type_reflector,
      cel::runtime_internal::IssueCollector& issue_collector,
      ProgramBuilder& program_builder,
      std::shared_ptr<google::protobuf::Arena>& arena ABSL_ATTRIBUTE_LIFETIME_BOUND,
      std::shared_ptr<google::protobuf::MessageFactory> message_factory = nullptr)
      : environment_(std::move(environment)),
        resolver_(resolver),
        type_reflector_(type_reflector),
        options_(options),
        issue_collector_(issue_collector),
        program_builder_(program_builder),
        arena_(arena),
        message_factory_(std::move(message_factory)),
        explicit_arena_(arena_ != nullptr),
        resolve_types_(options_.enable_typed_field_access) {}

  ProgramBuilder& program_builder() { return program_builder_; }

  // Returns true if the subplan is inspectable.
  //
  // If false, the node is not mapped to a subexpression in the program builder.
  bool NodeHasInspectablePlan(const cel::Expr& node) const;

  // Return a view to the current subplan representing node.
  //
  // Note: this is invalidated after a sibling or parent is updated.
  //
  // This operation forces the subexpression to flatten which removes the
  // expr->program mapping for any descendants.
  ExecutionPathView GetSubplan(const cel::Expr& node);

  // Extract the plan steps for the given expr.
  //
  // After successful extraction, the subexpression is still inspectable, but
  // empty.
  absl::StatusOr<ExecutionPath> ExtractSubplan(const cel::Expr& node);

  // Replace the subplan associated with node with a new subplan.
  //
  // This operation forces the subexpression to flatten which removes the
  // expr->program mapping for any descendants.
  absl::Status ReplaceSubplan(const cel::Expr& node, ExecutionPath path);

  // Extend the current subplan with the given expression step.
  absl::Status AddSubplanStep(const cel::Expr& node, ExpressionStep step);
  absl::Status AddSubplanStep(const cel::Expr& node,
                              std::unique_ptr<ExpressionStepLogic> step,
                              int64_t expr_id = -1, int64_t stack_delta = 1) {
    return AddSubplanStep(node, ExpressionStep::MakeGenericStep(
                                    std::move(step), expr_id, stack_delta));
  }

  const Resolver& resolver() const { return resolver_; }
  const cel::TypeReflector& type_reflector() const { return type_reflector_; }
  const cel::RuntimeOptions& options() const { return options_; }
  cel::runtime_internal::IssueCollector& issue_collector() {
    return issue_collector_;
  }

  const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool() const {
    return environment_->descriptor_pool.get();
  }

  // Returns `true` if an arena was explicitly provided during planning.
  bool HasExplicitArena() const { return explicit_arena_; }

  google::protobuf::Arena* absl_nonnull MutableArena() {
    if (!explicit_arena_ && arena_ == nullptr) {
      arena_ = std::make_shared<google::protobuf::Arena>();
    }
    ABSL_DCHECK(arena_ != nullptr);
    return arena_.get();
  }

  // Whether the planner should resolve types in the AST to concrete cel::Type
  // values.
  bool resolve_types() const { return resolve_types_; }

  void set_resolve_types() { resolve_types_ = true; }

  // Returns the resolved type if available.
  //
  // This can only be used as a hint since we can't guarantee that the AST was
  // type checked or can be fully trusted (e.g. post-processing broke IDs).
  //
  // This is populated by FlatExprBuilder on the PreVisit pass for the Expr.
  std::optional<cel::Type> GetResolvedType(int64_t expr_id) const {
    auto it = resolved_types_.find(expr_id);
    if (it != resolved_types_.end()) {
      return it->second;
    }
    return std::nullopt;
  }

  std::optional<cel::Type> TryResolveType(int64_t expr_id,
                                          const cel::TypeSpec& ast_type);

  void SetResolvedType(int64_t expr_id, cel::Type type) {
    resolved_types_[expr_id] = std::move(type);
  }

  // Returns `true` if a message factory was explicitly provided during
  // planning.
  bool HasExplicitMessageFactory() const { return message_factory_ != nullptr; }

  google::protobuf::MessageFactory* absl_nonnull MutableMessageFactory() {
    return HasExplicitMessageFactory() ? message_factory_.get()
                                       : environment_->MutableMessageFactory();
  }

 private:
  const std::shared_ptr<const cel::runtime_internal::RuntimeEnv> environment_;
  const Resolver& resolver_;
  const cel::TypeReflector& type_reflector_;
  const cel::RuntimeOptions& options_;
  cel::runtime_internal::IssueCollector& issue_collector_;
  ProgramBuilder& program_builder_;
  std::shared_ptr<google::protobuf::Arena>& arena_;
  const std::shared_ptr<google::protobuf::MessageFactory> message_factory_;
  const bool explicit_arena_;
  bool resolve_types_;
  absl::flat_hash_map<int64_t, cel::Type> resolved_types_;
};

// Interface for Ast Transforms.
// If any are present, the FlatExprBuilder will apply the Ast Transforms in
// order on a copy of the relevant input expressions before planning the
// program.
class AstTransform {
 public:
  virtual ~AstTransform() = default;

  virtual absl::Status UpdateAst(PlannerContext& context,
                                 cel::Ast& ast) const = 0;
};

// Interface for program optimizers.
//
// If any are present, the FlatExprBuilder will notify the implementations in
// order as it traverses the input ast.
//
// Note: implementations must correctly check that subprograms are available
// before accessing (i.e. they have not already been edited).
class ProgramOptimizer {
 public:
  virtual ~ProgramOptimizer() = default;

  // Called before planning the given expr node.
  virtual absl::Status OnPreVisit(PlannerContext& context,
                                  const cel::Expr& node) = 0;

  // Called after planning the given expr node.
  virtual absl::Status OnPostVisit(PlannerContext& context,
                                   const cel::Expr& node) = 0;
};

// Type definition for ProgramOptimizer factories.
//
// The expression builder must remain thread compatible, but ProgramOptimizers
// are often stateful for a given expression. To avoid requiring the optimizer
// implementation to handle concurrent planning, the builder creates a new
// instance per expression planned.
//
// The factory must be thread safe, but the returned instance may assume
// it is called from a synchronous context.
using ProgramOptimizerFactory =
    absl::AnyInvocable<absl::StatusOr<std::unique_ptr<ProgramOptimizer>>(
        PlannerContext&, const cel::Ast&) const>;

}  // namespace google::api::expr::runtime

#endif  // THIRD_PARTY_CEL_CPP_EVAL_COMPILER_FLAT_EXPR_BUILDER_EXTENSIONS_H_