[CIR] Make AST Function attributes constrained.

Author: xlauko

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -1590,7 +1590,7 @@ def FuncOp : CIR_Op<"func", [
                        OptionalAttr<DictArrayAttr>:$arg_attrs,
                        OptionalAttr<DictArrayAttr>:$res_attrs,
                        OptionalAttr<FlatSymbolRefAttr>:$aliasee,
-                       OptionalAttr<AnyAttr>:$ast);
+                       OptionalAttr<AnyASTFunctionDeclAttr>:$ast);
   let regions = (region AnyRegion:$body);
   let skipDefaultBuilders = 1;
 

clang/include/clang/CIR/Interfaces/ASTAttrInterfaces.h

@@ -21,6 +21,9 @@ namespace cir {
 mlir::Attribute makeAstDeclAttr(const clang::Decl *decl,
                                 mlir::MLIRContext *ctx);
 
+mlir::Attribute makeFuncDeclAttr(const clang::Decl *decl,
+                                 mlir::MLIRContext *ctx);
+
 } // namespace cir
 } // namespace mlir
 

clang/include/clang/CIR/Interfaces/ASTAttrInterfaces.td

@@ -190,6 +190,16 @@ let cppNamespace = "::mlir::cir" in {
       AttrInterface<"ASTClassTemplatePartialSpecializationDeclInterface",
                     [ASTClassTemplateSpecializationDeclInterface]>;
 
+
+  def AnyASTFunctionDeclAttr : Attr<
+      CPred<"::mlir::isa<::mlir::cir::ASTFunctionDeclInterface>($_self)">,
+            "AST Function attribute"> {
+    let storageType = "::mlir::Attribute";
+    let returnType = "::mlir::Attribute";
+    let convertFromStorage = "$_self";
+    let constBuilderCall = "$0";
+  }
+
 } // namespace mlir::cir
 
 #endif // MLIR_CIR_INTERFACES_AST_ATTR_INTERFACES

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -1822,7 +1822,7 @@ CIRGenModule::createCIRFunction(mlir::Location loc, StringRef name,
     f = builder.create<mlir::cir::FuncOp>(loc, name, Ty);
 
     if (FD)
-      f.setAstAttr(makeAstDeclAttr(FD, builder.getContext()));
+      f.setAstAttr(makeFuncDeclAttr(FD, builder.getContext()));
 
     if (FD && !FD->hasPrototype())
       f.setNoProtoAttr(builder.getUnitAttr());

clang/lib/CIR/Dialect/IR/CIRAttrs.cpp

@@ -46,8 +46,11 @@ using namespace mlir::cir;
 // CIR AST Attr helpers
 //===----------------------------------------------------------------------===//
 
-mlir::Attribute mlir::cir::makeAstDeclAttr(const clang::Decl *decl,
-                                           mlir::MLIRContext *ctx) {
+namespace mlir {
+namespace cir {
+
+mlir::Attribute makeAstDeclAttr(const clang::Decl *decl,
+                                mlir::MLIRContext *ctx) {
   if (auto ast = clang::dyn_cast<clang::CXXConstructorDecl>(decl))
     return ASTCXXConstructorDeclAttr::get(ctx, ast);
   if (auto ast = clang::dyn_cast<clang::CXXConversionDecl>(decl))
@@ -78,6 +81,33 @@ mlir::Attribute mlir::cir::makeAstDeclAttr(const clang::Decl *decl,
   return ASTDeclAttr::get(ctx, decl);
 };
 
+mlir::Attribute makeFuncDeclAttr(const clang::Decl *decl,
+                                 mlir::MLIRContext *ctx) {
+  return llvm::TypeSwitch<const clang::Decl *, mlir::Attribute>(decl)
+      .Case([ctx](const clang::CXXConstructorDecl *ast) {
+        return ASTCXXConstructorDeclAttr::get(ctx, ast);
+      })
+      .Case([ctx](const clang::CXXConversionDecl *ast) {
+        return ASTCXXConversionDeclAttr::get(ctx, ast);
+      })
+      .Case([ctx](const clang::CXXDestructorDecl *ast) {
+        return ASTCXXDestructorDeclAttr::get(ctx, ast);
+      })
+      .Case([ctx](const clang::CXXMethodDecl *ast) {
+        return ASTCXXMethodDeclAttr::get(ctx, ast);
+      })
+      .Case([ctx](const clang::FunctionDecl *ast) {
+        return ASTFunctionDeclAttr::get(ctx, ast);
+      })
+      .Default([](auto) {
+        llvm_unreachable("unexpected Decl kind");
+        return mlir::Attribute();
+      });
+}
+
+} // namespace cir
+} // namespace mlir
+
 //===----------------------------------------------------------------------===//
 // General CIR parsing / printing
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -42,6 +42,7 @@ using namespace mlir::cir;
 #include "clang/CIR/Dialect/IR/CIROpsStructs.cpp.inc"
 
 #include "clang/CIR/Dialect/IR/CIROpsDialect.cpp.inc"
+#include "clang/CIR/Interfaces/ASTAttrInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // CIR Dialect