From 9ca6d48295553d82c5045bc97775657c9e3e08d3 Mon Sep 17 00:00:00 2001 From: Ryan Zuklie Date: Wed, 17 Jul 2024 13:34:55 -0700 Subject: [PATCH] Add intrinsic for applying macro to token list This applies a macro to a token list of token lists. ``` __intrinsic_token_apply!( ( (1, 2, 3), (4, 5, 6) ), fn, __intrinsic_token_comma!()) ``` Becomes ``` fn!(1, 2, 3), fn!(4, 5, 6) ``` Bug: 353706622 Change-Id: Id802eafa7b3879923d09cd13f49578c9b0e7d57a --- .../engine/perfetto_sql_preprocessor.cc | 64 ++++++++++++++++ .../engine/perfetto_sql_preprocessor.h | 5 ++ .../perfetto_sql_preprocessor_unittest.cc | 74 +++++++++++++++++++ 3 files changed, 143 insertions(+) diff --git a/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.cc b/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.cc index 9ac4ae842a..a255da7fa1 100644 --- a/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.cc +++ b/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.cc @@ -248,6 +248,16 @@ base::StatusOr PerfettoSqlPreprocessor::RewriteInternal( } continue; } + if (macro_name == "__intrinsic_token_apply") { + ASSIGN_OR_RETURN( + std::optional res, + ExecuteTokenApply(tokenizer, name_token, std::move(token_list))); + if (res) { + tokenizer.Rewrite(rewriter, prev, tok, *std::move(res), + SqliteTokenizer::EndToken::kInclusive); + } + continue; + } if (macro_name == "__intrinsic_token_comma") { if (!token_list.empty()) { return ErrorAtToken(tokenizer, name_token, @@ -310,6 +320,11 @@ base::StatusOr PerfettoSqlPreprocessor::ExecuteMacroInvocation( return ErrorAtToken(tokenizer, name_token, "Macro invoked with too few args"); } + if (token_list.size() > macro->args.size()) { + // TODO(lalitm): add a link to macro documentation. + return ErrorAtToken(tokenizer, name_token, + "Macro invoked with too many args"); + } std::unordered_map inner_bindings; for (auto& t : token_list) { inner_bindings.emplace(macro->args[inner_bindings.size()], std::move(t)); @@ -376,4 +391,53 @@ PerfettoSqlPreprocessor::ExecuteTokenZipJoin( return {SqlSource::FromTraceProcessorImplementation(zipped)}; } +base::StatusOr> +PerfettoSqlPreprocessor::ExecuteTokenApply( + const SqliteTokenizer& tokenizer, + const SqliteTokenizer::Token& name_token, + std::vector token_list) { + if (token_list.size() != 3) { + return ErrorAtToken(tokenizer, name_token, + "token_apply: must have exactly three args"); + } + + SqliteTokenizer arg_list_tokenizer(token_list[0]); + SqliteTokenizer::Token inner_tok = arg_list_tokenizer.NextNonWhitespace(); + if (inner_tok.token_type == SqliteTokenType::TK_VARIABLE) { + return {std::nullopt}; + } + ASSIGN_OR_RETURN(std::vector arg_list_sources, + ParseTokenList(arg_list_tokenizer, inner_tok, {})); + + SqliteTokenizer name_tokenizer(token_list[1]); + inner_tok = name_tokenizer.NextNonWhitespace(); + if (inner_tok.token_type == SqliteTokenType::TK_VARIABLE) { + return {std::nullopt}; + } + + std::vector res; + for (uint32_t i = 0; i < arg_list_sources.size(); ++i) { + SqliteTokenizer args_tokenizer(arg_list_sources[i]); + inner_tok = args_tokenizer.NextNonWhitespace(); + if (inner_tok.token_type == SqliteTokenType::TK_VARIABLE) { + return {std::nullopt}; + } + + ASSIGN_OR_RETURN(std::vector args_sources, + ParseTokenList(args_tokenizer, inner_tok, {})); + + ASSIGN_OR_RETURN(SqlSource invocation_res, + ExecuteMacroInvocation(tokenizer, name_token, + token_list[1].sql(), args_sources)); + res.push_back(invocation_res.sql()); + } + + if (res.empty()) { + return {SqlSource::FromTraceProcessorImplementation("")}; + } + + std::string zipped = base::Join(res, " " + token_list[2].sql() + " "); + return {SqlSource::FromTraceProcessorImplementation(zipped)}; +} + } // namespace perfetto::trace_processor diff --git a/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.h b/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.h index c9a0cce252..17377fcc08 100644 --- a/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.h +++ b/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor.h @@ -89,6 +89,11 @@ class PerfettoSqlPreprocessor { std::vector token_list, bool prefixed); + base::StatusOr> ExecuteTokenApply( + const SqliteTokenizer& tokenizer, + const SqliteTokenizer::Token& name_token, + std::vector token_list); + SqliteTokenizer global_tokenizer_; const base::FlatHashMap* macros_ = nullptr; std::unordered_set seen_macros_; diff --git a/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor_unittest.cc b/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor_unittest.cc index 0dc70955ca..ee233ea416 100644 --- a/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor_unittest.cc +++ b/src/trace_processor/perfetto_sql/engine/perfetto_sql_preprocessor_unittest.cc @@ -26,6 +26,8 @@ namespace perfetto::trace_processor { namespace { +using ::testing::HasSubstr; + using Macro = PerfettoSqlPreprocessor::Macro; class PerfettoSqlPreprocessorUnittest : public ::testing::Test { @@ -290,5 +292,77 @@ TEST_F(PerfettoSqlPreprocessorUnittest, ZipJoin) { } } +TEST_F(PerfettoSqlPreprocessorUnittest, TokenApply) { + auto foo = SqlSource::FromExecuteQuery( + "CREATE PERFETTO MACRO G(a Expr, b Expr) Returns Expr AS $a AS $b"); + macros_.Insert("G", Macro{ + false, + "G", + {"a", "b"}, + FindSubstr(foo, "$a AS $b"), + }); + + auto tp = SqlSource::FromExecuteQuery( + "CREATE PERFETTO MACRO TokApply(a Expr, b Expr, c Expr) Returns Expr " + "AS __intrinsic_token_apply!($a, $b, $c)"); + macros_.Insert("TokApply", + Macro{ + false, + "TokApply", + {"a", "b", "c"}, + FindSubstr(tp, "__intrinsic_token_apply!($a, $b, $c)"), + }); + { + auto source = + SqlSource::FromExecuteQuery("__intrinsic_token_apply!((), G, AND)"); + PerfettoSqlPreprocessor preprocessor(source, macros_); + ASSERT_TRUE(preprocessor.NextStatement()) + << preprocessor.status().message(); + ASSERT_EQ(preprocessor.statement().sql(), ""); + ASSERT_FALSE(preprocessor.NextStatement()); + } + { + auto source = SqlSource::FromExecuteQuery( + "__intrinsic_token_apply!(((foo, bar)), G, AND)"); + PerfettoSqlPreprocessor preprocessor(source, macros_); + ASSERT_TRUE(preprocessor.NextStatement()) + << preprocessor.status().message(); + ASSERT_EQ(preprocessor.statement().sql(), "foo AS bar"); + ASSERT_FALSE(preprocessor.NextStatement()); + } + { + auto source = SqlSource::FromExecuteQuery( + "__intrinsic_token_apply!(((foo, bar), (baz, bat)), G, AND)"); + PerfettoSqlPreprocessor preprocessor(source, macros_); + ASSERT_TRUE(preprocessor.NextStatement()) + << preprocessor.status().message(); + ASSERT_EQ(preprocessor.statement().sql(), "foo AS bar AND baz AS bat"); + ASSERT_FALSE(preprocessor.NextStatement()); + } + { + auto source = SqlSource::FromExecuteQuery( + "__intrinsic_token_apply!(((foo, bar), (baz, bat, bada)), G, AND)"); + PerfettoSqlPreprocessor preprocessor(source, macros_); + ASSERT_FALSE(preprocessor.NextStatement()); + ASSERT_THAT(preprocessor.status().message(), HasSubstr("too many args")); + } + { + auto source = SqlSource::FromExecuteQuery( + "__intrinsic_token_apply!(((foo, bar), (baz)), G, AND)"); + PerfettoSqlPreprocessor preprocessor(source, macros_); + ASSERT_FALSE(preprocessor.NextStatement()); + ASSERT_THAT(preprocessor.status().message(), HasSubstr("too few args")); + } + { + auto source = SqlSource::FromExecuteQuery( + "TokApply!(((foo, bar), (baz, bat)), G, AND)"); + PerfettoSqlPreprocessor preprocessor(source, macros_); + ASSERT_TRUE(preprocessor.NextStatement()) + << preprocessor.status().message(); + ASSERT_EQ(preprocessor.statement().sql(), "foo AS bar AND baz AS bat"); + ASSERT_FALSE(preprocessor.NextStatement()); + } +} + } // namespace } // namespace perfetto::trace_processor