From f007235ad4afd761236e69fb3fd4673e0e08e727 Mon Sep 17 00:00:00 2001
From: v01dxyz <v01dxyz@v01d.xyz>
Date: Thu, 12 Dec 2024 01:59:15 +0100
Subject: [PATCH] Make getFunctionTypeLoc support AttributedType

---
 clang/include/clang/AST/TypeLoc.h |  4 ++++
 clang/lib/AST/Decl.cpp            | 19 +++++++++++++++++--
 clang/lib/Sema/SemaDeclAttr.cpp   | 15 +--------------
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/clang/include/clang/AST/TypeLoc.h b/clang/include/clang/AST/TypeLoc.h
index d2d4f1b6e2e0..54cc13f02385 100644
--- a/clang/include/clang/AST/TypeLoc.h
+++ b/clang/include/clang/AST/TypeLoc.h
@@ -883,6 +883,10 @@ class AttributedTypeLoc : public ConcreteTypeLoc<UnqualTypeLoc,
     return getInnerTypeLoc();
   }
 
+  TypeLoc getEquivalentTypeLoc() const {
+    return TypeLoc(getTypePtr()->getEquivalentType(), getNonLocalData());
+  }
+
   /// The type attribute.
   const Attr *getAttr() const {
     return getLocalData()->TypeAttr;
diff --git a/clang/lib/AST/Decl.cpp b/clang/lib/AST/Decl.cpp
index 85fe7d63775d..e897b9946e72 100644
--- a/clang/lib/AST/Decl.cpp
+++ b/clang/lib/AST/Decl.cpp
@@ -3745,8 +3745,23 @@ bool FunctionDecl::doesDeclarationForceExternallyVisibleDefinition() const {
 
 FunctionTypeLoc FunctionDecl::getFunctionTypeLoc() const {
   const TypeSourceInfo *TSI = getTypeSourceInfo();
-  return TSI ? TSI->getTypeLoc().IgnoreParens().getAs<FunctionTypeLoc>()
-             : FunctionTypeLoc();
+
+  if (!TSI)
+    return FunctionTypeLoc();
+
+  TypeLoc TL = TSI->getTypeLoc();
+  FunctionTypeLoc FTL;
+
+  while (!(FTL = TL.getAs<FunctionTypeLoc>())) {
+    if (auto PTL = TL.getAs<ParenTypeLoc>())
+      TL = PTL.getInnerLoc();
+    else if (auto ATL = TL.getAs<AttributedTypeLoc>())
+      TL = ATL.getEquivalentTypeLoc();
+    else
+      break;
+  }
+
+  return FTL;
 }
 
 SourceRange FunctionDecl::getReturnTypeSourceRange() const {
diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index 8edbb33865d3..247287199173 100644
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -2413,23 +2413,10 @@ static void handleCHERICompartmentName(Sema &S, Decl *D, const ParsedAttr &Attr,
   if (FD && FD->getReturnType()->isVoidType()) {
     S.Diag(Attr.getLoc(), diag::warn_cheri_compartment_void_return_type);
 
-    const TypeSourceInfo *TSI = FD->getTypeSourceInfo();
-    TypeLoc TL = TSI->getTypeLoc().IgnoreParens();
-
-    // ignore function type attributes
-    while (auto ATL = TL.getAs<AttributedTypeLoc>())
-      TL = ATL.getModifiedLoc();
-
-    if (auto FTL = TL.getAs<FunctionTypeLoc>()) {
-      SourceRange SR = FTL.getReturnLoc().getSourceRange();
+    if (SourceRange SR = FD->getReturnTypeSourceRange(); SR.isValid()) {
       S.Diag(SR.getBegin(), diag::note_cheri_compartment_void_return_type)
           << FixItHint::CreateReplacement(SR, "int");
     }
-    //    FunctionTypeLoc FTL = FD->getFunctionTypeLoc();
-
-    //    TypeLoc RL = FTL.getReturnLoc();
-    //    SourceRange SR = FD->getReturnTypeSourceRange();
-
   } else
     D->addAttr(::new (S.Context) WarnUnusedResultAttr(
         S.Context, Attr, "CHERI compartment call"));