Skip to content

Commit 23c9913

Browse files
jnthntatumcopybara-github
jnthntatum
authored and
copybara-github
committed
Add option to specify an arena and persist the resolved types from type checking.
PiperOrigin-RevId: 907236091
1 parent f8ed1bd commit 23c9913

10 files changed

Lines changed: 139 additions & 19 deletions

File tree

‎checker/BUILD‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,9 @@ cc_library(
5050
":type_check_issue",
5151
"//common:ast",
5252
"//common:source",
53+
"//common:type",
5354
"@com_google_absl//absl/base:nullability",
55+
"@com_google_absl//absl/container:flat_hash_map",
5456
"@com_google_absl//absl/status",
5557
"@com_google_absl//absl/status:statusor",
5658
"@com_google_absl//absl/strings",
@@ -74,11 +76,14 @@ cc_test(
7476

7577
cc_library(
7678
name = "type_checker",
79+
srcs = ["type_checker.cc"],
7780
hdrs = ["type_checker.h"],
7881
deps = [
7982
":validation_result",
8083
"//common:ast",
84+
"@com_google_absl//absl/base:nullability",
8185
"@com_google_absl//absl/status:statusor",
86+
"@com_google_protobuf//:protobuf",
8287
],
8388
)
8489

‎checker/internal/type_checker_impl.cc‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include <cstddef>
1818
#include <cstdint>
1919
#include <memory>
20+
#include <optional>
2021
#include <string>
2122
#include <utility>
2223
#include <vector>
@@ -1176,11 +1177,13 @@ class ResolveRewriter : public AstRewriterBase {
11761177
explicit ResolveRewriter(const ResolveVisitor& visitor,
11771178
const TypeInferenceContext& inference_context,
11781179
const CheckerOptions& options,
1179-
Ast::ReferenceMap& references, Ast::TypeMap& types)
1180+
Ast::ReferenceMap& references, Ast::TypeMap& types,
1181+
ValidationResult::TypeMap& resolved_types)
11801182
: visitor_(visitor),
11811183
inference_context_(inference_context),
11821184
reference_map_(references),
11831185
type_map_(types),
1186+
resolved_types_(resolved_types),
11841187
options_(options) {}
11851188
bool PostVisitRewrite(Expr& expr) override {
11861189
bool rewritten = false;
@@ -1235,6 +1238,7 @@ class ResolveRewriter : public AstRewriterBase {
12351238
return rewritten;
12361239
}
12371240
type_map_[expr.id()] = *std::move(flattened_type);
1241+
resolved_types_[expr.id()] = iter->second;
12381242
rewritten = true;
12391243
}
12401244

@@ -1249,23 +1253,28 @@ class ResolveRewriter : public AstRewriterBase {
12491253
const TypeInferenceContext& inference_context_;
12501254
Ast::ReferenceMap& reference_map_;
12511255
Ast::TypeMap& type_map_;
1256+
ValidationResult::TypeMap& resolved_types_;
12521257
const CheckerOptions& options_;
12531258
};
12541259

12551260
} // namespace
12561261

1257-
absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
1258-
std::unique_ptr<Ast> ast) const {
1259-
google::protobuf::Arena type_arena;
1262+
absl::StatusOr<ValidationResult> TypeCheckerImpl::CheckImpl(
1263+
std::unique_ptr<Ast> ast, google::protobuf::Arena* arena) const {
1264+
std::optional<google::protobuf::Arena> type_arena;
1265+
if (arena == nullptr) {
1266+
type_arena.emplace();
1267+
arena = &(*type_arena);
1268+
}
12601269

12611270
std::vector<TypeCheckIssue> issues;
12621271
CEL_ASSIGN_OR_RETURN(auto generator,
12631272
NamespaceGenerator::Create(env_.container()));
12641273

12651274
TypeInferenceContext type_inference_context(
1266-
&type_arena, options_.enable_legacy_null_assignment);
1275+
arena, options_.enable_legacy_null_assignment);
12671276
ResolveVisitor visitor(std::move(generator), env_, *ast,
1268-
type_inference_context, issues, &type_arena);
1277+
type_inference_context, issues, arena);
12691278

12701279
TraversalOptions opts;
12711280
opts.use_comprehension_callbacks = true;
@@ -1310,9 +1319,10 @@ absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
13101319
// Apply updates as needed.
13111320
// Happens in a second pass to simplify validating that pointers haven't
13121321
// been invalidated by other updates.
1322+
ValidationResult::TypeMap resolved_types;
13131323
ResolveRewriter rewriter(visitor, type_inference_context, options_,
13141324
ast->mutable_reference_map(),
1315-
ast->mutable_type_map());
1325+
ast->mutable_type_map(), resolved_types);
13161326
AstRewrite(ast->mutable_root_expr(), rewriter);
13171327

13181328
CEL_RETURN_IF_ERROR(rewriter.status());
@@ -1325,7 +1335,15 @@ absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
13251335
{cel::ExtensionSpec::Component::kRuntime}));
13261336
}
13271337

1328-
return ValidationResult(std::move(ast), std::move(issues));
1338+
auto result = ValidationResult(std::move(ast), std::move(issues));
1339+
if (!type_arena.has_value()) {
1340+
// cel::Type values will expire after this function returns when the local
1341+
// arena is destructed. Only set the resolved type map if we're using the
1342+
// caller's arena.
1343+
result.SetResolvedTypeMap(std::move(resolved_types));
1344+
}
1345+
1346+
return result;
13291347
}
13301348

13311349
std::unique_ptr<TypeCheckerBuilder> TypeCheckerImpl::ToBuilder() const {

‎checker/internal/type_checker_impl.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,8 @@ class TypeCheckerImpl : public TypeChecker {
4242
TypeCheckerImpl(TypeCheckerImpl&&) = delete;
4343
TypeCheckerImpl& operator=(TypeCheckerImpl&&) = delete;
4444

45-
absl::StatusOr<ValidationResult> Check(
46-
std::unique_ptr<Ast> ast) const override;
45+
absl::StatusOr<ValidationResult> CheckImpl(
46+
std::unique_ptr<Ast> ast, google::protobuf::Arena* arena) const override;
4747

4848
std::unique_ptr<TypeCheckerBuilder> ToBuilder() const override;
4949

‎checker/type_checker.cc‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
// Copyright 2024 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "checker/type_checker.h"
16+
17+
namespace cel {
18+
absl::StatusOr<ValidationResult> TypeChecker::Check(
19+
std::unique_ptr<Ast> ast) const {
20+
return CheckImpl(std::move(ast), nullptr);
21+
}
22+
23+
absl::StatusOr<ValidationResult> TypeChecker::Check(
24+
std::unique_ptr<Ast> ast, google::protobuf::Arena* arena) const {
25+
return CheckImpl(std::move(ast), arena);
26+
}
27+
28+
absl::StatusOr<ValidationResult> TypeChecker::Check(const Ast& ast) const {
29+
return CheckImpl(std::make_unique<Ast>(ast), nullptr);
30+
}
31+
32+
absl::StatusOr<ValidationResult> TypeChecker::Check(
33+
const Ast& ast, google::protobuf::Arena* arena) const {
34+
return CheckImpl(std::make_unique<Ast>(ast), arena);
35+
}
36+
} // namespace cel

‎checker/type_checker.h‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@
1616
#define THIRD_PARTY_CEL_CPP_CHECKER_TYPE_CHECKER_H_
1717

1818
#include <memory>
19+
#include <utility>
1920

21+
#include "absl/base/nullability.h"
2022
#include "absl/status/statusor.h"
2123
#include "checker/validation_result.h"
2224
#include "common/ast.h"
25+
#include "google/protobuf/arena.h"
2326

2427
namespace cel {
2528

@@ -42,13 +45,19 @@ class TypeChecker {
4245
// A non-ok status is returned if type checking can't reasonably complete
4346
// (e.g. if an internal precondition is violated or an extension returns an
4447
// error).
45-
virtual absl::StatusOr<ValidationResult> Check(
46-
std::unique_ptr<Ast> ast) const = 0;
48+
absl::StatusOr<ValidationResult> Check(std::unique_ptr<Ast> ast) const;
49+
absl::StatusOr<ValidationResult> Check(std::unique_ptr<Ast> ast,
50+
google::protobuf::Arena* arena) const;
51+
absl::StatusOr<ValidationResult> Check(const Ast& ast) const;
52+
absl::StatusOr<ValidationResult> Check(const Ast& ast,
53+
google::protobuf::Arena* arena) const;
4754

4855
// Returns a builder initialized with the configuration of this type checker.
4956
virtual std::unique_ptr<TypeCheckerBuilder> ToBuilder() const = 0;
5057

51-
// TODO(uncreated-issue/73): add overload for cref AST.
58+
private:
59+
virtual absl::StatusOr<ValidationResult> CheckImpl(
60+
std::unique_ptr<Ast> ast, google::protobuf::Arena* absl_nullable arena) const = 0;
5261
};
5362

5463
} // namespace cel

‎checker/validation_result.h‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,26 +15,31 @@
1515
#ifndef THIRD_PARTY_CEL_CPP_CHECKER_VALIDATION_RESULT_H_
1616
#define THIRD_PARTY_CEL_CPP_CHECKER_VALIDATION_RESULT_H_
1717

18+
#include <cstdint>
1819
#include <memory>
1920
#include <string>
2021
#include <utility>
2122
#include <vector>
2223

2324
#include "absl/base/nullability.h"
25+
#include "absl/container/flat_hash_map.h"
2426
#include "absl/status/status.h"
2527
#include "absl/status/statusor.h"
2628
#include "absl/types/span.h"
2729
#include "checker/type_check_issue.h"
2830
#include "common/ast.h"
2931
#include "common/source.h"
32+
#include "common/type.h"
3033

3134
namespace cel {
3235

33-
// ValidationResult holds the result of TypeChecking.
36+
// ValidationResult holds the result of type checking.
3437
//
3538
// Error states are captured as type check issues where possible.
3639
class ValidationResult {
3740
public:
41+
using TypeMap = absl::flat_hash_map<int64_t, Type>;
42+
3843
ValidationResult(std::unique_ptr<Ast> ast, std::vector<TypeCheckIssue> issues)
3944
: ast_(std::move(ast)), issues_(std::move(issues)) {}
4045

@@ -71,6 +76,18 @@ class ValidationResult {
7176
return std::move(source_);
7277
}
7378

79+
// Returns the resolved type map for the AST.
80+
//
81+
// Only populated if the AST was checked with an explicit arena.
82+
//
83+
// The type entries may have storage in the arena or reference type
84+
// information from the type checker that produced the AST. This means the map
85+
// is only valid as long as both the type checker and the arena are valid.
86+
const TypeMap& GetResolvedTypeMap() const { return resolved_type_map_; }
87+
void SetResolvedTypeMap(TypeMap resolved_type_map) {
88+
resolved_type_map_ = std::move(resolved_type_map);
89+
}
90+
7491
// Returns a string representation of the issues in the result suitable for
7592
// display.
7693
//
@@ -89,6 +106,7 @@ class ValidationResult {
89106

90107
private:
91108
absl_nullable std::unique_ptr<Ast> ast_;
109+
TypeMap resolved_type_map_;
92110
std::vector<TypeCheckIssue> issues_;
93111
absl_nullable std::unique_ptr<Source> source_;
94112
};

‎compiler/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,11 @@ cc_library(
2828
"//parser:options",
2929
"//parser:parser_interface",
3030
"//validator",
31+
"@com_google_absl//absl/base:nullability",
3132
"@com_google_absl//absl/status",
3233
"@com_google_absl//absl/status:statusor",
3334
"@com_google_absl//absl/strings:string_view",
35+
"@com_google_protobuf//:protobuf",
3436
],
3537
)
3638

‎compiler/compiler.h‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include <string>
2020
#include <utility>
2121

22+
#include "absl/base/nullability.h"
2223
#include "absl/status/status.h"
2324
#include "absl/status/statusor.h"
2425
#include "absl/strings/string_view.h"
@@ -29,6 +30,7 @@
2930
#include "parser/options.h"
3031
#include "parser/parser_interface.h"
3132
#include "validator/validator.h"
33+
#include "google/protobuf/arena.h"
3234

3335
namespace cel {
3436

@@ -126,10 +128,16 @@ class Compiler {
126128
virtual ~Compiler() = default;
127129

128130
virtual absl::StatusOr<ValidationResult> Compile(
129-
absl::string_view source, absl::string_view description) const = 0;
131+
absl::string_view source, absl::string_view description,
132+
google::protobuf::Arena* absl_nullable arena) const = 0;
130133

131134
absl::StatusOr<ValidationResult> Compile(absl::string_view source) const {
132-
return Compile(source, "<input>");
135+
return Compile(source, "<input>", nullptr);
136+
}
137+
138+
absl::StatusOr<ValidationResult> Compile(
139+
absl::string_view source, absl::string_view description) const {
140+
return Compile(source, description, nullptr);
133141
}
134142

135143
// Accessor for the underlying type checker.

‎compiler/compiler_factory.cc‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
#include "parser/parser.h"
3434
#include "parser/parser_interface.h"
3535
#include "validator/validator.h"
36+
#include "google/protobuf/arena.h"
3637
#include "google/protobuf/descriptor.h"
3738

3839
namespace cel {
@@ -50,13 +51,13 @@ class CompilerImpl : public Compiler {
5051
validator_(std::move(validator)) {}
5152

5253
absl::StatusOr<ValidationResult> Compile(
53-
absl::string_view expression,
54-
absl::string_view description) const override {
54+
absl::string_view expression, absl::string_view description,
55+
google::protobuf::Arena* arena) const override {
5556
CEL_ASSIGN_OR_RETURN(auto source,
5657
cel::NewSource(expression, std::string(description)));
5758
CEL_ASSIGN_OR_RETURN(auto ast, parser_->Parse(*source));
5859
CEL_ASSIGN_OR_RETURN(ValidationResult result,
59-
type_checker_->Check(std::move(ast)));
60+
type_checker_->Check(std::move(ast), arena));
6061

6162
result.SetSource(std::move(source));
6263
if (!validator_.validations().empty()) {

‎compiler/compiler_factory_test.cc‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
#include "parser/parser_interface.h"
3838
#include "testutil/baseline_tests.h"
3939
#include "validator/timestamp_literal_validator.h"
40+
#include "google/protobuf/arena.h"
4041
#include "google/protobuf/descriptor.h"
4142

4243
namespace cel {
@@ -390,5 +391,27 @@ TEST(CompilerFactoryTest, ToBuilderWorks) {
390391
EXPECT_TRUE(result.IsValid());
391392
}
392393

394+
TEST(CompilerFactoryTest, SpecifyArenaKeepsResolvedTypes) {
395+
ASSERT_OK_AND_ASSIGN(
396+
auto builder,
397+
NewCompilerBuilder(cel::internal::GetSharedTestingDescriptorPool()));
398+
399+
ASSERT_THAT(builder->AddLibrary(StandardCompilerLibrary()), IsOk());
400+
ASSERT_THAT(builder->AddLibrary(OptionalCompilerLibrary()), IsOk());
401+
402+
ASSERT_OK_AND_ASSIGN(auto compiler, builder->Build());
403+
404+
google::protobuf::Arena arena;
405+
ASSERT_OK_AND_ASSIGN(ValidationResult result,
406+
compiler->Compile("[[1, 2, 3]][?0]", "<input>", &arena));
407+
ASSERT_OK_AND_ASSIGN(auto ast, result.ReleaseAst());
408+
auto it = result.GetResolvedTypeMap().find(ast->root_expr().id());
409+
ASSERT_TRUE(it != result.GetResolvedTypeMap().end());
410+
EXPECT_TRUE(
411+
it->second.IsOptional() &&
412+
it->second.GetOptional().GetParameter().IsList() &&
413+
it->second.GetOptional().GetParameter().GetList().GetElement().IsInt());
414+
}
415+
393416
} // namespace
394417
} // namespace cel

0 commit comments

Comments
 (0)