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 66ac4d7a6e..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<SqlSource> PerfettoSqlPreprocessor::RewriteInternal( } continue; } + if (macro_name == "__intrinsic_token_apply") { + ASSIGN_OR_RETURN( + std::optional<SqlSource> 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, @@ -381,4 +391,53 @@ PerfettoSqlPreprocessor::ExecuteTokenZipJoin( return {SqlSource::FromTraceProcessorImplementation(zipped)}; } +base::StatusOr<std::optional<SqlSource>> +PerfettoSqlPreprocessor::ExecuteTokenApply( + const SqliteTokenizer& tokenizer, + const SqliteTokenizer::Token& name_token, + std::vector<SqlSource> 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<SqlSource> 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<std::string> 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<SqlSource> 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<SqlSource> token_list, bool prefixed); + base::StatusOr<std::optional<SqlSource>> ExecuteTokenApply( + const SqliteTokenizer& tokenizer, + const SqliteTokenizer::Token& name_token, + std::vector<SqlSource> token_list); + SqliteTokenizer global_tokenizer_; const base::FlatHashMap<std::string, Macro>* macros_ = nullptr; std::unordered_set<std::string> 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