1//===- AllocToken.cpp - Allocation token instrumentation ------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements AllocToken, an instrumentation pass that
10// replaces allocation calls with token-enabled versions.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/Transforms/Instrumentation/AllocToken.h"
15#include "InstrumentationOptions.h"
16#include "llvm/ADT/DenseMap.h"
17#include "llvm/ADT/SmallVector.h"
18#include "llvm/ADT/Statistic.h"
19#include "llvm/ADT/StringExtras.h"
20#include "llvm/ADT/StringRef.h"
21#include "llvm/Analysis/MemoryBuiltins.h"
22#include "llvm/Analysis/OptimizationRemarkEmitter.h"
23#include "llvm/Analysis/TargetLibraryInfo.h"
24#include "llvm/IR/Analysis.h"
25#include "llvm/IR/Attributes.h"
26#include "llvm/IR/Constants.h"
27#include "llvm/IR/DerivedTypes.h"
28#include "llvm/IR/Function.h"
29#include "llvm/IR/GlobalValue.h"
30#include "llvm/IR/IRBuilder.h"
31#include "llvm/IR/InstIterator.h"
32#include "llvm/IR/InstrTypes.h"
33#include "llvm/IR/Instructions.h"
34#include "llvm/IR/IntrinsicInst.h"
35#include "llvm/IR/Metadata.h"
36#include "llvm/IR/Module.h"
37#include "llvm/IR/PassManager.h"
38#include "llvm/IR/Type.h"
39#include "llvm/Support/AllocToken.h"
40#include "llvm/Support/Casting.h"
41#include "llvm/Support/Compiler.h"
42#include "llvm/Support/ErrorHandling.h"
43#include "llvm/Support/RandomNumberGenerator.h"
44#include "llvm/Support/SipHash.h"
45#include <cassert>
46#include <cstdint>
47#include <memory>
48#include <optional>
49#include <string>
50#include <utility>
51#include <variant>
52
53using namespace llvm;
54using TokenMode = AllocTokenMode;
55
56#define DEBUG_TYPE "alloc-token"
57
58namespace {
59
60//===--- Statistics -------------------------------------------------------===//
61
62STATISTIC(NumFunctionsModified, "Functions modified");
63STATISTIC(NumAllocationsInstrumented, "Allocations instrumented");
64
65//===----------------------------------------------------------------------===//
66
67/// Returns the !alloc_token metadata if available.
68///
69/// Expected format is: !{<type-name>, <contains-pointer>[, <function-name>]}
70MDNode *getAllocTokenMetadata(const CallBase &CB) {
71 MDNode *Ret = nullptr;
72 if (auto *II = dyn_cast<IntrinsicInst>(Val: &CB);
73 II && II->getIntrinsicID() == Intrinsic::alloc_token_id) {
74 auto *MDV = cast<MetadataAsValue>(Val: II->getArgOperand(i: 0));
75 Ret = cast<MDNode>(Val: MDV->getMetadata());
76 // If the intrinsic has an empty MDNode, type inference failed.
77 if (Ret->getNumOperands() == 0)
78 return nullptr;
79 } else {
80 Ret = CB.getMetadata(KindID: LLVMContext::MD_alloc_token);
81 if (!Ret)
82 return nullptr;
83 }
84 assert((Ret->getNumOperands() == 2 || Ret->getNumOperands() == 3) &&
85 "bad !alloc_token");
86 assert(isa<MDString>(Ret->getOperand(0)));
87 assert(isa<ConstantAsMetadata>(Ret->getOperand(1)));
88 assert(Ret->getNumOperands() == 2 || isa<MDString>(Ret->getOperand(2)));
89 return Ret;
90}
91
92bool containsPointer(const MDNode *MD) {
93 ConstantAsMetadata *C = cast<ConstantAsMetadata>(Val: MD->getOperand(I: 1));
94 auto *CI = cast<ConstantInt>(Val: C->getValue());
95 return CI->getValue().getBoolValue();
96}
97
98class ModeBase {
99public:
100 explicit ModeBase(const IntegerType &TokenTy, uint64_t MaxTokens)
101 : MaxTokens(MaxTokens ? MaxTokens : TokenTy.getBitMask()) {
102 assert(MaxTokens <= TokenTy.getBitMask());
103 }
104
105protected:
106 uint64_t boundedToken(uint64_t Val) const {
107 assert(MaxTokens != 0);
108 return Val % MaxTokens;
109 }
110
111 const uint64_t MaxTokens;
112};
113
114/// Implementation for TokenMode::Increment.
115class IncrementMode : public ModeBase {
116public:
117 using ModeBase::ModeBase;
118
119 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &) {
120 return boundedToken(Val: Counter++);
121 }
122
123private:
124 uint64_t Counter = 0;
125};
126
127/// Implementation for TokenMode::Random.
128class RandomMode : public ModeBase {
129public:
130 RandomMode(const IntegerType &TokenTy, uint64_t MaxTokens,
131 std::unique_ptr<RandomNumberGenerator> RNG)
132 : ModeBase(TokenTy, MaxTokens), RNG(std::move(RNG)) {}
133 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &) {
134 return boundedToken(Val: (*RNG)());
135 }
136
137private:
138 std::unique_ptr<RandomNumberGenerator> RNG;
139};
140
141/// Implementation for TokenMode::TypeHash. The implementation ensures
142/// hashes are stable across different compiler invocations. Uses SipHash as the
143/// hash function.
144class TypeHashMode : public ModeBase {
145public:
146 TypeHashMode(const IntegerType &TokenTy, uint64_t MaxTokens,
147 uint64_t Fallback)
148 : ModeBase(TokenTy, MaxTokens), Fallback(Fallback) {}
149
150 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
151
152 if (MDNode *N = getAllocTokenMetadata(CB)) {
153 MDString *S = cast<MDString>(Val: N->getOperand(I: 0));
154 AllocTokenMetadata Metadata{.TypeName: S->getString(), .ContainsPointer: containsPointer(MD: N)};
155 if (auto Token = getAllocToken(Mode: TokenMode::TypeHash, Metadata, MaxTokens))
156 return *Token;
157 }
158 // Fallback.
159 remarkNoMetadata(CB, ORE);
160 return Fallback;
161 }
162
163protected:
164 /// Remark that there was no precise type information.
165 static void remarkNoMetadata(const CallBase &CB,
166 OptimizationRemarkEmitter &ORE) {
167 ORE.emit(RemarkBuilder: [&] {
168 ore::NV FuncNV("Function", CB.getParent()->getParent());
169 const Function *Callee = CB.getCalledFunction();
170 ore::NV CalleeNV("Callee", Callee ? Callee->getName() : "<unknown>");
171 return OptimizationRemark(DEBUG_TYPE, "NoAllocToken", &CB)
172 << "Call to '" << CalleeNV << "' in '" << FuncNV
173 << "' without source-level type token";
174 });
175 }
176
177 const uint64_t Fallback;
178};
179
180/// Implementation for TokenMode::TypeHashPointerSplit.
181class TypeHashPointerSplitMode : public TypeHashMode {
182public:
183 using TypeHashMode::TypeHashMode;
184
185 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
186 if (MDNode *N = getAllocTokenMetadata(CB)) {
187 MDString *S = cast<MDString>(Val: N->getOperand(I: 0));
188 AllocTokenMetadata Metadata{.TypeName: S->getString(), .ContainsPointer: containsPointer(MD: N)};
189 if (auto Token = getAllocToken(Mode: TokenMode::TypeHashPointerSplit, Metadata,
190 MaxTokens))
191 return *Token;
192 }
193 // Pick the fallback token, which by default is 0, meaning it'll fall into
194 // the pointer-less bucket. Override by setting -alloc-token-fallback if
195 // that is the wrong choice.
196 remarkNoMetadata(CB, ORE);
197 return Fallback;
198 }
199};
200
201/// Implementation for TokenMode::TypeFuncHash and
202/// TokenMode::TypeFuncHashPointerSplit.
203class TypeFuncHashMode : public TypeHashMode {
204public:
205 TypeFuncHashMode(const IntegerType &TokenTy, uint64_t MaxTokens,
206 uint64_t Fallback, TokenMode Mode)
207 : TypeHashMode(TokenTy, MaxTokens, Fallback), Mode(Mode) {}
208
209 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
210 MDNode *N = getAllocTokenMetadata(CB);
211 if (!N) {
212 remarkNoMetadata(CB, ORE);
213 return Fallback;
214 }
215 // Generated for another mode, whose tokens may already be in the module.
216 if (N->getNumOperands() != 3) {
217 CB.getContext().emitError(
218 I: &CB, ErrorStr: "!alloc_token without function name is incompatible with mode " +
219 getAllocTokenModeAsString(Mode));
220 return Fallback;
221 }
222 AllocTokenMetadata Metadata{.TypeName: cast<MDString>(Val: N->getOperand(I: 0))->getString(),
223 .ContainsPointer: containsPointer(MD: N),
224 .FunctionName: cast<MDString>(Val: N->getOperand(I: 2))->getString()};
225 if (Metadata.TypeName.empty())
226 remarkNoMetadata(CB, ORE);
227 return *getAllocToken(Mode, Metadata, MaxTokens);
228 }
229
230private:
231 const TokenMode Mode;
232};
233
234// Apply opt overrides and module flags.
235static AllocTokenOptions resolveOptions(const InstrumentationOptions &CLOpts,
236 AllocTokenOptions Opts,
237 const Module &M) {
238 auto IntModuleFlagOrNull = [&](StringRef Key) {
239 return mdconst::extract_or_null<ConstantInt>(MD: M.getModuleFlag(Key));
240 };
241
242 if (auto *S = dyn_cast_or_null<MDString>(Val: M.getModuleFlag(Key: "alloc-token-mode")))
243 if (auto Mode = getAllocTokenModeFromString(Name: S->getString()))
244 Opts.Mode = *Mode;
245 if (auto *Val = IntModuleFlagOrNull("alloc-token-max"))
246 Opts.MaxTokens = Val->getZExtValue();
247 if (auto *Val = IntModuleFlagOrNull("alloc-token-fast-abi"))
248 Opts.FastABI |= Val->isOne();
249 if (auto *Val = IntModuleFlagOrNull("alloc-token-extended"))
250 Opts.Extended |= Val->isOne();
251
252 // Allow overriding options from command line options.
253 Opts.MaxTokens = CLOpts.alloc_token_max.value_or(u&: Opts.MaxTokens);
254 Opts.FastABI = valueOr(X: CLOpts.alloc_token_fast_abi, Default: Opts.FastABI);
255 Opts.Extended = valueOr(X: CLOpts.alloc_token_extended, Default: Opts.Extended);
256
257 return Opts;
258}
259
260class AllocToken {
261public:
262 AllocToken(const InstrumentationOptions &CLOpts, AllocTokenOptions Opts,
263 Module &M, ModuleAnalysisManager &MAM)
264 : CLOpts(CLOpts), Options(resolveOptions(CLOpts, Opts: std::move(Opts), M)),
265 Mod(M),
266 FAM(MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager()),
267 Mode(IncrementMode(*IntPtrTy, Options.MaxTokens)) {
268 switch (Options.Mode) {
269 case TokenMode::Increment:
270 break;
271 case TokenMode::Random:
272 Mode.emplace<RandomMode>(args&: *IntPtrTy, args: Options.MaxTokens,
273 args: M.createRNG(DEBUG_TYPE));
274 break;
275 case TokenMode::TypeHash:
276 Mode.emplace<TypeHashMode>(args&: *IntPtrTy, args: Options.MaxTokens,
277 args: CLOpts.alloc_token_fallback);
278 break;
279 case TokenMode::TypeHashPointerSplit:
280 Mode.emplace<TypeHashPointerSplitMode>(args&: *IntPtrTy, args: Options.MaxTokens,
281 args: CLOpts.alloc_token_fallback);
282 break;
283 case TokenMode::TypeFuncHash:
284 case TokenMode::TypeFuncHashPointerSplit:
285 Mode.emplace<TypeFuncHashMode>(args&: *IntPtrTy, args: Options.MaxTokens,
286 args: CLOpts.alloc_token_fallback, args: Options.Mode);
287 break;
288 }
289 }
290
291 bool instrumentFunction(Function &F);
292
293private:
294 /// Returns the LibFunc (or NotLibFunc) if this call should be instrumented.
295 std::optional<LibFunc>
296 shouldInstrumentCall(const CallBase &CB, const TargetLibraryInfo &TLI) const;
297
298 /// Returns true for functions that are eligible for instrumentation.
299 bool isInstrumentableLibFunc(LibFunc Func, const CallBase &CB,
300 const TargetLibraryInfo &TLI) const;
301
302 /// Returns true for isAllocationFn() functions that we should ignore.
303 static bool ignoreInstrumentableLibFunc(LibFunc Func);
304
305 /// Replace a call/invoke with a call/invoke to the allocation function
306 /// with token ID.
307 bool replaceAllocationCall(CallBase *CB, LibFunc Func,
308 OptimizationRemarkEmitter &ORE,
309 const TargetLibraryInfo &TLI);
310
311 /// Return replacement function for a LibFunc that takes a token ID.
312 FunctionCallee getTokenAllocFunction(const CallBase &CB, uint64_t TokenID,
313 LibFunc OriginalFunc);
314
315 /// Lower alloc_token_* intrinsics.
316 void replaceIntrinsicInst(IntrinsicInst *II, OptimizationRemarkEmitter &ORE);
317
318 /// Return the token ID from metadata in the call.
319 uint64_t getToken(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
320 return std::visit(visitor: [&](auto &&Mode) { return Mode(CB, ORE); }, variants&: Mode);
321 }
322
323 const InstrumentationOptions &CLOpts;
324 const AllocTokenOptions Options;
325 Module &Mod;
326 IntegerType *IntPtrTy = Mod.getDataLayout().getIntPtrType(C&: Mod.getContext());
327 FunctionAnalysisManager &FAM;
328 // Cache for replacement functions.
329 DenseMap<std::pair<LibFunc, uint64_t>, FunctionCallee> TokenAllocFunctions;
330 // Selected mode.
331 std::variant<IncrementMode, RandomMode, TypeHashMode,
332 TypeHashPointerSplitMode, TypeFuncHashMode>
333 Mode;
334};
335
336bool AllocToken::instrumentFunction(Function &F) {
337 // Do not apply any instrumentation for naked functions.
338 if (F.hasFnAttribute(Kind: Attribute::Naked))
339 return false;
340 // Don't touch available_externally functions, their actual body is elsewhere.
341 if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage)
342 return false;
343
344 SmallVector<std::pair<CallBase *, LibFunc>, 4> AllocCalls;
345 SmallVector<IntrinsicInst *, 4> IntrinsicInsts;
346
347 // Only instrument functions that have the sanitize_alloc_token attribute.
348 const bool InstrumentFunction =
349 F.hasFnAttribute(Kind: Attribute::SanitizeAllocToken) &&
350 !F.hasFnAttribute(Kind: Attribute::DisableSanitizerInstrumentation);
351
352 // Get TLI only when required.
353 const TargetLibraryInfo *TLI =
354 InstrumentFunction ? &FAM.getResult<TargetLibraryAnalysis>(IR&: F) : nullptr;
355
356 // Collect all allocation calls to avoid iterator invalidation.
357 for (Instruction &I : instructions(F)) {
358 // Collect all alloc_token_* intrinsics.
359 if (auto *II = dyn_cast<IntrinsicInst>(Val: &I);
360 II && II->getIntrinsicID() == Intrinsic::alloc_token_id) {
361 IntrinsicInsts.emplace_back(Args&: II);
362 continue;
363 }
364
365 if (!InstrumentFunction)
366 continue;
367
368 auto *CB = dyn_cast<CallBase>(Val: &I);
369 if (!CB)
370 continue;
371 if (std::optional<LibFunc> Func = shouldInstrumentCall(CB: *CB, TLI: *TLI))
372 AllocCalls.emplace_back(Args&: CB, Args&: Func.value());
373 }
374
375 // Return early to avoid unnecessarily instantiating the ORE.
376 if (AllocCalls.empty() && IntrinsicInsts.empty())
377 return false;
378
379 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(IR&: F);
380 bool Modified = false;
381
382 for (auto &[CB, Func] : AllocCalls)
383 Modified |= replaceAllocationCall(CB, Func, ORE, TLI: *TLI);
384
385 for (auto *II : IntrinsicInsts) {
386 replaceIntrinsicInst(II, ORE);
387 Modified = true;
388 }
389
390 if (Modified)
391 NumFunctionsModified++;
392
393 return Modified;
394}
395
396std::optional<LibFunc>
397AllocToken::shouldInstrumentCall(const CallBase &CB,
398 const TargetLibraryInfo &TLI) const {
399 const Function *Callee = CB.getCalledFunction();
400 if (!Callee)
401 return std::nullopt;
402
403 // Ignore nobuiltin of the CallBase, so that we can cover nobuiltin libcalls
404 // if requested via isInstrumentableLibFunc(). Note that isAllocationFn() is
405 // returning false for nobuiltin calls.
406 LibFunc Func = TLI.getLibFunc(FDecl: *Callee);
407 if (Func != NotLibFunc) {
408 if (isInstrumentableLibFunc(Func, CB, TLI))
409 return Func;
410 } else if (Options.Extended && CB.getMetadata(KindID: LLVMContext::MD_alloc_token)) {
411 return NotLibFunc;
412 }
413
414 return std::nullopt;
415}
416
417bool AllocToken::isInstrumentableLibFunc(LibFunc Func, const CallBase &CB,
418 const TargetLibraryInfo &TLI) const {
419 if (ignoreInstrumentableLibFunc(Func))
420 return false;
421
422 if (isAllocationFn(V: &CB, TLI: &TLI))
423 return true;
424
425 switch (Func) {
426 // These libfuncs don't return normal pointers, and are therefore not handled
427 // by isAllocationFn().
428 case LibFunc_posix_memalign:
429 case LibFunc_size_returning_new:
430 case LibFunc_size_returning_new_hot_cold:
431 case LibFunc_size_returning_new_aligned:
432 case LibFunc_size_returning_new_aligned_hot_cold:
433 return true;
434
435 // See the comment on -alloc-token-cover-replaceable-new.
436 case LibFunc_Znwj:
437 case LibFunc_ZnwjRKSt9nothrow_t:
438 case LibFunc_ZnwjSt11align_val_t:
439 case LibFunc_ZnwjSt11align_val_tRKSt9nothrow_t:
440 case LibFunc_Znwm:
441 case LibFunc_Znwm12__hot_cold_t:
442 case LibFunc_ZnwmRKSt9nothrow_t:
443 case LibFunc_ZnwmRKSt9nothrow_t12__hot_cold_t:
444 case LibFunc_ZnwmSt11align_val_t:
445 case LibFunc_ZnwmSt11align_val_t12__hot_cold_t:
446 case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t:
447 case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t:
448 case LibFunc_Znaj:
449 case LibFunc_ZnajRKSt9nothrow_t:
450 case LibFunc_ZnajSt11align_val_t:
451 case LibFunc_ZnajSt11align_val_tRKSt9nothrow_t:
452 case LibFunc_Znam:
453 case LibFunc_Znam12__hot_cold_t:
454 case LibFunc_ZnamRKSt9nothrow_t:
455 case LibFunc_ZnamRKSt9nothrow_t12__hot_cold_t:
456 case LibFunc_ZnamSt11align_val_t:
457 case LibFunc_ZnamSt11align_val_t12__hot_cold_t:
458 case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t:
459 case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t:
460 case LibFunc_msvc_new_int:
461 case LibFunc_msvc_new_int_nothrow:
462 case LibFunc_msvc_new_longlong:
463 case LibFunc_msvc_new_longlong_nothrow:
464 case LibFunc_msvc_new_array_int:
465 case LibFunc_msvc_new_array_int_nothrow:
466 case LibFunc_msvc_new_array_longlong:
467 case LibFunc_msvc_new_array_longlong_nothrow:
468 return CLOpts.alloc_token_cover_replaceable_new;
469
470 default:
471 return false;
472 }
473}
474
475bool AllocToken::ignoreInstrumentableLibFunc(LibFunc Func) {
476 switch (Func) {
477 case LibFunc_strdup:
478 case LibFunc_dunder_strdup:
479 case LibFunc_strndup:
480 case LibFunc_dunder_strndup:
481 return true;
482 default:
483 return false;
484 }
485}
486
487bool AllocToken::replaceAllocationCall(CallBase *CB, LibFunc Func,
488 OptimizationRemarkEmitter &ORE,
489 const TargetLibraryInfo &TLI) {
490 uint64_t TokenID = getToken(CB: *CB, ORE);
491
492 FunctionCallee TokenAlloc = getTokenAllocFunction(CB: *CB, TokenID, OriginalFunc: Func);
493 if (!TokenAlloc)
494 return false;
495 NumAllocationsInstrumented++;
496
497 if (Options.FastABI) {
498 assert(TokenAlloc.getFunctionType()->getNumParams() == CB->arg_size());
499 CB->setCalledFunction(TokenAlloc);
500 return true;
501 }
502
503 IRBuilder<> IRB(CB);
504 // Original args.
505 SmallVector<Value *, 4> NewArgs{CB->args()};
506 // Add token ID, truncated to IntPtrTy width.
507 NewArgs.push_back(Elt: ConstantInt::get(Ty: IntPtrTy, V: TokenID));
508 assert(TokenAlloc.getFunctionType()->getNumParams() == NewArgs.size());
509
510 // Preserve invoke vs call semantics for exception handling.
511 CallBase *NewCall;
512 if (auto *II = dyn_cast<InvokeInst>(Val: CB)) {
513 NewCall = IRB.CreateInvoke(Callee: TokenAlloc, NormalDest: II->getNormalDest(),
514 UnwindDest: II->getUnwindDest(), Args: NewArgs);
515 } else {
516 NewCall = IRB.CreateCall(Callee: TokenAlloc, Args: NewArgs);
517 cast<CallInst>(Val: NewCall)->setTailCall(CB->isTailCall());
518 }
519 NewCall->setCallingConv(CB->getCallingConv());
520 NewCall->copyMetadata(SrcInst: *CB);
521 NewCall->setAttributes(CB->getAttributes());
522
523 // Replace all uses and delete the old call.
524 CB->replaceAllUsesWith(V: NewCall);
525 CB->eraseFromParent();
526 return true;
527}
528
529FunctionCallee AllocToken::getTokenAllocFunction(const CallBase &CB,
530 uint64_t TokenID,
531 LibFunc OriginalFunc) {
532 std::optional<std::pair<LibFunc, uint64_t>> Key;
533 if (OriginalFunc != NotLibFunc) {
534 Key = std::make_pair(x&: OriginalFunc, y: Options.FastABI ? TokenID : 0);
535 auto It = TokenAllocFunctions.find(Val: *Key);
536 if (It != TokenAllocFunctions.end())
537 return It->second;
538 }
539
540 const Function *Callee = CB.getCalledFunction();
541 if (!Callee)
542 return FunctionCallee();
543 const FunctionType *OldFTy = Callee->getFunctionType();
544 if (OldFTy->isVarArg())
545 return FunctionCallee();
546 // Copy params, and append token ID type.
547 Type *RetTy = OldFTy->getReturnType();
548 SmallVector<Type *, 4> NewParams{OldFTy->params()};
549 std::string TokenAllocName = CLOpts.alloc_token_prefix.str();
550 if (Options.FastABI)
551 TokenAllocName += utostr(X: TokenID) + "_";
552 else
553 NewParams.push_back(Elt: IntPtrTy); // token ID
554 TokenAllocName += Callee->getName();
555 FunctionType *NewFTy = FunctionType::get(Result: RetTy, Params: NewParams, isVarArg: false);
556 AttributeList NewAttrs = Callee->getAttributes();
557 FunctionCallee TokenAlloc =
558 Mod.getOrInsertFunction(Name: TokenAllocName, T: NewFTy, AttributeList: NewAttrs);
559
560 if (Key.has_value())
561 TokenAllocFunctions[*Key] = TokenAlloc;
562 return TokenAlloc;
563}
564
565void AllocToken::replaceIntrinsicInst(IntrinsicInst *II,
566 OptimizationRemarkEmitter &ORE) {
567 assert(II->getIntrinsicID() == Intrinsic::alloc_token_id);
568
569 uint64_t TokenID = getToken(CB: *II, ORE);
570 Value *V = ConstantInt::get(Ty: IntPtrTy, V: TokenID);
571 II->replaceAllUsesWith(V);
572 II->eraseFromParent();
573}
574
575} // namespace
576
577AllocTokenPass::AllocTokenPass(AllocTokenOptions Opts)
578 : Options(std::move(Opts)) {}
579
580PreservedAnalyses AllocTokenPass::run(Module &M, ModuleAnalysisManager &MAM) {
581 const InstrumentationOptions &CLOpts =
582 M.getContext().getLibraryOptions().get<InstrumentationOptions>();
583 AllocToken Pass(CLOpts, Options, M, MAM);
584 bool Modified = false;
585
586 for (Function &F : M) {
587 if (F.empty())
588 continue; // declaration
589 Modified |= Pass.instrumentFunction(F);
590 }
591
592 return Modified ? PreservedAnalyses::none().preserveSet<CFGAnalyses>()
593 : PreservedAnalyses::all();
594}
595