[llvm-branch-commits] [llvm] [AllocToken, Clang] Implement __builtin_infer_alloc_token() and llvm.alloc.token.id (PR #156842)

[Vitaly Buka via llvm-branch-commits]

================
@@ -334,21 +348,30 @@ bool AllocToken::instrumentFunction(Function &F) {
   // Do not apply any instrumentation for naked functions.
   if (F.hasFnAttribute(Attribute::Naked))
     return false;
-  if (F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation))
-    return false;
   // Don't touch available_externally functions, their actual body is elsewhere.
   if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage)
     return false;
-  // Only instrument functions that have the sanitize_alloc_token attribute.
-  if (!F.hasFnAttribute(Attribute::SanitizeAllocToken))
-    return false;
 
   auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
   auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
   SmallVector<std::pair<CallBase *, LibFunc>, 4> AllocCalls;
+  SmallVector<IntrinsicInst *, 4> IntrinsicInsts;
 
   // Collect all allocation calls to avoid iterator invalidation.
   for (Instruction &I : instructions(F)) {
+    // Collect all alloc_token_* intrinsics.
+    if (auto *II = dyn_cast<IntrinsicInst>(&I);
+        II && II->getIntrinsicID() == Intrinsic::alloc_token_id) {
+      IntrinsicInsts.emplace_back(II);
+      continue;
+    }
+
+    // Only instrument functions that have the sanitize_alloc_token attribute.
----------------
vitalybuka wrote:

It stead of continue on each iteration of look, cleanner would be

Maybe move outside of the loop and store into 

```
bool InstrumentFuonction = !F.hasFnAttribute(Attribute::SanitizeAllocToken) && !F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation);

```

https://github.com/llvm/llvm-project/pull/156842