[AllocToken, Clang] Implement TypeHashPointerSplit mode

clang/docs/AllocToken.rst

@@ -31,20 +31,20 @@ Token Assignment Mode
 
 The default mode to calculate tokens is:
 
-* *TypeHashPointerSplit* (mode=3): This mode assigns a token ID based on
-  the hash of the allocated type's name, where the top half ID-space is
-  reserved for types that contain pointers and the bottom half for types that
-  do not contain pointers.
+* ``typehashpointersplit``: This mode assigns a token ID based on the hash of
+  the allocated type's name, where the top half ID-space is reserved for types
+  that contain pointers and the bottom half for types that do not contain
+  pointers.
 
 Other token ID assignment modes are supported, but they may be subject to
 change or removal. These may (experimentally) be selected with ``-mllvm
 -alloc-token-mode=<mode>``:
 
-* *TypeHash* (mode=2): This mode assigns a token ID based on the hash of
-  the allocated type's name.
+* ``typehash``: This mode assigns a token ID based on the hash of the allocated
+  type's name.
 
-* *Random* (mode=1): This mode assigns a statically-determined random token ID
-  to each allocation site.
+* ``random``: This mode assigns a statically-determined random token ID to each
+  allocation site.
 
 * ``increment``: This mode assigns a simple, incrementally increasing token ID
   to each allocation site.

clang/lib/CodeGen/CGExpr.cpp

@@ -1272,6 +1272,57 @@ void CodeGenFunction::EmitBoundsCheckImpl(const Expr *E, llvm::Value *Bound,
   EmitCheck(std::make_pair(Check, CheckKind), CheckHandler, StaticData, Index);
 }
 
+static bool
+typeContainsPointer(QualType T,
+                    llvm::SmallPtrSet<const RecordDecl *, 4> &VisitedRD,
+                    bool &IncompleteType) {
+  QualType CanonicalType = T.getCanonicalType();
+  if (CanonicalType->isPointerType())
+    return true; // base case
+
+  // Look through typedef chain to check for special types.
+  for (QualType CurrentT = T; const auto *TT = CurrentT->getAs<TypedefType>();
+       CurrentT = TT->getDecl()->getUnderlyingType()) {
+    const IdentifierInfo *II = TT->getDecl()->getIdentifier();
+    if (!II)
+      continue;
+    // Special Case: Syntactically uintptr_t is not a pointer; semantically,
+    // however, very likely used as such. Therefore, classify uintptr_t as a
+    // pointer, too.
+    if (II->isStr("uintptr_t"))
+      return true;
+  }
+
+  // The type is an array; check the element type.
+  if (const ArrayType *AT = CanonicalType->getAsArrayTypeUnsafe())
+    return typeContainsPointer(AT->getElementType(), VisitedRD, IncompleteType);
+  // The type is a struct, class, or union.
+  if (const RecordDecl *RD = CanonicalType->getAsRecordDecl()) {
+    if (!RD->isCompleteDefinition()) {
+      IncompleteType = true;
+      return false;
+    }
+    if (!VisitedRD.insert(RD).second)
+      return false; // already visited
+    // Check all fields.
+    for (const FieldDecl *Field : RD->fields()) {
+      if (typeContainsPointer(Field->getType(), VisitedRD, IncompleteType))
+        return true;
+    }
+    // For C++ classes, also check base classes.
+    if (const CXXRecordDecl *CXXRD = dyn_cast<CXXRecordDecl>(RD)) {
+      // Polymorphic types require a vptr.
+      if (CXXRD->isPolymorphic())
+        return true;
+      for (const CXXBaseSpecifier &Base : CXXRD->bases()) {
+        if (typeContainsPointer(Base.getType(), VisitedRD, IncompleteType))
+          return true;
+      }
+    }
+  }
+  return false;
+}
+
 void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) {
   assert(SanOpts.has(SanitizerKind::AllocToken) &&
          "Only needed with -fsanitize=alloc-token");
@@ -1289,54 +1340,8 @@ void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) {
   // recursively check if a type contains a pointer type.
   llvm::SmallPtrSet<const RecordDecl *, 4> VisitedRD;
   bool IncompleteType = false;
-  auto TypeContainsPtr = [&](auto &&self, QualType T) -> bool {
-    QualType CanonicalType = T.getCanonicalType();
-    if (CanonicalType->isPointerType())
-      return true; // base case
-
-    // Look through typedef chain to check for special types.
-    for (QualType CurrentT = T; const auto *TT = CurrentT->getAs<TypedefType>();
-         CurrentT = TT->getDecl()->getUnderlyingType()) {
-      const IdentifierInfo *II = TT->getDecl()->getIdentifier();
-      if (!II)
-        continue;
-      // Special Case: Syntactically uintptr_t is not a pointer; semantically,
-      // however, very likely used as such. Therefore, classify uintptr_t as a
-      // pointer, too.
-      if (II->isStr("uintptr_t"))
-        return true;
-    }
-
-    // The type is an array; check the element type.
-    if (const ArrayType *AT = CanonicalType->getAsArrayTypeUnsafe())
-      return self(self, AT->getElementType());
-    // The type is a struct, class, or union.
-    if (const RecordDecl *RD = CanonicalType->getAsRecordDecl()) {
-      if (!RD->isCompleteDefinition()) {
-        IncompleteType = true;
-        return false;
-      }
-      if (!VisitedRD.insert(RD).second)
-        return false; // already visited
-      // Check all fields.
-      for (const FieldDecl *Field : RD->fields()) {
-        if (self(self, Field->getType()))
-          return true;
-      }
-      // For C++ classes, also check base classes.
-      if (const CXXRecordDecl *CXXRD = dyn_cast<CXXRecordDecl>(RD)) {
-        // Polymorphic types require a vptr.
-        if (CXXRD->isPolymorphic())
-          return true;
-        for (const CXXBaseSpecifier &Base : CXXRD->bases()) {
-          if (self(self, Base.getType()))
-            return true;
-        }
-      }
-    }
-    return false;
-  };
-  const bool ContainsPtr = TypeContainsPtr(TypeContainsPtr, AllocType);
+  const bool ContainsPtr =
+      typeContainsPointer(AllocType, VisitedRD, IncompleteType);
   if (!ContainsPtr && IncompleteType)
     return;
   auto *ContainsPtrC = Builder.getInt1(ContainsPtr);
@@ -1345,7 +1350,7 @@ void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) {
   // Format: !{<type-name>, <contains-pointer>}
   auto *MDN =
       llvm::MDNode::get(CGM.getLLVMContext(), {TypeNameMD, ContainsPtrMD});
-  CB->setMetadata(llvm::LLVMContext::MD_alloc_token_hint, MDN);
+  CB->setMetadata(llvm::LLVMContext::MD_alloc_token, MDN);
 }
 
 CodeGenFunction::ComplexPairTy CodeGenFunction::

clang/test/CodeGenCXX/alloc-token-pointer.cpp

@@ -31,25 +31,25 @@ int **test_malloc_ptr() {
 
 // CHECK-LABEL: @_Z12test_new_intv(
 int *test_new_int() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 4, i64 0){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 4, i64 0){{.*}} !alloc_token
   return new int;
 }
 
 // CHECK-LABEL: @_Z20test_new_ulong_arrayv(
 unsigned long *test_new_ulong_array() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 80, i64 0){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 80, i64 0){{.*}} !alloc_token
   return new unsigned long[10];
 }
 
 // CHECK-LABEL: @_Z12test_new_ptrv(
 int **test_new_ptr() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 8, i64 1){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 8, i64 1){{.*}} !alloc_token
   return new int*;
 }
 
 // CHECK-LABEL: @_Z18test_new_ptr_arrayv(
 int **test_new_ptr_array() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 80, i64 1){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 80, i64 1){{.*}} !alloc_token
   return new int*[10];
 }
 
@@ -108,13 +108,13 @@ ContainsPtr *test_operatornew_struct_array_with_ptr2() {
 
 // CHECK-LABEL: @_Z24test_new_struct_with_ptrv(
 ContainsPtr *test_new_struct_with_ptr() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 16, i64 1){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 16, i64 1){{.*}} !alloc_token
   return new ContainsPtr;
 }
 
 // CHECK-LABEL: @_Z30test_new_struct_array_with_ptrv(
 ContainsPtr *test_new_struct_array_with_ptr() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 160, i64 1){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 160, i64 1){{.*}} !alloc_token
   return new ContainsPtr[10];
 }
 
@@ -127,13 +127,13 @@ class TestClass {
 
 // CHECK-LABEL: @_Z14test_new_classv(
 TestClass *test_new_class() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 64, i64 0){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 64, i64 0){{.*}} !alloc_token
   return new TestClass();
 }
 
 // CHECK-LABEL: @_Z20test_new_class_arrayv(
 TestClass *test_new_class_array() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 648, i64 0){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 648, i64 0){{.*}} !alloc_token
   return new TestClass[10];
 }
 
@@ -147,13 +147,13 @@ class VirtualTestClass {
 
 // CHECK-LABEL: @_Z22test_new_virtual_classv(
 VirtualTestClass *test_new_virtual_class() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 72, i64 1){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znwm(i64 noundef 72, i64 1){{.*}} !alloc_token
   return new VirtualTestClass();
 }
 
 // CHECK-LABEL: @_Z28test_new_virtual_class_arrayv(
 VirtualTestClass *test_new_virtual_class_array() {
-  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 728, i64 1){{.*}} !alloc_token_hint
+  // CHECK: call {{.*}} ptr @__alloc_token_Znam(i64 noundef 728, i64 1){{.*}} !alloc_token
   return new VirtualTestClass[10];
 }
 

llvm/lib/Transforms/Instrumentation/AllocToken.cpp

@@ -74,27 +74,25 @@ enum class TokenMode : unsigned {
   /// reserved for types that contain pointers and the bottom half for types
   /// that do not contain pointers.
   TypeHashPointerSplit = 3,
-
-  // Mode count - keep last
-  ModeCount
 };
 
 //===--- Command-line options ---------------------------------------------===//
 
-struct ModeParser : public cl::parser<unsigned> {
-  ModeParser(cl::Option &O) : cl::parser<unsigned>(O) {}
-  bool parse(cl::Option &O, StringRef ArgName, StringRef Arg, unsigned &Value) {
-    if (cl::parser<unsigned>::parse(O, ArgName, Arg, Value))
-      return true;
-    if (Value >= static_cast<unsigned>(TokenMode::ModeCount))
-      return O.error("'" + Arg + "' value invalid");
-    return false;
-  }
-};
-
-cl::opt<unsigned, false, ModeParser>
-    ClMode("alloc-token-mode", cl::desc("Token assignment mode"), cl::Hidden,
-           cl::init(static_cast<unsigned>(TokenMode::TypeHashPointerSplit)));
+cl::opt<TokenMode> ClMode(
+    "alloc-token-mode", cl::Hidden, cl::desc("Token assignment mode"),
+    cl::init(TokenMode::TypeHashPointerSplit),
+    cl::values(
+        clEnumValN(TokenMode::Increment, "increment",
+                   "Incrementally increasing token ID"),
+        clEnumValN(TokenMode::Random, "random",
+                   "Statically-assigned random token ID"),
+        clEnumValN(TokenMode::TypeHash, "typehash",
+                   "Token ID based on allocated type hash"),
+        clEnumValN(
+            TokenMode::TypeHashPointerSplit, "typehashpointersplit",
+            "Token ID based on allocated type hash, where the top half "
+            "ID-space is reserved for types that contain pointers and the "
+            "bottom half for types that do not contain pointers. ")));
 
 cl::opt<std::string> ClFuncPrefix("alloc-token-prefix",
                                   cl::desc("The allocation function prefix"),
@@ -148,11 +146,11 @@ STATISTIC(NumAllocations, "Allocations found");
 /// Returns the !alloc_token metadata if available.
 ///
 /// Expected format is: !{<type-name>, <contains-pointer>}
-MDNode *getAllocTokenHintMetadata(const CallBase &CB) {
-  MDNode *Ret = CB.getMetadata(LLVMContext::MD_alloc_token_hint);
+MDNode *getAllocTokenMetadata(const CallBase &CB) {
+  MDNode *Ret = CB.getMetadata(LLVMContext::MD_alloc_token);
   if (!Ret)
     return nullptr;
-  assert(Ret->getNumOperands() == 2 && "bad !alloc_token_hint");
+  assert(Ret->getNumOperands() == 2 && "bad !alloc_token");
   assert(isa<MDString>(Ret->getOperand(0)));
   assert(isa<ConstantAsMetadata>(Ret->getOperand(1)));
   return Ret;
@@ -217,7 +215,7 @@ class TypeHashMode : public ModeBase {
 protected:
   std::pair<MDNode *, uint64_t> getHash(const CallBase &CB,
                                         OptimizationRemarkEmitter &ORE) {
-    if (MDNode *N = getAllocTokenHintMetadata(CB)) {
+    if (MDNode *N = getAllocTokenMetadata(CB)) {
       MDString *S = cast<MDString>(N->getOperand(0));
       return {N, xxHash64(S->getString())};
     }
@@ -287,9 +285,6 @@ class AllocToken {
     case TokenMode::TypeHashPointerSplit:
       Mode.emplace<TypeHashPointerSplitMode>(*Options.MaxTokens);
       break;
-    case TokenMode::ModeCount:
-      llvm_unreachable("");
-      break;
     }
   }