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

#include "eval/compiler/regex_precompilation_optimization.h"

#include <cstddef>

#include <cstdint>

#include <memory>

#include <optional>

#include <string>

#include <utility>

#include <vector>

#include "absl/base/nullability.h"

#include "absl/container/flat_hash_map.h"

#include "absl/status/status.h"

#include "absl/status/statusor.h"

#include "absl/strings/string_view.h"

#include "base/builtins.h"

#include "common/ast.h"

#include "common/expr.h"

#include "common/native_type.h"

#include "common/value.h"

#include "eval/compiler/flat_expr_builder_extensions.h"

#include "eval/eval/evaluator_core.h"

#include "eval/eval/regex_match_step.h"

#include "internal/re2_options.h"

#include "internal/status_macros.h"

#include "re2/re2.h"

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

namespace {

using ::cel::Ast;

using ::cel::CallExpr;

using ::cel::Expr;

using ::cel::NativeTypeId;

using ::cel::Reference;

using ::cel::Value;

using ReferenceMap = absl::flat_hash_map<int64_t, Reference>;

bool IsFunctionOverload(const Expr& expr, absl::string_view function,

                        absl::string_view overload, size_t arity,

                        const ReferenceMap& reference_map) {

  if (!expr.has_call_expr()) {

    return false;

  }

  const auto& call_expr = expr.call_expr();

  if (call_expr.function() != function) {

    return false;

  }

  if (call_expr.args().size() + (call_expr.has_target() ? 1 : 0) != arity) {

    return false;

  }

  // If parse-only and opted in to the optimization, assume this is the intended

  // overload. This will still only change the evaluation plan if the second arg

  // is a constant string.

  if (reference_map.empty()) {

    return true;

  }

  auto reference = reference_map.find(expr.id());

  if (reference != reference_map.end() &&

      reference->second.overload_id().size() == 1 &&

      reference->second.overload_id().front() == overload) {

    return true;

  }

  return false;

}

// Abstraction for deduplicating regular expressions over the course of a single

// create expression call. Should not be used during evaluation. Uses

// std::shared_ptr and std::weak_ptr.

class RegexProgramBuilder final {

 public:

  explicit RegexProgramBuilder(int max_program_size)

      : max_program_size_(max_program_size) {}

  absl::StatusOr<std::shared_ptr<const RE2>> BuildRegexProgram(

      std::string pattern) {

    auto existing = programs_.find(pattern);

    if (existing != programs_.end()) {

      if (auto program = existing->second.lock(); program) {

        return program;

      }

      programs_.erase(existing);

    }

    auto program =

        std::make_shared<RE2>(pattern, cel::internal::MakeRE2Options());

    CEL_RETURN_IF_ERROR(cel::internal::CheckRE2(*program, max_program_size_));

    programs_.insert({std::move(pattern), program});

    return program;

  }

 private:

  const int max_program_size_;

  absl::flat_hash_map<std::string, std::weak_ptr<const RE2>> programs_;

};

class RegexPrecompilationOptimization : public ProgramOptimizer {

 public:

  explicit RegexPrecompilationOptimization(const ReferenceMap& reference_map,

                                           int regex_max_program_size)

      : reference_map_(reference_map),

        regex_program_builder_(regex_max_program_size) {}

  absl::Status OnPreVisit(PlannerContext& context, const Expr& node) override {

    return absl::OkStatus();

  }

  absl::Status OnPostVisit(PlannerContext& context, const Expr& node) override {

    // Check that this is the correct matches overload instead of a user defined

    // overload.

    if (!IsFunctionOverload(node, cel::builtin::kRegexMatch, "matches_string",

                            2, reference_map_)) {

      return absl::OkStatus();

    }

    ProgramBuilder::Subexpression* subexpression =

        context.program_builder().GetSubexpression(&node);

    const CallExpr& call_expr = node.call_expr();

    const Expr& pattern_expr = call_expr.args().back();

    // Try to check if the regex is valid, whether or not we can actually update

    // the plan.

    std::optional<std::string> pattern =

        GetConstantString(context, subexpression, node, pattern_expr);

    if (!pattern.has_value()) {

      return absl::OkStatus();

    }

    CEL_ASSIGN_OR_RETURN(

        std::shared_ptr<const RE2> regex_program,

        regex_program_builder_.BuildRegexProgram(std::move(pattern).value()));

    if (subexpression == nullptr || subexpression->IsFlattened()) {

      // Already modified, can't update further.

      return absl::OkStatus();

    }

    const Expr& subject_expr =

        call_expr.has_target() ? call_expr.target() : call_expr.args().front();

    return RewritePlan(context, subexpression, node, subject_expr,

                       std::move(regex_program));

  }

 private:

  std::optional<std::string> GetConstantString(

      PlannerContext& context,

      ProgramBuilder::Subexpression* absl_nullable subexpression,

      const Expr& call_expr, const Expr& re_expr) const {

    if (re_expr.has_const_expr() && re_expr.const_expr().has_string_value()) {

      return re_expr.const_expr().string_value();

    }

    if (subexpression == nullptr || subexpression->IsFlattened()) {

      // Already modified, can't recover the input pattern.

      return std::nullopt;

    }

    std::optional<Value> constant;

      // otherwise stack-machine program.

      ExecutionPathView re_plan = context.GetSubplan(re_expr);

      if (re_plan.size() == 1) {

        cel::Value val;

        if (GetIfConstant(re_plan[0], val)) {

          constant = std::move(val);

        }

      }

    if (constant.has_value() && constant->IsString()) {

      return constant->GetString().ToString();

    }

    return std::nullopt;

  }

  absl::Status RewritePlan(

      PlannerContext& context,

      ProgramBuilder::Subexpression* absl_nonnull subexpression,

      const Expr& call, const Expr& subject,

      std::shared_ptr<const RE2> regex_program) {

    return RewriteStackMachinePlan(context, call, subject,

                                   std::move(regex_program));

  }

  absl::Status RewriteStackMachinePlan(

      PlannerContext& context, const Expr& call, const Expr& subject,

      std::shared_ptr<const RE2> regex_program) {

    if (context.GetSubplan(subject).empty()) {

      // This subexpression was already optimized, nothing to do.

      return absl::OkStatus();

    }

    CEL_ASSIGN_OR_RETURN(ExecutionPath new_plan,

                         context.ExtractSubplan(subject));

    CEL_ASSIGN_OR_RETURN(auto step,

                         CreateRegexMatchStep(std::move(regex_program)));

    new_plan.push_back(

        ExpressionStep::MakeGenericStep(std::move(step), call.id()));

    return context.ReplaceSubplan(call, std::move(new_plan));

  }

  const ReferenceMap& reference_map_;

  RegexProgramBuilder regex_program_builder_;

};

}  // namespace

ProgramOptimizerFactory CreateRegexPrecompilationExtension(

    int regex_max_program_size) {

  return [=](PlannerContext& context, const Ast& ast) {

    return std::make_unique<RegexPrecompilationOptimization>(

        ast.reference_map(), regex_max_program_size);

  };

}

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