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