Integrate AST rewriter

Author: robertoraggi

packages/cxx-gen-ast/src/new_ast_rewriter_cc.ts

@@ -359,7 +359,7 @@ if (auto param = symbol_cast<NonTypeParameterSymbol>(ast->symbol);
       std::optional<int> index{n};
       std::swap(rewrite.elementIndex_, index);
 
-      auto expression = rewrite(ast->expression);
+      auto expression = rewrite.expression(ast->expression);
       if (!current) {
         current = expression;
       } else {

src/mlir/cxx/mlir/codegen.cc

@@ -109,6 +109,17 @@ auto Codegen::findOrCreateFunction(FunctionSymbol* functionSymbol)
   std::vector<mlir::Type> inputTypes;
   std::vector<mlir::Type> resultTypes;
 
+  if (!functionSymbol->isStatic() &&
+      functionSymbol->enclosingSymbol()->isClass()) {
+    // if it is a non static member function, we need to add the `this` pointer
+
+    auto classSymbol =
+        symbol_cast<ClassSymbol>(functionSymbol->enclosingSymbol());
+
+    inputTypes.push_back(builder_.getType<mlir::cxx::PointerType>(
+        convertType(classSymbol->type())));
+  }
+
   for (auto paramTy : functionType->parameterTypes()) {
     inputTypes.push_back(convertType(paramTy));
   }

src/mlir/cxx/mlir/codegen.h

@@ -314,8 +314,10 @@ class Codegen {
   mlir::ModuleOp module_;
   mlir::cxx::FuncOp function_;
   TranslationUnit* unit_ = nullptr;
+  mlir::Block* entryBlock_ = nullptr;
   mlir::Block* exitBlock_ = nullptr;
   mlir::cxx::AllocaOp exitValue_;
+  mlir::Value thisValue_;
   std::unordered_map<ClassSymbol*, mlir::Type> classNames_;
   std::unordered_map<Symbol*, mlir::Value> locals_;
   std::unordered_map<FunctionSymbol*, mlir::cxx::FuncOp> funcOps_;

src/mlir/cxx/mlir/codegen_declarations.cc

@@ -348,7 +348,9 @@ auto Codegen::DeclarationVisitor::operator()(FunctionDefinitionAST* ast)
 
   // Add the function body.
   auto entryBlock = gen.builder_.createBlock(&func.getBody());
-  for (const auto& input : func.getFunctionType().getInputs()) {
+  auto inputs = func.getFunctionType().getInputs();
+
+  for (const auto& input : inputs) {
     entryBlock->addArgument(input, loc);
   }
 
@@ -370,17 +372,36 @@ auto Codegen::DeclarationVisitor::operator()(FunctionDefinitionAST* ast)
 
   // function state
   std::swap(gen.function_, func);
+  std::swap(gen.entryBlock_, entryBlock);
   std::swap(gen.exitBlock_, exitBlock);
   std::swap(gen.exitValue_, exitValue);
   std::swap(gen.locals_, locals);
 
+  mlir::Value thisValue;
+
+  // if this is a non static member function, we need to allocate the `this`
+  if (!functionSymbol->isStatic() &&
+      functionSymbol->enclosingSymbol()->isClass()) {
+    auto classSymbol =
+        symbol_cast<ClassSymbol>(functionSymbol->enclosingSymbol());
+
+    auto thisType = gen.convertType(classSymbol->type());
+    auto ptrType = gen.builder_.getType<mlir::cxx::PointerType>(thisType);
+
+    thisValue = gen.newTemp(classSymbol->type(), ast->firstSourceLocation());
+
+    // store the `this` pointer in the entry block
+    gen.builder_.create<mlir::cxx::StoreOp>(
+        loc, gen.entryBlock_->getArgument(0), thisValue);
+  }
+
   FunctionParametersSymbol* params = nullptr;
   for (auto member : ast->symbol->scope()->symbols()) {
     params = symbol_cast<FunctionParametersSymbol>(member);
     if (!params) continue;
 
     int argc = 0;
-    auto args = entryBlock->getArguments();
+    auto args = gen.entryBlock_->getArguments();
     for (auto param : params->scope()->symbols()) {
       auto arg = symbol_cast<ParameterSymbol>(param);
       if (!arg) continue;
@@ -399,6 +420,8 @@ auto Codegen::DeclarationVisitor::operator()(FunctionDefinitionAST* ast)
     }
   }
 
+  std::swap(gen.thisValue_, thisValue);
+
   allocateLocals(functionSymbol);
 
   // generate code for the function body
@@ -428,7 +451,10 @@ auto Codegen::DeclarationVisitor::operator()(FunctionDefinitionAST* ast)
   }
 
   // restore the state
+  std::swap(gen.thisValue_, thisValue);
+
   std::swap(gen.function_, func);
+  std::swap(gen.entryBlock_, entryBlock);
   std::swap(gen.exitBlock_, exitBlock);
   std::swap(gen.exitValue_, exitValue);
   std::swap(gen.locals_, locals);

src/mlir/cxx/mlir/codegen_expressions.cc

@@ -299,9 +299,15 @@ auto Codegen::ExpressionVisitor::operator()(ObjectLiteralExpressionAST* ast)
 
 auto Codegen::ExpressionVisitor::operator()(ThisExpressionAST* ast)
     -> ExpressionResult {
-  auto op =
-      gen.emitTodoExpr(ast->firstSourceLocation(), to_string(ast->kind()));
-  return {op};
+  auto type = gen.convertType(ast->type);
+  auto loc = gen.getLocation(ast->firstSourceLocation());
+
+  auto loadOp =
+      gen.builder_.create<mlir::cxx::LoadOp>(loc, type, gen.thisValue_);
+  return {loadOp};
+  // auto op =
+  //     gen.emitTodoExpr(ast->firstSourceLocation(), to_string(ast->kind()));
+  // return {op};
 }
 
 auto Codegen::ExpressionVisitor::operator()(GenericSelectionExpressionAST* ast)
@@ -499,6 +505,38 @@ auto Codegen::ExpressionVisitor::operator()(CallExpressionAST* ast)
       func = nested->expression;
     }
 
+    if (auto member = ast_cast<MemberExpressionAST>(func)) {
+      auto thisValue = gen.expression(member->baseExpression);
+      auto functionSymbol = symbol_cast<FunctionSymbol>(member->symbol);
+
+      auto funcOp = gen.findOrCreateFunction(functionSymbol);
+
+      mlir::SmallVector<mlir::Value> arguments;
+      arguments.push_back(thisValue.value);
+      for (auto node : ListView{ast->expressionList}) {
+        auto value = gen.expression(node);
+        arguments.push_back(value.value);
+      }
+
+      auto loc = gen.getLocation(ast->lparenLoc);
+
+      auto functionType = type_cast<FunctionType>(functionSymbol->type());
+      mlir::SmallVector<mlir::Type> resultTypes;
+      if (!control()->is_void(functionType->returnType())) {
+        resultTypes.push_back(gen.convertType(functionType->returnType()));
+      }
+
+      auto op = gen.builder_.create<mlir::cxx::CallOp>(
+          loc, resultTypes, funcOp.getSymName(), arguments, mlir::TypeAttr{});
+
+      if (functionType->isVariadic()) {
+        op.setVarCalleeType(
+            cast<mlir::cxx::FunctionType>(gen.convertType(functionType)));
+      }
+
+      return ExpressionResult{op.getResult()};
+    }
+
     auto id = ast_cast<IdExpressionAST>(func);
     if (!id) return {};
 

src/mlir/cxx/mlir/codegen_units.cc

@@ -25,6 +25,8 @@
 #include <cxx/ast_visitor.h>
 #include <cxx/control.h>
 #include <cxx/memory_layout.h>
+#include <cxx/scope.h>
+#include <cxx/symbols.h>
 #include <cxx/translation_unit.h>
 
 // mlir
@@ -62,6 +64,38 @@ struct Codegen::UnitVisitor {
 
   auto operator()(TranslationUnitAST* ast) -> UnitResult;
   auto operator()(ModuleUnitAST* ast) -> UnitResult;
+
+  struct VisitSymbols {
+    UnitVisitor& p;
+
+    void operator()(NamespaceSymbol* symbol) {
+      for (auto member : symbol->scope()->symbols()) {
+        visit(*this, member);
+      }
+    }
+
+    void operator()(FunctionSymbol* symbol) {
+      if (auto funcDecl = symbol->declaration()) {
+        p.gen.declaration(funcDecl);
+      }
+    }
+
+    void operator()(ClassSymbol* symbol) {
+      for (auto specialization : symbol->specializations()) {
+        visit(*this, specialization.symbol);
+      }
+
+      if (!symbol->templateParameters()) {
+        for (auto member : symbol->scope()->symbols()) {
+          visit(*this, member);
+        }
+      }
+    }
+
+    void operator()(Symbol*) {
+      // Do nothing for other symbols.
+    }
+  } visitor{*this};
 };
 
 auto Codegen::operator()(UnitAST* ast) -> UnitResult {
@@ -115,6 +149,9 @@ auto Codegen::UnitVisitor::operator()(TranslationUnitAST* ast) -> UnitResult {
 
   std::swap(gen.module_, module);
 
+  visit(visitor, gen.unit_->globalScope()->owner());
+
+#if false
   ForEachExternalDefinition forEachExternalDefinition;
 
   forEachExternalDefinition.functionCallback =
@@ -125,6 +162,7 @@ auto Codegen::UnitVisitor::operator()(TranslationUnitAST* ast) -> UnitResult {
   for (auto node : ListView{ast->declarationList}) {
     forEachExternalDefinition.accept(node);
   }
+#endif
 
   std::swap(gen.module_, module);
 

src/parser/CMakeLists.txt

@@ -25,6 +25,13 @@ add_library(cxx-parser
     cxx/ast_interpreter.cc
     cxx/ast_pretty_printer.cc
     cxx/ast_printer.cc
+    cxx/ast_rewriter_declarations.cc
+    cxx/ast_rewriter_declarators.cc
+    cxx/ast_rewriter_expressions.cc
+    cxx/ast_rewriter_names.cc
+    cxx/ast_rewriter_specifiers.cc
+    cxx/ast_rewriter_statements.cc
+    cxx/ast_rewriter_units.cc
     cxx/ast_rewriter.cc
     cxx/ast_slot.cc
     cxx/ast_visitor.cc
@@ -54,7 +61,6 @@ add_library(cxx-parser
     cxx/scope.cc
     cxx/source_location.cc
     cxx/symbol_chain_view.cc
-    cxx/symbol_instantiation.cc
     cxx/symbol_printer.cc
     cxx/symbols.cc
     cxx/token.cc

src/parser/cxx/ast.h

@@ -3703,7 +3703,7 @@ class SimpleTemplateIdAST final : public UnqualifiedIdAST {
   List<TemplateArgumentAST*>* templateArgumentList = nullptr;
   SourceLocation greaterLoc;
   const Identifier* identifier = nullptr;
-  Symbol* primaryTemplateSymbol = nullptr;
+  Symbol* symbol = nullptr;
 
   void accept(ASTVisitor* visitor) override { visitor->visit(this); }