[clang][RISCV] Introduce CodeGenModule::calcRISCVZicfilpFuncSigLabel()

This method calculates a CFI label used in the RISC-V Zicfilp func-sig CFI
scheme for a given function type/declaration. The scheme, according to psABI,
encodes the label based on function signature, and the rules are modified from
the Itanium C++ ABI mangling rule to allow functions (callees) that are called
indirectly to have the expected label as indicated by the function pointer type
seen at the call site (caller).

Author: mylai-mtk

clang/include/clang/AST/ASTContext.h

@@ -1876,6 +1876,9 @@ class ASTContext : public RefCountedBase<ASTContext> {
   /// (struct/union/class/enum) decl.
   QualType getTagDeclType(const TagDecl *Decl) const;
 
+  /// Return the type for "void *"
+  QualType getVoidPtrType() const { return VoidPtrTy; }
+
   /// Return the unique type for "size_t" (C99 7.17), defined in
   /// <stddef.h>.
   ///
@@ -1903,6 +1906,10 @@ class ASTContext : public RefCountedBase<ASTContext> {
   /// defined in <stddef.h> as defined by the target.
   QualType getWideCharType() const { return WideCharTy; }
 
+  /// Return the type of wide characters in C context, no matter whether it's C
+  /// or C++ being compiled.
+  QualType getWCharTypeInC() const;
+
   /// Return the type of "signed wchar_t".
   ///
   /// Used when in C++, as a GCC extension.

clang/include/clang/AST/Mangle.h

@@ -212,6 +212,11 @@ class ItaniumMangleContext : public MangleContext {
 
   virtual void mangleModuleInitializer(const Module *Module, raw_ostream &) = 0;
 
+  virtual void mangleForRISCVZicfilpFuncSigLabel(const FunctionType &FT,
+                                                 const bool IsCXXInstanceMethod,
+                                                 const bool IsCXXVirtualMethod,
+                                                 raw_ostream &) = 0;
+
   // This has to live here, otherwise the CXXNameMangler won't have access to
   // it.
   virtual DiscriminatorOverrideTy getDiscriminatorOverride() const = 0;

clang/lib/AST/ASTContext.cpp

@@ -6442,6 +6442,12 @@ CanQualType ASTContext::getUIntMaxType() const {
   return getFromTargetType(Target->getUIntMaxType());
 }
 
+/// Return the type of wide characters in C context, no matter whether it's C
+/// or C++ being compiled.
+QualType ASTContext::getWCharTypeInC() const {
+  return getFromTargetType(Target->getWCharType());
+}
+
 /// getSignedWCharType - Return the type of "signed wchar_t".
 /// Used when in C++, as a GCC extension.
 QualType ASTContext::getSignedWCharType() const {

clang/lib/AST/ItaniumMangle.cpp

@@ -43,6 +43,14 @@ using namespace clang;
 
 namespace {
 
+static bool mayBeCovariant(const Type &Ty) {
+  if (auto *const PT = Ty.getAs<PointerType>())
+    return PT->getPointeeType()->isStructureOrClassType();
+  if (auto *const RT = Ty.getAs<ReferenceType>())
+    return RT->getPointeeType()->isStructureOrClassType();
+  return false;
+}
+
 static bool isLocalContainerContext(const DeclContext *DC) {
   return isa<FunctionDecl>(DC) || isa<ObjCMethodDecl>(DC) || isa<BlockDecl>(DC);
 }
@@ -136,6 +144,11 @@ class ItaniumMangleContextImpl : public ItaniumMangleContext {
 
   void mangleModuleInitializer(const Module *Module, raw_ostream &) override;
 
+  void mangleForRISCVZicfilpFuncSigLabel(const FunctionType &,
+                                         const bool IsCXXInstanceMethod,
+                                         const bool IsCXXVirtualMethod,
+                                         raw_ostream &) override;
+
   bool getNextDiscriminator(const NamedDecl *ND, unsigned &disc) {
     // Lambda closure types are already numbered.
     if (isLambda(ND))
@@ -395,8 +408,10 @@ class CXXNameMangler {
   llvm::DenseMap<uintptr_t, unsigned> Substitutions;
   llvm::DenseMap<StringRef, unsigned> ModuleSubstitutions;
 
+protected:
   ASTContext &getASTContext() const { return Context.getASTContext(); }
 
+private:
   bool isCompatibleWith(LangOptions::ClangABI Ver) {
     return Context.getASTContext().getLangOpts().getClangABICompat() <= Ver;
   }
@@ -443,6 +458,8 @@ class CXXNameMangler {
     NullOut = true;
   }
 
+  virtual ~CXXNameMangler() = default;
+
   struct WithTemplateDepthOffset { unsigned Offset; };
   CXXNameMangler(ItaniumMangleContextImpl &C, raw_ostream &Out,
                  WithTemplateDepthOffset Offset)
@@ -559,9 +576,12 @@ class CXXNameMangler {
                                       StringRef Prefix = "");
   void mangleOperatorName(DeclarationName Name, unsigned Arity);
   void mangleOperatorName(OverloadedOperatorKind OO, unsigned Arity);
+
+protected:
   void mangleQualifiers(Qualifiers Quals, const DependentAddressSpaceType *DAST = nullptr);
   void mangleRefQualifier(RefQualifierKind RefQualifier);
 
+private:
   void mangleObjCMethodName(const ObjCMethodDecl *MD);
 
   // Declare manglers for every type class.
@@ -572,11 +592,24 @@ class CXXNameMangler {
 
   void mangleType(const TagType*);
   void mangleType(TemplateName);
+
+protected:
+  // Use the `Impl` scheme instead of directly virtualizing `mangleType`s since
+  // `mangleType`s are declared by tables
+  virtual void mangleTypeImpl(const BuiltinType *T);
+  virtual void mangleTypeImpl(const FunctionProtoType *T);
+  virtual void mangleTypeImpl(const FunctionNoProtoType *T);
+
+private:
   static StringRef getCallingConvQualifierName(CallingConv CC);
   void mangleExtParameterInfo(FunctionProtoType::ExtParameterInfo info);
   void mangleExtFunctionInfo(const FunctionType *T);
+
+protected:
   void mangleBareFunctionType(const FunctionProtoType *T, bool MangleReturnType,
                               const FunctionDecl *FD = nullptr);
+
+private:
   void mangleNeonVectorType(const VectorType *T);
   void mangleNeonVectorType(const DependentVectorType *T);
   void mangleAArch64NeonVectorType(const VectorType *T);
@@ -3058,7 +3091,9 @@ void CXXNameMangler::mangleCXXRecordDecl(const CXXRecordDecl *Record) {
   addSubstitution(Record);
 }
 
-void CXXNameMangler::mangleType(const BuiltinType *T) {
+void CXXNameMangler::mangleType(const BuiltinType *T) { mangleTypeImpl(T); }
+
+void CXXNameMangler::mangleTypeImpl(const BuiltinType *T) {
   //  <type>         ::= <builtin-type>
   //  <builtin-type> ::= v  # void
   //                 ::= w  # wchar_t
@@ -3563,10 +3598,14 @@ CXXNameMangler::mangleExtParameterInfo(FunctionProtoType::ExtParameterInfo PI) {
     mangleVendorQualifier("noescape");
 }
 
+void CXXNameMangler::mangleType(const FunctionProtoType *T) {
+  return mangleTypeImpl(T);
+}
+
 // <type>          ::= <function-type>
 // <function-type> ::= [<CV-qualifiers>] F [Y]
 //                      <bare-function-type> [<ref-qualifier>] E
-void CXXNameMangler::mangleType(const FunctionProtoType *T) {
+void CXXNameMangler::mangleTypeImpl(const FunctionProtoType *T) {
   mangleExtFunctionInfo(T);
 
   // Mangle CV-qualifiers, if present.  These are 'this' qualifiers,
@@ -3604,6 +3643,10 @@ void CXXNameMangler::mangleType(const FunctionProtoType *T) {
 }
 
 void CXXNameMangler::mangleType(const FunctionNoProtoType *T) {
+  return mangleTypeImpl(T);
+}
+
+void CXXNameMangler::mangleTypeImpl(const FunctionNoProtoType *T) {
   // Function types without prototypes can arise when mangling a function type
   // within an overloadable function in C. We mangle these as the absence of any
   // parameter types (not even an empty parameter list).
@@ -7074,6 +7117,85 @@ bool CXXNameMangler::shouldHaveAbiTags(ItaniumMangleContextImpl &C,
   return TrackAbiTags.AbiTagsRoot.getUsedAbiTags().size();
 }
 
+namespace {
+
+class RISCVZicfilpFuncSigLabelMangler : public CXXNameMangler {
+  bool IsTopLevelAndCXXVirtualMethod;
+
+public:
+  RISCVZicfilpFuncSigLabelMangler(ItaniumMangleContextImpl &C, raw_ostream &Out,
+                                  const bool IsCXXVirtualMethod)
+      : CXXNameMangler(C, Out),
+        IsTopLevelAndCXXVirtualMethod(/*IsTopLevel=*/true &&
+                                      IsCXXVirtualMethod) {}
+
+  void mangleTypeImpl(const BuiltinType *T) override {
+    if (T->getKind() == BuiltinType::WChar_S ||
+        T->getKind() == BuiltinType::WChar_U) {
+      const Type *const OverrideT =
+          getASTContext().getWCharTypeInC().getTypePtr();
+      assert(isa<BuiltinType>(OverrideT) &&
+             "`wchar_t' in C is expected to be defined to a built-in type");
+      T = static_cast<const BuiltinType *>(OverrideT);
+    }
+    return CXXNameMangler::mangleTypeImpl(T);
+  }
+
+  // This <function-type> is the RISC-V psABI modified version
+  // <function-type> ::= [<CV-qualifiers>] [Dx] F <bare-function-type>
+  //                     [<ref-qualifier>] E
+  void mangleTypeImpl(const FunctionProtoType *T) override {
+    // Mangle CV-qualifiers, if present.  These are 'this' qualifiers,
+    // e.g. "const" in "int (A::*)() const".
+    mangleQualifiers(T->getMethodQuals());
+
+    getStream() << 'F';
+
+    bool MangleReturnType = true;
+    if (const Type &RetT = *T->getReturnType().getTypePtr();
+        IsTopLevelAndCXXVirtualMethod && mayBeCovariant(RetT)) {
+      // Possible covariant types mangle dummy cv-unqualified `class v` as its
+      // class type
+      if (RetT.isPointerType())
+        getStream() << "P1v";
+      else if (RetT.isLValueReferenceType())
+        getStream() << "R1v";
+      else {
+        assert(RetT.isRValueReferenceType() &&
+               "Expect an r-value ref for covariant return type that is not a "
+               "pointer or an l-value ref");
+        getStream() << "O1v";
+      }
+
+      IsTopLevelAndCXXVirtualMethod = false; // Not top-level anymore
+      MangleReturnType = false;
+    }
+    mangleBareFunctionType(T, MangleReturnType);
+
+    // Mangle the ref-qualifier, if present.
+    mangleRefQualifier(T->getRefQualifier());
+
+    getStream() << 'E';
+  }
+
+  void mangleTypeImpl(const FunctionNoProtoType *T) override {
+    return CXXNameMangler::mangleTypeImpl(toFunctionProtoType(T));
+  }
+
+private:
+  const FunctionProtoType *
+  toFunctionProtoType(const FunctionNoProtoType *const T) {
+    FunctionProtoType::ExtProtoInfo EPI;
+    EPI.ExtInfo = T->getExtInfo();
+    const Type *const NewT = getASTContext()
+                                 .getFunctionType(T->getReturnType(), {}, EPI)
+                                 .getTypePtr();
+    return static_cast<const FunctionProtoType *>(NewT);
+  }
+}; // class RISCVZicfilpFuncSigLabelMangler
+
+} // anonymous namespace
+
 //
 
 /// Mangles the name of the declaration D and emits that name to the given
@@ -7412,6 +7534,17 @@ void ItaniumMangleContextImpl::mangleModuleInitializer(const Module *M,
   }
 }
 
+void ItaniumMangleContextImpl::mangleForRISCVZicfilpFuncSigLabel(
+    const FunctionType &FT, const bool IsCXXInstanceMethod,
+    const bool IsCXXVirtualMethod, raw_ostream &Out) {
+  if (IsCXXInstanceMethod)
+    // member methods uses a dummy class named `v` in place of real classes
+    Out << "M1v";
+
+  RISCVZicfilpFuncSigLabelMangler Mangler(*this, Out, IsCXXVirtualMethod);
+  Mangler.mangleType(QualType(&FT, 0));
+}
+
 ItaniumMangleContext *ItaniumMangleContext::create(ASTContext &Context,
                                                    DiagnosticsEngine &Diags,
                                                    bool IsAux) {

clang/lib/CodeGen/CodeGenModule.cpp

@@ -2829,6 +2829,56 @@ void CodeGenModule::CreateFunctionTypeMetadataForIcall(const FunctionDecl *FD,
       F->addTypeMetadata(0, llvm::ConstantAsMetadata::get(CrossDsoTypeId));
 }
 
+uint32_t
+CodeGenModule::calcRISCVZicfilpFuncSigLabel(const FunctionType &FT,
+                                            const bool IsCXXInstanceMethod,
+                                            const bool IsCXXVirtualMethod) {
+  std::string OutName;
+  llvm::raw_string_ostream Out(OutName);
+  MangleContext &MC = getCXXABI().getMangleContext();
+  cast<ItaniumMangleContext>(MC).mangleForRISCVZicfilpFuncSigLabel(
+      FT, IsCXXInstanceMethod, IsCXXVirtualMethod, Out);
+
+  const llvm::MD5::MD5Result MD5Result =
+      llvm::MD5::hash({(const uint8_t *)OutName.data(), OutName.size()});
+
+  // Take 20 bits of MD5 result with the approach specified by psABI
+  uint64_t MD5High = MD5Result.high();
+  uint64_t MD5Low = MD5Result.low();
+  while (MD5High && MD5Low) {
+    const uint32_t Low20Bits = MD5Low & 0xFFFFFULL;
+    if (Low20Bits)
+      return Low20Bits;
+
+    // Logical right shift MD5 result by 20 bits
+    MD5Low = (MD5High & 0xFFFFF) << 44 | MD5Low >> 20;
+    MD5High >>= 20;
+  }
+
+  return llvm::MD5Hash("RISC-V") & 0xFFFFFULL;
+}
+
+uint32_t CodeGenModule::calcRISCVZicfilpFuncSigLabel(const FunctionDecl &FD) {
+  if (FD.isMain())
+    // All main functions use `int main(int, char**)` for label calculation
+    // according to psABI spec. This value is the pre-calculated label.
+    return 0xd0639;
+
+  if (isa<CXXDestructorDecl>(FD))
+    // All destructors use `void (void*)` for label calculation according to the
+    // psABI spec. This value is the pre-calculated label.
+    return 0x639c2;
+
+  bool IsCXXInstanceMethod = false;
+  bool IsCXXVirtualMethod = false;
+  if (const auto *const MD = dyn_cast<CXXMethodDecl>(&FD)) {
+    IsCXXInstanceMethod = MD->isInstance();
+    IsCXXVirtualMethod = MD->isVirtual();
+  }
+  return calcRISCVZicfilpFuncSigLabel(*FD.getFunctionType(),
+                                      IsCXXInstanceMethod, IsCXXVirtualMethod);
+}
+
 void CodeGenModule::setKCFIType(const FunctionDecl *FD, llvm::Function *F) {
   llvm::LLVMContext &Ctx = F->getContext();
   llvm::MDBuilder MDB(Ctx);

clang/lib/CodeGen/CodeGenModule.h

@@ -1559,6 +1559,14 @@ class CodeGenModule : public CodeGenTypeCache {
   /// Emit KCFI type identifier constants and remove unused identifiers.
   void finalizeKCFITypes();
 
+  /// Calculate RISC-V Zicfilp func-sig scheme CFI label
+  uint32_t calcRISCVZicfilpFuncSigLabel(const FunctionType &FT,
+                                        const bool IsCXXInstanceMethod,
+                                        const bool IsCXXVirtualMethod);
+
+  /// Calculate RISC-V Zicfilp func-sig scheme CFI label
+  uint32_t calcRISCVZicfilpFuncSigLabel(const FunctionDecl &FD);
+
   /// Whether this function's return type has no side effects, and thus may
   /// be trivially discarded if it is unused.
   bool MayDropFunctionReturn(const ASTContext &Context,