Inline trivial lazy subexpressions when planning CEL expressions.

eval/compiler/BUILD

@@ -61,6 +61,7 @@ cc_test(
         ":resolver",
         "//common:expr",
         "//common:type",
+        "//common:value",
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_logic",
         "//internal:testing",

eval/compiler/flat_expr_builder.cc

@@ -818,16 +818,24 @@ class FlatExprVisitor : public cel::AstVisitor {
     SlotLookupResult slot = LookupSlot(path);
 
     if (slot.subexpression >= 0) {
-      auto* subexpression =
+      const ProgramBuilder::Subexpression* subexpression =
           program_builder_.GetExtractedSubexpression(slot.subexpression);
       if (subexpression == nullptr) {
         SetProgressStatusIfError(
             absl::InternalError("bad subexpression reference"));
         return;
       }
-        // Off by one since mainline expression will be index 0.
-        AddStep(ExpressionStep::MakeLazyInitStep(
-            slot.slot, slot.subexpression + 1, expr.id()));
+      if (options_.enable_subexpression_inlining) {
+        if (std::optional<ExpressionStep> inlined_step =
+                subexpression->MaybeInline(expr.id());
+            inlined_step.has_value()) {
+          AddStep(*std::move(inlined_step));
+          return;
+        }
+      }
+      // Off by one since mainline expression will be index 0.
+      AddStep(ExpressionStep::MakeLazyInitStep(
+          slot.slot, slot.subexpression + 1, expr.id()));
       return;
     } else if (slot.slot >= 0) {
         AddStep(ExpressionStep::MakeReadSlotStep(slot.slot, expr.id()));

eval/compiler/flat_expr_builder_extensions.cc

@@ -14,8 +14,10 @@
 #include "eval/compiler/flat_expr_builder_extensions.h"
 
 #include <cstddef>
+#include <cstdint>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -28,14 +30,47 @@
 #include "absl/status/statusor.h"
 #include "absl/types/variant.h"
 #include "common/expr.h"
+#include "common/value.h"
 #include "eval/eval/evaluator_core.h"
+#include "eval/eval/ident_step.h"
 
 namespace google::api::expr::runtime {
 
 namespace {
 
 using Subexpression = google::api::expr::runtime::ProgramBuilder::Subexpression;
 
+const ExpressionStep* absl_nullable TryGetSingleStep(
+    const Subexpression* absl_nonnull subexpr) {
+  while (subexpr != nullptr) {
+    if (subexpr->IsFlattened()) {
+      const ExecutionPath& elements = subexpr->flattened_elements();
+      if (elements.size() != 1) {
+        return nullptr;
+      }
+      return &elements[0];
+    }
+    const auto& elements = subexpr->elements();
+    if (elements.size() != 1) {
+      return nullptr;
+    }
+    const auto& element = elements[0];
+    if (const ExpressionStep* step = absl::get_if<ExpressionStep>(&element);
+        step != nullptr) {
+      return step;
+    }
+    // This might happen if an AST node is interpreted as a no-op, but we still
+    // create a subexpression for it.
+    if (Subexpression* const* child = absl::get_if<Subexpression*>(&element);
+        child != nullptr) {
+      subexpr = *child;
+      continue;
+    }
+    return nullptr;
+  }
+  return nullptr;
+}
+
 }  // namespace
 
 Subexpression::Subexpression(const cel::Expr* self, ProgramBuilder* owner)
@@ -185,7 +220,33 @@ bool Subexpression::ExtractTo(ExecutionPath& out) {
   return true;
 }
 
-ExecutionPath ProgramBuilder::FlattenSubexpression(Subexpression* expr) {
+bool Subexpression::IsInlinable() const {
+  const ExpressionStep* step = TryGetSingleStep(this);
+  if (step == nullptr) {
+    return false;
+  }
+  return IsConstant(*step) || GetIfIdentStep(*step) != nullptr;
+}
+
+std::optional<ExpressionStep> Subexpression::MaybeInline(
+    int64_t expr_id) const {
+  const ExpressionStep* step = TryGetSingleStep(this);
+  if (step == nullptr) {
+    return std::nullopt;
+  }
+  cel::Value value;
+  if (GetIfConstant(*step, value)) {
+    return ExpressionStep::MakeConstant(value, expr_id);
+  }
+  if (const IdentStepInfo* ident = GetIfIdentStep(*step); ident != nullptr) {
+    return ExpressionStep::MakeIdentifierStep(ident->ident_index,
+                                              ident->slot_index, expr_id);
+  }
+  return std::nullopt;
+}
+
+ExecutionPath ProgramBuilder::FlattenSubexpression(
+    Subexpression* absl_nonnull expr) {
   ExecutionPath out;
 
   if (!expr) {

eval/compiler/flat_expr_builder_extensions.h

@@ -179,6 +179,14 @@ class ProgramBuilder {
     // 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);
 
@@ -256,6 +264,15 @@ class ProgramBuilder {
     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.

eval/compiler/flat_expr_builder_extensions_test.cc

@@ -23,10 +23,12 @@
 #include "absl/status/statusor.h"
 #include "common/expr.h"
 #include "common/type.h"
+#include "common/value.h"
 #include "eval/compiler/resolver.h"
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_logic.h"
 #include "eval/eval/function_step.h"
+#include "eval/eval/ident_step.h"
 #include "internal/testing.h"
 #include "runtime/function_registry.h"
 #include "runtime/internal/issue_collector.h"
@@ -539,5 +541,100 @@ TEST_F(ProgramBuilderTest, ExtractToRequiresFlatten) {
                                 UniquePtrHolds(step_ptrs.a)));
 }
 
+TEST_F(ProgramBuilderTest, IsInlinableAndMaybeInlineConstant) {
+  Expr a;
+  Expr b;
+  ProgramBuilder program_builder;
+
+  program_builder.EnterSubexpression(&a);
+  program_builder.EnterSubexpression(&b);
+  program_builder.AddStep(ExpressionStep::MakeConstant(cel::IntValue(42), 10));
+  program_builder.ExitSubexpression(&b);
+  program_builder.ExitSubexpression(&a);
+
+  Subexpression* subexpr_a = program_builder.GetSubexpression(&a);
+  Subexpression* subexpr_b = program_builder.GetSubexpression(&b);
+  ASSERT_TRUE(subexpr_a != nullptr);
+  ASSERT_TRUE(subexpr_b != nullptr);
+
+  // Unflattened direct constant step and nested single-child subexpression.
+  EXPECT_TRUE(subexpr_b->IsInlinable());
+  EXPECT_TRUE(subexpr_a->IsInlinable());
+
+  std::optional<ExpressionStep> inlined_b = subexpr_b->MaybeInline(99);
+  ASSERT_TRUE(inlined_b.has_value());
+  EXPECT_EQ(inlined_b->id(), 99);
+  cel::Value val;
+  ASSERT_TRUE(GetIfConstant(*inlined_b, val));
+  EXPECT_EQ(val.GetInt().NativeValue(), 42);
+
+  // Flattened constant step remains inlinable.
+  subexpr_a->Flatten();
+  EXPECT_TRUE(subexpr_a->IsInlinable());
+  std::optional<ExpressionStep> inlined_a = subexpr_a->MaybeInline(100);
+  ASSERT_TRUE(inlined_a.has_value());
+  EXPECT_EQ(inlined_a->id(), 100);
+  ASSERT_TRUE(GetIfConstant(*inlined_a, val));
+  EXPECT_EQ(val.GetInt().NativeValue(), 42);
+}
+
+TEST_F(ProgramBuilderTest, IsInlinableAndMaybeInlineIdentifier) {
+  Expr a;
+  ProgramBuilder program_builder;
+
+  program_builder.EnterSubexpression(&a);
+  program_builder.AddStep(ExpressionStep::MakeIdentifierStep(3, 7, 10));
+  program_builder.ExitSubexpression(&a);
+
+  Subexpression* subexpr_a = program_builder.GetSubexpression(&a);
+  ASSERT_TRUE(subexpr_a != nullptr);
+
+  EXPECT_TRUE(subexpr_a->IsInlinable());
+  std::optional<ExpressionStep> inlined = subexpr_a->MaybeInline(55);
+  ASSERT_TRUE(inlined.has_value());
+  EXPECT_EQ(inlined->id(), 55);
+  const IdentStepInfo* ident_info = GetIfIdentStep(*inlined);
+  ASSERT_TRUE(ident_info != nullptr);
+  EXPECT_EQ(ident_info->ident_index, 3);
+  EXPECT_EQ(ident_info->slot_index, 7);
+
+  subexpr_a->Flatten();
+  EXPECT_TRUE(subexpr_a->IsInlinable());
+  std::optional<ExpressionStep> inlined_flat = subexpr_a->MaybeInline(56);
+  ASSERT_TRUE(inlined_flat.has_value());
+  EXPECT_EQ(inlined_flat->id(), 56);
+  ident_info = GetIfIdentStep(*inlined_flat);
+  ASSERT_TRUE(ident_info != nullptr);
+  EXPECT_EQ(ident_info->ident_index, 3);
+  EXPECT_EQ(ident_info->slot_index, 7);
+}
+
+TEST_F(ProgramBuilderTest, NotInlinableWhenNonTrivial) {
+  Expr a;
+  Expr b;
+  Expr c;
+  ProgramBuilder program_builder;
+
+  ASSERT_OK_AND_ASSIGN(SimpleTreeSteps step_ptrs,
+                       InitSimpleTree(a, b, c, program_builder));
+  (void)step_ptrs;
+
+  Subexpression* subexpr_a = program_builder.GetSubexpression(&a);
+  Subexpression* subexpr_b = program_builder.GetSubexpression(&b);
+  ASSERT_TRUE(subexpr_a != nullptr);
+  ASSERT_TRUE(subexpr_b != nullptr);
+
+  // Single generic step is not inlinable.
+  EXPECT_FALSE(subexpr_b->IsInlinable());
+  EXPECT_EQ(subexpr_b->MaybeInline(1), std::nullopt);
+
+  // Multi-step subexpression is not inlinable (unflattened and flattened).
+  EXPECT_FALSE(subexpr_a->IsInlinable());
+  EXPECT_EQ(subexpr_a->MaybeInline(1), std::nullopt);
+  subexpr_a->Flatten();
+  EXPECT_FALSE(subexpr_a->IsInlinable());
+  EXPECT_EQ(subexpr_a->MaybeInline(1), std::nullopt);
+}
+
 }  // namespace
 }  // namespace google::api::expr::runtime