Simulate template argument pack expansion

robertoraggi · robertoraggi · commit ea7c30b3330c · 2025-03-20T22:43:54.000+01:00
diff --git a/src/parser/cxx/ast_rewriter.cc b/src/parser/cxx/ast_rewriter.cc
@@ -22,6 +22,7 @@
 
 // cxx
 #include <cxx/ast.h>
+#include <cxx/ast_cursor.h>
 #include <cxx/binder.h>
 #include <cxx/control.h>
 #include <cxx/decl.h>
@@ -31,6 +32,8 @@
 #include <cxx/type_checker.h>
 #include <cxx/types.h>
 
+#include <format>
+
 namespace cxx {
 
 ASTRewriter::ASTRewriter(TypeChecker* typeChcker,
@@ -54,6 +57,29 @@ void ASTRewriter::setRestrictedToDeclarations(bool restrictedToDeclarations) {
   restrictedToDeclarations_ = restrictedToDeclarations;
 }
 
+auto ASTRewriter::getParameterPack(ExpressionAST* ast) -> ParameterPackSymbol* {
+  for (auto cursor = ASTCursor{ast, {}}; cursor; ++cursor) {
+    const auto& current = *cursor;
+    if (!std::holds_alternative<AST*>(current.node)) continue;
+
+    auto id = ast_cast<IdExpressionAST>(std::get<AST*>(current.node));
+    if (!id) continue;
+
+    auto param = symbol_cast<NonTypeParameterSymbol>(id->symbol);
+    if (!param) continue;
+
+    if (param->depth() != 0) continue;
+
+    auto arg = templateArguments_[param->index()];
+    auto argSymbol = std::get<Symbol*>(arg);
+
+    auto parameterPack = symbol_cast<ParameterPackSymbol>(argSymbol);
+    if (parameterPack) return parameterPack;
+  }
+
+  return nullptr;
+}
+
 struct ASTRewriter::UnitVisitor {
   ASTRewriter& rewrite;
   [[nodiscard]] auto translationUnit() const -> TranslationUnit* {
@@ -2355,16 +2381,27 @@ auto ASTRewriter::ExpressionVisitor::operator()(IdExpressionAST* ast)
 
   copy->symbol = ast->symbol;
 
-  if (auto x = symbol_cast<NonTypeParameterSymbol>(copy->symbol);
-      x && x->depth() == 0 && x->index() < rewrite.templateArguments_.size()) {
+  if (auto param = symbol_cast<NonTypeParameterSymbol>(copy->symbol);
+      param && param->depth() == 0 &&
+      param->index() < rewrite.templateArguments_.size()) {
     auto initializerPtr =
-        std::get_if<Symbol*>(&rewrite.templateArguments_[x->index()]);
+        std::get_if<Symbol*>(&rewrite.templateArguments_[param->index()]);
+
     if (!initializerPtr) {
       cxx_runtime_error("expected initializer for non-type template parameter");
     }
 
     copy->symbol = *initializerPtr;
     copy->type = copy->symbol->type();
+
+    auto parameterPack = symbol_cast<ParameterPackSymbol>(copy->symbol);
+    if (parameterPack && rewrite.elementIndex_.has_value()) {
+      // ### TODO: check that idx is related to this non-type parameter.
+      auto idx = rewrite.elementIndex_.value();
+      auto element = parameterPack->elements()[idx];
+      copy->symbol = element;
+      copy->type = element->type();
+    }
   }
 
   copy->isTemplateIntroduced = ast->isTemplateIntroduced;
@@ -2473,6 +2510,33 @@ auto ASTRewriter::ExpressionVisitor::operator()(RightFoldExpressionAST* ast)
 
 auto ASTRewriter::ExpressionVisitor::operator()(LeftFoldExpressionAST* ast)
     -> ExpressionAST* {
+  if (auto parameterPack = rewrite.getParameterPack(ast->expression)) {
+    std::vector<ExpressionAST*> instantiations;
+    ExpressionAST* current = nullptr;
+    int n = 0;
+    for (auto element : parameterPack->elements()) {
+      std::optional<int> index{n};
+      std::swap(rewrite.elementIndex_, index);
+      auto expression = rewrite(ast->expression);
+      if (!current) {
+        current = expression;
+      } else {
+        auto binop = make_node<BinaryExpressionAST>(arena());
+        binop->valueCategory = current->valueCategory;
+        binop->type = current->type;
+        binop->leftExpression = current;
+        binop->op = ast->op;
+        binop->opLoc = ast->opLoc;
+        binop->rightExpression = expression;
+        current = binop;
+      }
+      std::swap(rewrite.elementIndex_, index);
+      ++n;
+    }
+
+    return current;
+  }
+
   auto copy = make_node<LeftFoldExpressionAST>(arena());
 
   copy->valueCategory = ast->valueCategory;
diff --git a/src/parser/cxx/ast_rewriter.h b/src/parser/cxx/ast_rewriter.h
@@ -151,8 +151,12 @@ class ASTRewriter {
  private:
   [[nodiscard]] auto rewriter() -> ASTRewriter* { return this; }
 
+  [[nodiscard]] auto getParameterPack(ExpressionAST* ast)
+      -> ParameterPackSymbol*;
+
   TypeChecker* typeChecker_ = nullptr;
   const std::vector<TemplateArgument>& templateArguments_;
+  std::optional<int> elementIndex_;
   TranslationUnit* unit_ = nullptr;
   Binder binder_;
   bool restrictedToDeclarations_ = false;
diff --git a/src/parser/cxx/control.cc b/src/parser/cxx/control.cc
@@ -143,6 +143,7 @@ struct Control::Private {
   std::forward_list<VariableSymbol> variableSymbols;
   std::forward_list<FieldSymbol> fieldSymbols;
   std::forward_list<ParameterSymbol> parameterSymbols;
+  std::forward_list<ParameterPackSymbol> parameterPackSymbols;
   std::forward_list<TypeParameterSymbol> typeParameterSymbols;
   std::forward_list<NonTypeParameterSymbol> nonTypeParameterSymbols;
   std::forward_list<TemplateTypeParameterSymbol> templateTypeParameterSymbols;
@@ -592,6 +593,13 @@ auto Control::newParameterSymbol(Scope* enclosingScope, SourceLocation loc)
   return symbol;
 }
 
+auto Control::newParameterPackSymbol(Scope* enclosingScope, SourceLocation loc)
+    -> ParameterPackSymbol* {
+  auto symbol = &d->parameterPackSymbols.emplace_front(enclosingScope);
+  symbol->setLocation(loc);
+  return symbol;
+}
+
 auto Control::newTypeParameterSymbol(Scope* enclosingScope, SourceLocation loc)
     -> TypeParameterSymbol* {
   auto symbol = &d->typeParameterSymbols.emplace_front(enclosingScope);
diff --git a/src/parser/cxx/control.h b/src/parser/cxx/control.h
@@ -202,6 +202,9 @@ class Control {
   [[nodiscard]] auto newParameterSymbol(Scope* enclosingScope,
                                         SourceLocation sourceLocation)
       -> ParameterSymbol*;
+  [[nodiscard]] auto newParameterPackSymbol(Scope* enclosingScope,
+                                            SourceLocation sourceLocation)
+      -> ParameterPackSymbol*;
   [[nodiscard]] auto newTypeParameterSymbol(Scope* enclosingScope,
                                             SourceLocation sourceLocation)
       -> TypeParameterSymbol*;
diff --git a/src/parser/cxx/external_name_encoder.cc b/src/parser/cxx/external_name_encoder.cc
@@ -156,6 +156,8 @@ struct ExternalNameEncoder::SymbolVisitor {
 
   void operator()(ParameterSymbol* symbol) {}
 
+  void operator()(ParameterPackSymbol* symbol) {}
+
   void operator()(EnumeratorSymbol* symbol) {}
 
   void operator()(FunctionParametersSymbol* symbol) {}
diff --git a/src/parser/cxx/symbol_instantiation.cc b/src/parser/cxx/symbol_instantiation.cc
@@ -62,6 +62,7 @@ struct SymbolInstantiation::VisitSymbol {
   [[nodiscard]] auto operator()(VariableSymbol* symbol) -> Symbol*;
   [[nodiscard]] auto operator()(FieldSymbol* symbol) -> Symbol*;
   [[nodiscard]] auto operator()(ParameterSymbol* symbol) -> Symbol*;
+  [[nodiscard]] auto operator()(ParameterPackSymbol* symbol) -> Symbol*;
   [[nodiscard]] auto operator()(EnumeratorSymbol* symbol) -> Symbol*;
   [[nodiscard]] auto operator()(FunctionParametersSymbol* symbol) -> Symbol*;
   [[nodiscard]] auto operator()(TemplateParametersSymbol* symbol) -> Symbol*;
@@ -312,6 +313,12 @@ auto SymbolInstantiation::VisitSymbol::operator()(ParameterSymbol* symbol)
   return newSymbol;
 }
 
+auto SymbolInstantiation::VisitSymbol::operator()(ParameterPackSymbol* symbol)
+    -> Symbol* {
+  auto newSymbol = self.replacement(symbol);
+  return newSymbol;
+}
+
 auto SymbolInstantiation::VisitSymbol::operator()(EnumeratorSymbol* symbol)
     -> Symbol* {
   auto newSymbol = self.replacement(symbol);
diff --git a/src/parser/cxx/symbol_printer.cc b/src/parser/cxx/symbol_printer.cc
@@ -269,6 +269,12 @@ struct DumpSymbols {
                        to_string(symbol->type(), symbol->name()));
   }
 
+  void operator()(ParameterPackSymbol* symbol) {
+    indent();
+    out << std::format("parameter pack {}\n",
+                       to_string(symbol->type(), symbol->name()));
+  }
+
   void operator()(TypeParameterSymbol* symbol) {
     std::string_view pack = symbol->isParameterPack() ? "..." : "";
     indent();
diff --git a/src/parser/cxx/symbols.cc b/src/parser/cxx/symbols.cc
@@ -615,6 +615,19 @@ ParameterSymbol::ParameterSymbol(Scope* enclosingScope)
 
 ParameterSymbol::~ParameterSymbol() {}
 
+ParameterPackSymbol::ParameterPackSymbol(Scope* enclosingScope)
+    : Symbol(Kind, enclosingScope) {}
+
+ParameterPackSymbol::~ParameterPackSymbol() {}
+
+auto ParameterPackSymbol::elements() const -> const std::vector<Symbol*>& {
+  return elements_;
+}
+
+void ParameterPackSymbol::addElement(Symbol* element) {
+  elements_.push_back(element);
+}
+
 TypeParameterSymbol::TypeParameterSymbol(Scope* enclosingScope)
     : Symbol(Kind, enclosingScope) {}
 
diff --git a/src/parser/cxx/symbols.h b/src/parser/cxx/symbols.h
@@ -624,6 +624,20 @@ class ParameterSymbol final : public Symbol {
   ~ParameterSymbol() override;
 };
 
+class ParameterPackSymbol final : public Symbol {
+ public:
+  constexpr static auto Kind = SymbolKind::kParameterPack;
+
+  explicit ParameterPackSymbol(Scope* enclosingScope);
+  ~ParameterPackSymbol() override;
+
+  [[nodiscard]] auto elements() const -> const std::vector<Symbol*>&;
+  void addElement(Symbol* element);
+
+ private:
+  std::vector<Symbol*> elements_;
+};
+
 class TypeParameterSymbol final : public Symbol {
  public:
   constexpr static auto Kind = SymbolKind::kTypeParameter;
diff --git a/src/parser/cxx/symbols_fwd.h b/src/parser/cxx/symbols_fwd.h
@@ -37,6 +37,7 @@ namespace cxx {
   V(Variable)                  \
   V(Field)                     \
   V(Parameter)                 \
+  V(ParameterPack)             \
   V(Enumerator)                \
   V(FunctionParameters)        \
   V(TemplateParameters)        \
diff --git a/tests/api_tests/test_rewriter.cc b/tests/api_tests/test_rewriter.cc
@@ -20,6 +20,7 @@
 
 #include <cxx/ast.h>
 #include <cxx/ast_interpreter.h>
+#include <cxx/ast_pretty_printer.h>
 #include <cxx/ast_rewriter.h>
 #include <cxx/control.h>
 #include <cxx/names.h>
@@ -109,13 +110,12 @@ using Func1 = Func<int, float>;
   auto func1 = source.getAs<TypeAliasSymbol>("Func1");
   ASSERT_TRUE(func1 != nullptr);
 
-  std::cout << "Func1: " << to_string(func1->type()) << "\n";
-
   auto func1Type = type_cast<UnresolvedNameType>(func1->type());
   ASSERT_TRUE(func1Type != nullptr);
 
   auto templateId = ast_cast<SimpleTemplateIdAST>(func1Type->unqualifiedId());
   ASSERT_TRUE(templateId != nullptr);
+
   auto templateSym =
       symbol_cast<TypeAliasSymbol>(templateId->primaryTemplateSymbol);
   ASSERT_TRUE(templateSym != nullptr);
@@ -128,6 +128,7 @@ using Func1 = Func<int, float>;
                                  getTemplateBodyAs<AliasDeclarationAST>(
                                      templateSym->templateDeclaration()),
                                  templateArguments);
+
   ASSERT_EQ(to_string(funcInstance->typeId->type), "void (int, const float&)");
 }
 
@@ -145,6 +146,7 @@ constexpr int y = c<123 * 2>;
 
   auto y = source.getAs<VariableSymbol>("y");
   ASSERT_TRUE(y != nullptr);
+
   auto yinit = ast_cast<EqualInitializerAST>(y->initializer())->expression;
   ASSERT_TRUE(yinit != nullptr);
 
@@ -160,6 +162,7 @@ constexpr int y = c<123 * 2>;
   auto templateSym =
       symbol_cast<VariableSymbol>(templateId->primaryTemplateSymbol);
   ASSERT_TRUE(templateSym != nullptr);
+
   auto templateDecl = getTemplateBodyAs<SimpleDeclarationAST>(
       templateSym->templateDeclaration());
   ASSERT_TRUE(templateDecl != nullptr);
@@ -173,9 +176,87 @@ constexpr int y = c<123 * 2>;
 
   auto decl = instance->initDeclaratorList->value;
   ASSERT_TRUE(decl != nullptr);
+
   auto init = ast_cast<EqualInitializerAST>(decl->initializer);
   ASSERT_TRUE(init != nullptr);
+
   auto value = interp.evaluate(init->expression);
   ASSERT_TRUE(value.has_value());
+
   ASSERT_EQ(std::visit(ArithmeticCast<int>{}, *value), 123 * 2 + 321 + 123 * 2);
 }
+
+TEST(Rewriter, Pack) {
+  auto source = R"(
+template <int... xs>
+const auto S = (... + (xs * xs));
+
+const auto N = S<0, 1, 2>;
+)"_cxx_no_templates;
+
+  auto control = source.control();
+
+  auto S = source.getAs<VariableSymbol>("S");
+  ASSERT_TRUE(S != nullptr);
+
+  auto N = source.getAs<VariableSymbol>("N");
+  ASSERT_TRUE(N != nullptr);
+  auto Ninit = ast_cast<EqualInitializerAST>(N->initializer())->expression;
+  ASSERT_TRUE(Ninit != nullptr);
+
+  auto idExpr = ast_cast<IdExpressionAST>(Ninit);
+  ASSERT_TRUE(idExpr != nullptr);
+
+  auto templateId = ast_cast<SimpleTemplateIdAST>(idExpr->unqualifiedId);
+  ASSERT_TRUE(templateId != nullptr);
+  ASSERT_EQ(templateId->primaryTemplateSymbol, S);
+
+  auto parameterPack = control->newParameterPackSymbol(nullptr, {});
+
+  ASTInterpreter interp{&source.unit};
+
+  for (auto arg : ListView{templateId->templateArgumentList}) {
+    auto exprArg = ast_cast<ExpressionTemplateArgumentAST>(arg);
+    ASSERT_TRUE(exprArg != nullptr);
+
+    auto expr = exprArg->expression;
+    auto templArg = control->newVariableSymbol(nullptr, {});
+    templArg->setInitializer(expr);
+    templArg->setType(control->add_const(expr->type));
+    templArg->setConstValue(interp.evaluate(expr));
+    if (!templArg->constValue().has_value())
+      cxx_runtime_error("template argument is not a constant expression");
+
+    parameterPack->addElement(templArg);
+  }
+
+  ASSERT_EQ(parameterPack->elements().size(), 3);
+
+  std::vector<TemplateArgument> arguments;
+  arguments.push_back(parameterPack);
+
+  auto templateDecl =
+      getTemplateBodyAs<SimpleDeclarationAST>(S->templateDeclaration());
+  ASSERT_TRUE(templateDecl != nullptr);
+
+  auto instance = substitute(source, templateDecl, arguments);
+  ASSERT_TRUE(instance != nullptr);
+
+  auto decl_to_string = [&](DeclarationAST* ast) {
+    std::ostringstream os;
+    ASTPrettyPrinter printer{&source.unit, os};
+    printer(ast);
+    return os.str();
+  };
+
+  ASSERT_EQ(decl_to_string(instance),
+            "const auto S =(xs * xs) +(xs * xs) +(xs * xs);");
+
+  auto eq = ast_cast<EqualInitializerAST>(
+      instance->initDeclaratorList->value->initializer);
+  ASSERT_TRUE(eq != nullptr);
+
+  const auto value = interp.evaluate(eq->expression);
+  ASSERT_TRUE(value.has_value());
+  ASSERT_EQ(std::visit(ArithmeticCast<int>{}, *value), 0 * 0 + 1 * 1 + 2 * 2);
+}
