Merge "Add intrinsic for applying macro to token list" into main

google · Jul 18, 2024 · 861dde7 · 861dde7
2 parents c039d2a + 9ca6d48
commit 861dde7
Show file tree

Hide file tree

Showing 3 changed files with 138 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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