@@ -67,24 +67,24 @@ cl::opt<std::string> ClFuncPrefix("alloc-token-prefix",
6767 cl::desc (" The allocation function prefix"
6868 cl::Hidden, cl::init(" __alloc_token_"
6969
70- cl::opt<uint64_t >
70+ cl::opt<std::optional< uint64_t >, false , cl::parser< uint64_t > >
7171 ClMaxTokens (" alloc-token-max"
7272 cl::desc (" Maximum number of tokens (0 = target SIZE_MAX)"
73- cl::Hidden, cl::init(0 ));
73+ cl::Hidden, cl::init(std:: nullopt ));
7474
75- cl::opt<bool >
75+ cl::opt<std::optional< bool >, false , cl::parser< bool > >
7676 ClFastABI (" alloc-token-fast-abi"
7777 cl::desc (" The token ID is encoded in the function name"
78- cl::Hidden, cl::init(false ));
78+ cl::Hidden, cl::init(std:: nullopt ));
7979
8080// Instrument libcalls only by default - compatible allocators only need to take
8181// care of providing standard allocation functions. With extended coverage, also
8282// instrument non-libcall allocation function calls with !alloc_token
8383// metadata.
84- cl::opt<bool >
84+ cl::opt<std::optional< bool >, false , cl::parser< bool > >
8585 ClExtended (" alloc-token-extended"
8686 cl::desc (" Extend coverage to custom allocation functions"
87- cl::Hidden, cl::init(false ));
87+ cl::Hidden, cl::init(std:: nullopt ));
8888
8989// C++ defines ::operator new (and variants) as replaceable (vs. standard
9090// library versions), which are nobuiltin, and are therefore not covered by
@@ -237,28 +237,28 @@ class TypeHashPointerSplitMode : public TypeHashMode {
237237// Apply opt overrides and module flags.
238238static AllocTokenOptions resolveOptions (AllocTokenOptions Opts,
239239 const Module &M) {
240- if (!Opts.MaxTokens .has_value ())
241- Opts.MaxTokens = ClMaxTokens;
242- Opts.FastABI |= ClFastABI;
243- Opts.Extended |= ClExtended;
244-
245240 auto IntModuleFlagOrNull = [&](StringRef Key) {
246241 return mdconst::extract_or_null<ConstantInt>(M.getModuleFlag (Key));
247242 };
248243
249244 if (auto *S = dyn_cast_or_null<MDString>(M.getModuleFlag (" alloc-token-mode"
250245 if (auto Mode = getAllocTokenModeFromString (S->getString ()))
251246 Opts.Mode = *Mode;
252-
253247 if (auto *Val = IntModuleFlagOrNull (" alloc-token-max"
254248 Opts.MaxTokens = Val->getZExtValue ();
255-
256249 if (auto *Val = IntModuleFlagOrNull (" alloc-token-fast-abi"
257250 Opts.FastABI |= Val->isOne ();
258-
259251 if (auto *Val = IntModuleFlagOrNull (" alloc-token-extended"
260252 Opts.Extended |= Val->isOne ();
261253
254+ // Allow overriding options from command line options.
255+ if (ClMaxTokens.has_value ())
256+ Opts.MaxTokens = *ClMaxTokens;
257+ if (ClFastABI.has_value ())
258+ Opts.FastABI = *ClFastABI;
259+ if (ClExtended.has_value ())
260+ Opts.Extended = *ClExtended;
261+
262262 return Opts;
263263}
264264
@@ -268,19 +268,19 @@ class AllocToken {
268268 ModuleAnalysisManager &MAM)
269269 : Options(resolveOptions(std::move(Opts), M)), Mod(M),
270270 FAM(MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
271- Mode(IncrementMode(*IntPtrTy, * Options.MaxTokens)) {
271+ Mode(IncrementMode(*IntPtrTy, Options.MaxTokens)) {
272272 switch (Options.Mode ) {
273273 case TokenMode::Increment:
274274 break ;
275275 case TokenMode::Random:
276- Mode.emplace <RandomMode>(*IntPtrTy, * Options.MaxTokens ,
276+ Mode.emplace <RandomMode>(*IntPtrTy, Options.MaxTokens ,
277277 M.createRNG (DEBUG_TYPE));
278278 break ;
279279 case TokenMode::TypeHash:
280- Mode.emplace <TypeHashMode>(*IntPtrTy, * Options.MaxTokens );
280+ Mode.emplace <TypeHashMode>(*IntPtrTy, Options.MaxTokens );
281281 break ;
282282 case TokenMode::TypeHashPointerSplit:
283- Mode.emplace <TypeHashPointerSplitMode>(*IntPtrTy, * Options.MaxTokens );
283+ Mode.emplace <TypeHashPointerSplitMode>(*IntPtrTy, Options.MaxTokens );
284284 break ;
285285 }
286286 }
@@ -337,8 +337,6 @@ bool AllocToken::instrumentFunction(Function &F) {
337337 if (F.getLinkage () == GlobalValue::AvailableExternallyLinkage)
338338 return false ;
339339
340- auto &ORE = FAM.getResult <OptimizationRemarkEmitterAnalysis>(F);
341- auto &TLI = FAM.getResult <TargetLibraryAnalysis>(F);
342340 SmallVector<std::pair<CallBase *, LibFunc>, 4 > AllocCalls;
343341 SmallVector<IntrinsicInst *, 4 > IntrinsicInsts;
344342
@@ -347,6 +345,10 @@ bool AllocToken::instrumentFunction(Function &F) {
347345 F.hasFnAttribute (Attribute::SanitizeAllocToken) &&
348346 !F.hasFnAttribute (Attribute::DisableSanitizerInstrumentation);
349347
348+ // Get TLI only when required.
349+ const TargetLibraryInfo *TLI =
350+ InstrumentFunction ? &FAM.getResult <TargetLibraryAnalysis>(F) : nullptr ;
351+
350352 // Collect all allocation calls to avoid iterator invalidation.
351353 for (Instruction &I : instructions (F)) {
352354 // Collect all alloc_token_* intrinsics.
@@ -362,26 +364,28 @@ bool AllocToken::instrumentFunction(Function &F) {
362364 auto *CB = dyn_cast<CallBase>(&I);
363365 if (!CB)
364366 continue ;
365- if (std::optional<LibFunc> Func = shouldInstrumentCall (*CB, TLI))
367+ if (std::optional<LibFunc> Func = shouldInstrumentCall (*CB, * TLI))
366368 AllocCalls.emplace_back (CB, Func.value ());
367369 }
368370
371+ // Return early to avoid unnecessarily instantiating the ORE.
372+ if (AllocCalls.empty () && IntrinsicInsts.empty ())
373+ return false ;
374+
375+ auto &ORE = FAM.getResult <OptimizationRemarkEmitterAnalysis>(F);
369376 bool Modified = false ;
370377
371- if (!AllocCalls.empty ()) {
372- for (auto &[CB, Func] : AllocCalls)
373- Modified |= replaceAllocationCall (CB, Func, ORE, TLI);
374- if (Modified)
375- NumFunctionsModified++;
376- }
378+ for (auto &[CB, Func] : AllocCalls)
379+ Modified |= replaceAllocationCall (CB, Func, ORE, *TLI);
377380
378- if (!IntrinsicInsts.empty ()) {
379- for (auto *II : IntrinsicInsts)
380- replaceIntrinsicInst (II, ORE);
381+ for (auto *II : IntrinsicInsts) {
382+ replaceIntrinsicInst (II, ORE);
381383 Modified = true ;
382- NumFunctionsModified++;
383384 }
384385
386+ if (Modified)
387+ NumFunctionsModified++;
388+
385389 return Modified;
386390}
387391
0 commit comments