1 //===- StackSafetyAnalysis.cpp - Stack memory safety analysis -------------===//
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 //===----------------------------------------------------------------------===//
10 
11 #include "llvm/Analysis/StackSafetyAnalysis.h"
12 #include "llvm/ADT/APInt.h"
13 #include "llvm/ADT/SmallPtrSet.h"
14 #include "llvm/ADT/SmallVector.h"
15 #include "llvm/ADT/Statistic.h"
16 #include "llvm/Analysis/ModuleSummaryAnalysis.h"
17 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
18 #include "llvm/Analysis/StackLifetime.h"
19 #include "llvm/IR/ConstantRange.h"
20 #include "llvm/IR/DerivedTypes.h"
21 #include "llvm/IR/GlobalValue.h"
22 #include "llvm/IR/InstIterator.h"
23 #include "llvm/IR/Instructions.h"
24 #include "llvm/IR/IntrinsicInst.h"
25 #include "llvm/IR/ModuleSummaryIndex.h"
26 #include "llvm/InitializePasses.h"
27 #include "llvm/Support/Casting.h"
28 #include "llvm/Support/CommandLine.h"
29 #include "llvm/Support/FormatVariadic.h"
30 #include "llvm/Support/raw_ostream.h"
31 #include <algorithm>
32 #include <memory>
33 
34 using namespace llvm;
35 
36 #define DEBUG_TYPE "stack-safety"
37 
38 STATISTIC(NumAllocaStackSafe, "Number of safe allocas");
39 STATISTIC(NumAllocaTotal, "Number of total allocas");
40 
41 STATISTIC(NumCombinedCalleeLookupTotal,
42           "Number of total callee lookups on combined index.");
43 STATISTIC(NumCombinedCalleeLookupFailed,
44           "Number of failed callee lookups on combined index.");
45 STATISTIC(NumModuleCalleeLookupTotal,
46           "Number of total callee lookups on module index.");
47 STATISTIC(NumModuleCalleeLookupFailed,
48           "Number of failed callee lookups on module index.");
49 STATISTIC(NumCombinedParamAccessesBefore,
50           "Number of total param accesses before generateParamAccessSummary.");
51 STATISTIC(NumCombinedParamAccessesAfter,
52           "Number of total param accesses after generateParamAccessSummary.");
53 STATISTIC(NumCombinedDataFlowNodes,
54           "Number of total nodes in combined index for dataflow processing.");
55 STATISTIC(NumIndexCalleeUnhandled, "Number of index callee which are unhandled.");
56 STATISTIC(NumIndexCalleeMultipleWeak, "Number of index callee non-unique weak.");
57 STATISTIC(NumIndexCalleeMultipleExternal, "Number of index callee non-unique external.");
58 
59 
60 static cl::opt<int> StackSafetyMaxIterations("stack-safety-max-iterations",
61                                              cl::init(20), cl::Hidden);
62 
63 static cl::opt<bool> StackSafetyPrint("stack-safety-print", cl::init(false),
64                                       cl::Hidden);
65 
66 static cl::opt<bool> StackSafetyRun("stack-safety-run", cl::init(false),
67                                     cl::Hidden);
68 
69 namespace {
70 
71 // Check if we should bailout for such ranges.
isUnsafe(const ConstantRange & R)72 bool isUnsafe(const ConstantRange &R) {
73   return R.isEmptySet() || R.isFullSet() || R.isUpperSignWrapped();
74 }
75 
addOverflowNever(const ConstantRange & L,const ConstantRange & R)76 ConstantRange addOverflowNever(const ConstantRange &L, const ConstantRange &R) {
77   assert(!L.isSignWrappedSet());
78   assert(!R.isSignWrappedSet());
79   if (L.signedAddMayOverflow(R) !=
80       ConstantRange::OverflowResult::NeverOverflows)
81     return ConstantRange::getFull(L.getBitWidth());
82   ConstantRange Result = L.add(R);
83   assert(!Result.isSignWrappedSet());
84   return Result;
85 }
86 
unionNoWrap(const ConstantRange & L,const ConstantRange & R)87 ConstantRange unionNoWrap(const ConstantRange &L, const ConstantRange &R) {
88   assert(!L.isSignWrappedSet());
89   assert(!R.isSignWrappedSet());
90   auto Result = L.unionWith(R);
91   // Two non-wrapped sets can produce wrapped.
92   if (Result.isSignWrappedSet())
93     Result = ConstantRange::getFull(Result.getBitWidth());
94   return Result;
95 }
96 
97 /// Describes use of address in as a function call argument.
98 template <typename CalleeTy> struct CallInfo {
99   /// Function being called.
100   const CalleeTy *Callee = nullptr;
101   /// Index of argument which pass address.
102   size_t ParamNo = 0;
103 
CallInfo__anon9ddec65d0111::CallInfo104   CallInfo(const CalleeTy *Callee, size_t ParamNo)
105       : Callee(Callee), ParamNo(ParamNo) {}
106 
107   struct Less {
operator ()__anon9ddec65d0111::CallInfo::Less108     bool operator()(const CallInfo &L, const CallInfo &R) const {
109       return std::tie(L.ParamNo, L.Callee) < std::tie(R.ParamNo, R.Callee);
110     }
111   };
112 };
113 
114 /// Describe uses of address (alloca or parameter) inside of the function.
115 template <typename CalleeTy> struct UseInfo {
116   // Access range if the address (alloca or parameters).
117   // It is allowed to be empty-set when there are no known accesses.
118   ConstantRange Range;
119 
120   // List of calls which pass address as an argument.
121   // Value is offset range of address from base address (alloca or calling
122   // function argument). Range should never set to empty-set, that is an invalid
123   // access range that can cause empty-set to be propagated with
124   // ConstantRange::add
125   using CallsTy = std::map<CallInfo<CalleeTy>, ConstantRange,
126                            typename CallInfo<CalleeTy>::Less>;
127   CallsTy Calls;
128 
UseInfo__anon9ddec65d0111::UseInfo129   UseInfo(unsigned PointerSize) : Range{PointerSize, false} {}
130 
updateRange__anon9ddec65d0111::UseInfo131   void updateRange(const ConstantRange &R) { Range = unionNoWrap(Range, R); }
132 };
133 
134 template <typename CalleeTy>
operator <<(raw_ostream & OS,const UseInfo<CalleeTy> & U)135 raw_ostream &operator<<(raw_ostream &OS, const UseInfo<CalleeTy> &U) {
136   OS << U.Range;
137   for (auto &Call : U.Calls)
138     OS << ", "
139        << "@" << Call.first.Callee->getName() << "(arg" << Call.first.ParamNo
140        << ", " << Call.second << ")";
141   return OS;
142 }
143 
144 /// Calculate the allocation size of a given alloca. Returns empty range
145 // in case of confution.
getStaticAllocaSizeRange(const AllocaInst & AI)146 ConstantRange getStaticAllocaSizeRange(const AllocaInst &AI) {
147   const DataLayout &DL = AI.getModule()->getDataLayout();
148   TypeSize TS = DL.getTypeAllocSize(AI.getAllocatedType());
149   unsigned PointerSize = DL.getMaxPointerSizeInBits();
150   // Fallback to empty range for alloca size.
151   ConstantRange R = ConstantRange::getEmpty(PointerSize);
152   if (TS.isScalable())
153     return R;
154   APInt APSize(PointerSize, TS.getFixedSize(), true);
155   if (APSize.isNonPositive())
156     return R;
157   if (AI.isArrayAllocation()) {
158     const auto *C = dyn_cast<ConstantInt>(AI.getArraySize());
159     if (!C)
160       return R;
161     bool Overflow = false;
162     APInt Mul = C->getValue();
163     if (Mul.isNonPositive())
164       return R;
165     Mul = Mul.sextOrTrunc(PointerSize);
166     APSize = APSize.smul_ov(Mul, Overflow);
167     if (Overflow)
168       return R;
169   }
170   R = ConstantRange(APInt::getNullValue(PointerSize), APSize);
171   assert(!isUnsafe(R));
172   return R;
173 }
174 
175 template <typename CalleeTy> struct FunctionInfo {
176   std::map<const AllocaInst *, UseInfo<CalleeTy>> Allocas;
177   std::map<uint32_t, UseInfo<CalleeTy>> Params;
178   // TODO: describe return value as depending on one or more of its arguments.
179 
180   // StackSafetyDataFlowAnalysis counter stored here for faster access.
181   int UpdateCount = 0;
182 
print__anon9ddec65d0111::FunctionInfo183   void print(raw_ostream &O, StringRef Name, const Function *F) const {
184     // TODO: Consider different printout format after
185     // StackSafetyDataFlowAnalysis. Calls and parameters are irrelevant then.
186     O << "  @" << Name << ((F && F->isDSOLocal()) ? "" : " dso_preemptable")
187       << ((F && F->isInterposable()) ? " interposable" : "") << "\n";
188 
189     O << "    args uses:\n";
190     for (auto &KV : Params) {
191       O << "      ";
192       if (F)
193         O << F->getArg(KV.first)->getName();
194       else
195         O << formatv("arg{0}", KV.first);
196       O << "[]: " << KV.second << "\n";
197     }
198 
199     O << "    allocas uses:\n";
200     if (F) {
201       for (auto &I : instructions(F)) {
202         if (const AllocaInst *AI = dyn_cast<AllocaInst>(&I)) {
203           auto &AS = Allocas.find(AI)->second;
204           O << "      " << AI->getName() << "["
205             << getStaticAllocaSizeRange(*AI).getUpper() << "]: " << AS << "\n";
206         }
207       }
208     } else {
209       assert(Allocas.empty());
210     }
211     O << "\n";
212   }
213 };
214 
215 using GVToSSI = std::map<const GlobalValue *, FunctionInfo<GlobalValue>>;
216 
217 } // namespace
218 
219 struct StackSafetyInfo::InfoTy {
220   FunctionInfo<GlobalValue> Info;
221 };
222 
223 struct StackSafetyGlobalInfo::InfoTy {
224   GVToSSI Info;
225   SmallPtrSet<const AllocaInst *, 8> SafeAllocas;
226 };
227 
228 namespace {
229 
230 class StackSafetyLocalAnalysis {
231   Function &F;
232   const DataLayout &DL;
233   ScalarEvolution &SE;
234   unsigned PointerSize = 0;
235 
236   const ConstantRange UnknownRange;
237 
238   ConstantRange offsetFrom(Value *Addr, Value *Base);
239   ConstantRange getAccessRange(Value *Addr, Value *Base,
240                                const ConstantRange &SizeRange);
241   ConstantRange getAccessRange(Value *Addr, Value *Base, TypeSize Size);
242   ConstantRange getMemIntrinsicAccessRange(const MemIntrinsic *MI, const Use &U,
243                                            Value *Base);
244 
245   bool analyzeAllUses(Value *Ptr, UseInfo<GlobalValue> &AS,
246                       const StackLifetime &SL);
247 
248 public:
StackSafetyLocalAnalysis(Function & F,ScalarEvolution & SE)249   StackSafetyLocalAnalysis(Function &F, ScalarEvolution &SE)
250       : F(F), DL(F.getParent()->getDataLayout()), SE(SE),
251         PointerSize(DL.getPointerSizeInBits()),
252         UnknownRange(PointerSize, true) {}
253 
254   // Run the transformation on the associated function.
255   FunctionInfo<GlobalValue> run();
256 };
257 
offsetFrom(Value * Addr,Value * Base)258 ConstantRange StackSafetyLocalAnalysis::offsetFrom(Value *Addr, Value *Base) {
259   if (!SE.isSCEVable(Addr->getType()) || !SE.isSCEVable(Base->getType()))
260     return UnknownRange;
261 
262   auto *PtrTy = IntegerType::getInt8PtrTy(SE.getContext());
263   const SCEV *AddrExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Addr), PtrTy);
264   const SCEV *BaseExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Base), PtrTy);
265   const SCEV *Diff = SE.getMinusSCEV(AddrExp, BaseExp);
266   if (isa<SCEVCouldNotCompute>(Diff))
267     return UnknownRange;
268 
269   ConstantRange Offset = SE.getSignedRange(Diff);
270   if (isUnsafe(Offset))
271     return UnknownRange;
272   return Offset.sextOrTrunc(PointerSize);
273 }
274 
275 ConstantRange
getAccessRange(Value * Addr,Value * Base,const ConstantRange & SizeRange)276 StackSafetyLocalAnalysis::getAccessRange(Value *Addr, Value *Base,
277                                          const ConstantRange &SizeRange) {
278   // Zero-size loads and stores do not access memory.
279   if (SizeRange.isEmptySet())
280     return ConstantRange::getEmpty(PointerSize);
281   assert(!isUnsafe(SizeRange));
282 
283   ConstantRange Offsets = offsetFrom(Addr, Base);
284   if (isUnsafe(Offsets))
285     return UnknownRange;
286 
287   Offsets = addOverflowNever(Offsets, SizeRange);
288   if (isUnsafe(Offsets))
289     return UnknownRange;
290   return Offsets;
291 }
292 
getAccessRange(Value * Addr,Value * Base,TypeSize Size)293 ConstantRange StackSafetyLocalAnalysis::getAccessRange(Value *Addr, Value *Base,
294                                                        TypeSize Size) {
295   if (Size.isScalable())
296     return UnknownRange;
297   APInt APSize(PointerSize, Size.getFixedSize(), true);
298   if (APSize.isNegative())
299     return UnknownRange;
300   return getAccessRange(
301       Addr, Base, ConstantRange(APInt::getNullValue(PointerSize), APSize));
302 }
303 
getMemIntrinsicAccessRange(const MemIntrinsic * MI,const Use & U,Value * Base)304 ConstantRange StackSafetyLocalAnalysis::getMemIntrinsicAccessRange(
305     const MemIntrinsic *MI, const Use &U, Value *Base) {
306   if (const auto *MTI = dyn_cast<MemTransferInst>(MI)) {
307     if (MTI->getRawSource() != U && MTI->getRawDest() != U)
308       return ConstantRange::getEmpty(PointerSize);
309   } else {
310     if (MI->getRawDest() != U)
311       return ConstantRange::getEmpty(PointerSize);
312   }
313 
314   auto *CalculationTy = IntegerType::getIntNTy(SE.getContext(), PointerSize);
315   if (!SE.isSCEVable(MI->getLength()->getType()))
316     return UnknownRange;
317 
318   const SCEV *Expr =
319       SE.getTruncateOrZeroExtend(SE.getSCEV(MI->getLength()), CalculationTy);
320   ConstantRange Sizes = SE.getSignedRange(Expr);
321   if (Sizes.getUpper().isNegative() || isUnsafe(Sizes))
322     return UnknownRange;
323   Sizes = Sizes.sextOrTrunc(PointerSize);
324   ConstantRange SizeRange(APInt::getNullValue(PointerSize),
325                           Sizes.getUpper() - 1);
326   return getAccessRange(U, Base, SizeRange);
327 }
328 
329 /// The function analyzes all local uses of Ptr (alloca or argument) and
330 /// calculates local access range and all function calls where it was used.
analyzeAllUses(Value * Ptr,UseInfo<GlobalValue> & US,const StackLifetime & SL)331 bool StackSafetyLocalAnalysis::analyzeAllUses(Value *Ptr,
332                                               UseInfo<GlobalValue> &US,
333                                               const StackLifetime &SL) {
334   SmallPtrSet<const Value *, 16> Visited;
335   SmallVector<const Value *, 8> WorkList;
336   WorkList.push_back(Ptr);
337   const AllocaInst *AI = dyn_cast<AllocaInst>(Ptr);
338 
339   // A DFS search through all uses of the alloca in bitcasts/PHI/GEPs/etc.
340   while (!WorkList.empty()) {
341     const Value *V = WorkList.pop_back_val();
342     for (const Use &UI : V->uses()) {
343       const auto *I = cast<Instruction>(UI.getUser());
344       if (!SL.isReachable(I))
345         continue;
346 
347       assert(V == UI.get());
348 
349       switch (I->getOpcode()) {
350       case Instruction::Load: {
351         if (AI && !SL.isAliveAfter(AI, I)) {
352           US.updateRange(UnknownRange);
353           return false;
354         }
355         US.updateRange(
356             getAccessRange(UI, Ptr, DL.getTypeStoreSize(I->getType())));
357         break;
358       }
359 
360       case Instruction::VAArg:
361         // "va-arg" from a pointer is safe.
362         break;
363       case Instruction::Store: {
364         if (V == I->getOperand(0)) {
365           // Stored the pointer - conservatively assume it may be unsafe.
366           US.updateRange(UnknownRange);
367           return false;
368         }
369         if (AI && !SL.isAliveAfter(AI, I)) {
370           US.updateRange(UnknownRange);
371           return false;
372         }
373         US.updateRange(getAccessRange(
374             UI, Ptr, DL.getTypeStoreSize(I->getOperand(0)->getType())));
375         break;
376       }
377 
378       case Instruction::Ret:
379         // Information leak.
380         // FIXME: Process parameters correctly. This is a leak only if we return
381         // alloca.
382         US.updateRange(UnknownRange);
383         return false;
384 
385       case Instruction::Call:
386       case Instruction::Invoke: {
387         if (I->isLifetimeStartOrEnd())
388           break;
389 
390         if (AI && !SL.isAliveAfter(AI, I)) {
391           US.updateRange(UnknownRange);
392           return false;
393         }
394 
395         if (const MemIntrinsic *MI = dyn_cast<MemIntrinsic>(I)) {
396           US.updateRange(getMemIntrinsicAccessRange(MI, UI, Ptr));
397           break;
398         }
399 
400         const auto &CB = cast<CallBase>(*I);
401         if (!CB.isArgOperand(&UI)) {
402           US.updateRange(UnknownRange);
403           return false;
404         }
405 
406         unsigned ArgNo = CB.getArgOperandNo(&UI);
407         if (CB.isByValArgument(ArgNo)) {
408           US.updateRange(getAccessRange(
409               UI, Ptr, DL.getTypeStoreSize(CB.getParamByValType(ArgNo))));
410           break;
411         }
412 
413         // FIXME: consult devirt?
414         // Do not follow aliases, otherwise we could inadvertently follow
415         // dso_preemptable aliases or aliases with interposable linkage.
416         const GlobalValue *Callee =
417             dyn_cast<GlobalValue>(CB.getCalledOperand()->stripPointerCasts());
418         if (!Callee) {
419           US.updateRange(UnknownRange);
420           return false;
421         }
422 
423         assert(isa<Function>(Callee) || isa<GlobalAlias>(Callee));
424         ConstantRange Offsets = offsetFrom(UI, Ptr);
425         auto Insert =
426             US.Calls.emplace(CallInfo<GlobalValue>(Callee, ArgNo), Offsets);
427         if (!Insert.second)
428           Insert.first->second = Insert.first->second.unionWith(Offsets);
429         break;
430       }
431 
432       default:
433         if (Visited.insert(I).second)
434           WorkList.push_back(cast<const Instruction>(I));
435       }
436     }
437   }
438 
439   return true;
440 }
441 
run()442 FunctionInfo<GlobalValue> StackSafetyLocalAnalysis::run() {
443   FunctionInfo<GlobalValue> Info;
444   assert(!F.isDeclaration() &&
445          "Can't run StackSafety on a function declaration");
446 
447   LLVM_DEBUG(dbgs() << "[StackSafety] " << F.getName() << "\n");
448 
449   SmallVector<AllocaInst *, 64> Allocas;
450   for (auto &I : instructions(F))
451     if (auto *AI = dyn_cast<AllocaInst>(&I))
452       Allocas.push_back(AI);
453   StackLifetime SL(F, Allocas, StackLifetime::LivenessType::Must);
454   SL.run();
455 
456   for (auto *AI : Allocas) {
457     auto &UI = Info.Allocas.emplace(AI, PointerSize).first->second;
458     analyzeAllUses(AI, UI, SL);
459   }
460 
461   for (Argument &A : F.args()) {
462     // Non pointers and bypass arguments are not going to be used in any global
463     // processing.
464     if (A.getType()->isPointerTy() && !A.hasByValAttr()) {
465       auto &UI = Info.Params.emplace(A.getArgNo(), PointerSize).first->second;
466       analyzeAllUses(&A, UI, SL);
467     }
468   }
469 
470   LLVM_DEBUG(Info.print(dbgs(), F.getName(), &F));
471   LLVM_DEBUG(dbgs() << "[StackSafety] done\n");
472   return Info;
473 }
474 
475 template <typename CalleeTy> class StackSafetyDataFlowAnalysis {
476   using FunctionMap = std::map<const CalleeTy *, FunctionInfo<CalleeTy>>;
477 
478   FunctionMap Functions;
479   const ConstantRange UnknownRange;
480 
481   // Callee-to-Caller multimap.
482   DenseMap<const CalleeTy *, SmallVector<const CalleeTy *, 4>> Callers;
483   SetVector<const CalleeTy *> WorkList;
484 
485   bool updateOneUse(UseInfo<CalleeTy> &US, bool UpdateToFullSet);
486   void updateOneNode(const CalleeTy *Callee, FunctionInfo<CalleeTy> &FS);
updateOneNode(const CalleeTy * Callee)487   void updateOneNode(const CalleeTy *Callee) {
488     updateOneNode(Callee, Functions.find(Callee)->second);
489   }
updateAllNodes()490   void updateAllNodes() {
491     for (auto &F : Functions)
492       updateOneNode(F.first, F.second);
493   }
494   void runDataFlow();
495 #ifndef NDEBUG
496   void verifyFixedPoint();
497 #endif
498 
499 public:
StackSafetyDataFlowAnalysis(uint32_t PointerBitWidth,FunctionMap Functions)500   StackSafetyDataFlowAnalysis(uint32_t PointerBitWidth, FunctionMap Functions)
501       : Functions(std::move(Functions)),
502         UnknownRange(ConstantRange::getFull(PointerBitWidth)) {}
503 
504   const FunctionMap &run();
505 
506   ConstantRange getArgumentAccessRange(const CalleeTy *Callee, unsigned ParamNo,
507                                        const ConstantRange &Offsets) const;
508 };
509 
510 template <typename CalleeTy>
getArgumentAccessRange(const CalleeTy * Callee,unsigned ParamNo,const ConstantRange & Offsets) const511 ConstantRange StackSafetyDataFlowAnalysis<CalleeTy>::getArgumentAccessRange(
512     const CalleeTy *Callee, unsigned ParamNo,
513     const ConstantRange &Offsets) const {
514   auto FnIt = Functions.find(Callee);
515   // Unknown callee (outside of LTO domain or an indirect call).
516   if (FnIt == Functions.end())
517     return UnknownRange;
518   auto &FS = FnIt->second;
519   auto ParamIt = FS.Params.find(ParamNo);
520   if (ParamIt == FS.Params.end())
521     return UnknownRange;
522   auto &Access = ParamIt->second.Range;
523   if (Access.isEmptySet())
524     return Access;
525   if (Access.isFullSet())
526     return UnknownRange;
527   return addOverflowNever(Access, Offsets);
528 }
529 
530 template <typename CalleeTy>
updateOneUse(UseInfo<CalleeTy> & US,bool UpdateToFullSet)531 bool StackSafetyDataFlowAnalysis<CalleeTy>::updateOneUse(UseInfo<CalleeTy> &US,
532                                                          bool UpdateToFullSet) {
533   bool Changed = false;
534   for (auto &KV : US.Calls) {
535     assert(!KV.second.isEmptySet() &&
536            "Param range can't be empty-set, invalid offset range");
537 
538     ConstantRange CalleeRange =
539         getArgumentAccessRange(KV.first.Callee, KV.first.ParamNo, KV.second);
540     if (!US.Range.contains(CalleeRange)) {
541       Changed = true;
542       if (UpdateToFullSet)
543         US.Range = UnknownRange;
544       else
545         US.updateRange(CalleeRange);
546     }
547   }
548   return Changed;
549 }
550 
551 template <typename CalleeTy>
updateOneNode(const CalleeTy * Callee,FunctionInfo<CalleeTy> & FS)552 void StackSafetyDataFlowAnalysis<CalleeTy>::updateOneNode(
553     const CalleeTy *Callee, FunctionInfo<CalleeTy> &FS) {
554   bool UpdateToFullSet = FS.UpdateCount > StackSafetyMaxIterations;
555   bool Changed = false;
556   for (auto &KV : FS.Params)
557     Changed |= updateOneUse(KV.second, UpdateToFullSet);
558 
559   if (Changed) {
560     LLVM_DEBUG(dbgs() << "=== update [" << FS.UpdateCount
561                       << (UpdateToFullSet ? ", full-set" : "") << "] " << &FS
562                       << "\n");
563     // Callers of this function may need updating.
564     for (auto &CallerID : Callers[Callee])
565       WorkList.insert(CallerID);
566 
567     ++FS.UpdateCount;
568   }
569 }
570 
571 template <typename CalleeTy>
runDataFlow()572 void StackSafetyDataFlowAnalysis<CalleeTy>::runDataFlow() {
573   SmallVector<const CalleeTy *, 16> Callees;
574   for (auto &F : Functions) {
575     Callees.clear();
576     auto &FS = F.second;
577     for (auto &KV : FS.Params)
578       for (auto &CS : KV.second.Calls)
579         Callees.push_back(CS.first.Callee);
580 
581     llvm::sort(Callees);
582     Callees.erase(std::unique(Callees.begin(), Callees.end()), Callees.end());
583 
584     for (auto &Callee : Callees)
585       Callers[Callee].push_back(F.first);
586   }
587 
588   updateAllNodes();
589 
590   while (!WorkList.empty()) {
591     const CalleeTy *Callee = WorkList.back();
592     WorkList.pop_back();
593     updateOneNode(Callee);
594   }
595 }
596 
597 #ifndef NDEBUG
598 template <typename CalleeTy>
verifyFixedPoint()599 void StackSafetyDataFlowAnalysis<CalleeTy>::verifyFixedPoint() {
600   WorkList.clear();
601   updateAllNodes();
602   assert(WorkList.empty());
603 }
604 #endif
605 
606 template <typename CalleeTy>
607 const typename StackSafetyDataFlowAnalysis<CalleeTy>::FunctionMap &
run()608 StackSafetyDataFlowAnalysis<CalleeTy>::run() {
609   runDataFlow();
610   LLVM_DEBUG(verifyFixedPoint());
611   return Functions;
612 }
613 
findCalleeFunctionSummary(ValueInfo VI,StringRef ModuleId)614 FunctionSummary *findCalleeFunctionSummary(ValueInfo VI, StringRef ModuleId) {
615   if (!VI)
616     return nullptr;
617   auto SummaryList = VI.getSummaryList();
618   GlobalValueSummary* S = nullptr;
619   for (const auto& GVS : SummaryList) {
620     if (!GVS->isLive())
621       continue;
622     if (const AliasSummary *AS = dyn_cast<AliasSummary>(GVS.get()))
623       if (!AS->hasAliasee())
624         continue;
625     if (!isa<FunctionSummary>(GVS->getBaseObject()))
626       continue;
627     if (GlobalValue::isLocalLinkage(GVS->linkage())) {
628       if (GVS->modulePath() == ModuleId) {
629         S = GVS.get();
630         break;
631       }
632     } else if (GlobalValue::isExternalLinkage(GVS->linkage())) {
633       if (S) {
634         ++NumIndexCalleeMultipleExternal;
635         return nullptr;
636       }
637       S = GVS.get();
638     } else if (GlobalValue::isWeakLinkage(GVS->linkage())) {
639       if (S) {
640         ++NumIndexCalleeMultipleWeak;
641         return nullptr;
642       }
643       S = GVS.get();
644     } else if (GlobalValue::isAvailableExternallyLinkage(GVS->linkage()) ||
645                GlobalValue::isLinkOnceLinkage(GVS->linkage())) {
646       if (SummaryList.size() == 1)
647         S = GVS.get();
648       // According thinLTOResolvePrevailingGUID these are unlikely prevailing.
649     } else {
650       ++NumIndexCalleeUnhandled;
651     }
652   };
653   while (S) {
654     if (!S->isLive() || !S->isDSOLocal())
655       return nullptr;
656     if (FunctionSummary *FS = dyn_cast<FunctionSummary>(S))
657       return FS;
658     AliasSummary *AS = dyn_cast<AliasSummary>(S);
659     if (!AS || !AS->hasAliasee())
660       return nullptr;
661     S = AS->getBaseObject();
662     if (S == AS)
663       return nullptr;
664   }
665   return nullptr;
666 }
667 
findCalleeInModule(const GlobalValue * GV)668 const Function *findCalleeInModule(const GlobalValue *GV) {
669   while (GV) {
670     if (GV->isDeclaration() || GV->isInterposable() || !GV->isDSOLocal())
671       return nullptr;
672     if (const Function *F = dyn_cast<Function>(GV))
673       return F;
674     const GlobalAlias *A = dyn_cast<GlobalAlias>(GV);
675     if (!A)
676       return nullptr;
677     GV = A->getBaseObject();
678     if (GV == A)
679       return nullptr;
680   }
681   return nullptr;
682 }
683 
findParamAccess(const FunctionSummary & FS,uint32_t ParamNo)684 const ConstantRange *findParamAccess(const FunctionSummary &FS,
685                                      uint32_t ParamNo) {
686   assert(FS.isLive());
687   assert(FS.isDSOLocal());
688   for (auto &PS : FS.paramAccesses())
689     if (ParamNo == PS.ParamNo)
690       return &PS.Use;
691   return nullptr;
692 }
693 
resolveAllCalls(UseInfo<GlobalValue> & Use,const ModuleSummaryIndex * Index)694 void resolveAllCalls(UseInfo<GlobalValue> &Use,
695                      const ModuleSummaryIndex *Index) {
696   ConstantRange FullSet(Use.Range.getBitWidth(), true);
697   // Move Use.Calls to a temp storage and repopulate - don't use std::move as it
698   // leaves Use.Calls in an undefined state.
699   UseInfo<GlobalValue>::CallsTy TmpCalls;
700   std::swap(TmpCalls, Use.Calls);
701   for (const auto &C : TmpCalls) {
702     const Function *F = findCalleeInModule(C.first.Callee);
703     if (F) {
704       Use.Calls.emplace(CallInfo<GlobalValue>(F, C.first.ParamNo), C.second);
705       continue;
706     }
707 
708     if (!Index)
709       return Use.updateRange(FullSet);
710     FunctionSummary *FS =
711         findCalleeFunctionSummary(Index->getValueInfo(C.first.Callee->getGUID()),
712                                   C.first.Callee->getParent()->getModuleIdentifier());
713     ++NumModuleCalleeLookupTotal;
714     if (!FS) {
715       ++NumModuleCalleeLookupFailed;
716       return Use.updateRange(FullSet);
717     }
718     const ConstantRange *Found = findParamAccess(*FS, C.first.ParamNo);
719     if (!Found || Found->isFullSet())
720       return Use.updateRange(FullSet);
721     ConstantRange Access = Found->sextOrTrunc(Use.Range.getBitWidth());
722     if (!Access.isEmptySet())
723       Use.updateRange(addOverflowNever(Access, C.second));
724   }
725 }
726 
createGlobalStackSafetyInfo(std::map<const GlobalValue *,FunctionInfo<GlobalValue>> Functions,const ModuleSummaryIndex * Index)727 GVToSSI createGlobalStackSafetyInfo(
728     std::map<const GlobalValue *, FunctionInfo<GlobalValue>> Functions,
729     const ModuleSummaryIndex *Index) {
730   GVToSSI SSI;
731   if (Functions.empty())
732     return SSI;
733 
734   // FIXME: Simplify printing and remove copying here.
735   auto Copy = Functions;
736 
737   for (auto &FnKV : Copy)
738     for (auto &KV : FnKV.second.Params) {
739       resolveAllCalls(KV.second, Index);
740       if (KV.second.Range.isFullSet())
741         KV.second.Calls.clear();
742     }
743 
744   uint32_t PointerSize = Copy.begin()
745                              ->first->getParent()
746                              ->getDataLayout()
747                              .getMaxPointerSizeInBits();
748   StackSafetyDataFlowAnalysis<GlobalValue> SSDFA(PointerSize, std::move(Copy));
749 
750   for (auto &F : SSDFA.run()) {
751     auto FI = F.second;
752     auto &SrcF = Functions[F.first];
753     for (auto &KV : FI.Allocas) {
754       auto &A = KV.second;
755       resolveAllCalls(A, Index);
756       for (auto &C : A.Calls) {
757         A.updateRange(SSDFA.getArgumentAccessRange(C.first.Callee,
758                                                    C.first.ParamNo, C.second));
759       }
760       // FIXME: This is needed only to preserve calls in print() results.
761       A.Calls = SrcF.Allocas.find(KV.first)->second.Calls;
762     }
763     for (auto &KV : FI.Params) {
764       auto &P = KV.second;
765       P.Calls = SrcF.Params.find(KV.first)->second.Calls;
766     }
767     SSI[F.first] = std::move(FI);
768   }
769 
770   return SSI;
771 }
772 
773 } // end anonymous namespace
774 
775 StackSafetyInfo::StackSafetyInfo() = default;
776 
StackSafetyInfo(Function * F,std::function<ScalarEvolution & ()> GetSE)777 StackSafetyInfo::StackSafetyInfo(Function *F,
778                                  std::function<ScalarEvolution &()> GetSE)
779     : F(F), GetSE(GetSE) {}
780 
781 StackSafetyInfo::StackSafetyInfo(StackSafetyInfo &&) = default;
782 
783 StackSafetyInfo &StackSafetyInfo::operator=(StackSafetyInfo &&) = default;
784 
785 StackSafetyInfo::~StackSafetyInfo() = default;
786 
getInfo() const787 const StackSafetyInfo::InfoTy &StackSafetyInfo::getInfo() const {
788   if (!Info) {
789     StackSafetyLocalAnalysis SSLA(*F, GetSE());
790     Info.reset(new InfoTy{SSLA.run()});
791   }
792   return *Info;
793 }
794 
print(raw_ostream & O) const795 void StackSafetyInfo::print(raw_ostream &O) const {
796   getInfo().Info.print(O, F->getName(), dyn_cast<Function>(F));
797 }
798 
getInfo() const799 const StackSafetyGlobalInfo::InfoTy &StackSafetyGlobalInfo::getInfo() const {
800   if (!Info) {
801     std::map<const GlobalValue *, FunctionInfo<GlobalValue>> Functions;
802     for (auto &F : M->functions()) {
803       if (!F.isDeclaration()) {
804         auto FI = GetSSI(F).getInfo().Info;
805         Functions.emplace(&F, std::move(FI));
806       }
807     }
808     Info.reset(new InfoTy{
809         createGlobalStackSafetyInfo(std::move(Functions), Index), {}});
810     for (auto &FnKV : Info->Info) {
811       for (auto &KV : FnKV.second.Allocas) {
812         ++NumAllocaTotal;
813         const AllocaInst *AI = KV.first;
814         if (getStaticAllocaSizeRange(*AI).contains(KV.second.Range)) {
815           Info->SafeAllocas.insert(AI);
816           ++NumAllocaStackSafe;
817         }
818       }
819     }
820     if (StackSafetyPrint)
821       print(errs());
822   }
823   return *Info;
824 }
825 
826 std::vector<FunctionSummary::ParamAccess>
getParamAccesses(ModuleSummaryIndex & Index) const827 StackSafetyInfo::getParamAccesses(ModuleSummaryIndex &Index) const {
828   // Implementation transforms internal representation of parameter information
829   // into FunctionSummary format.
830   std::vector<FunctionSummary::ParamAccess> ParamAccesses;
831   for (const auto &KV : getInfo().Info.Params) {
832     auto &PS = KV.second;
833     // Parameter accessed by any or unknown offset, represented as FullSet by
834     // StackSafety, is handled as the parameter for which we have no
835     // StackSafety info at all. So drop it to reduce summary size.
836     if (PS.Range.isFullSet())
837       continue;
838 
839     ParamAccesses.emplace_back(KV.first, PS.Range);
840     FunctionSummary::ParamAccess &Param = ParamAccesses.back();
841 
842     Param.Calls.reserve(PS.Calls.size());
843     for (auto &C : PS.Calls) {
844       // Parameter forwarded into another function by any or unknown offset
845       // will make ParamAccess::Range as FullSet anyway. So we can drop the
846       // entire parameter like we did above.
847       // TODO(vitalybuka): Return already filtered parameters from getInfo().
848       if (C.second.isFullSet()) {
849         ParamAccesses.pop_back();
850         break;
851       }
852       Param.Calls.emplace_back(C.first.ParamNo,
853                                Index.getOrInsertValueInfo(C.first.Callee),
854                                C.second);
855     }
856   }
857   for (FunctionSummary::ParamAccess &Param : ParamAccesses) {
858     sort(Param.Calls, [](const FunctionSummary::ParamAccess::Call &L,
859                          const FunctionSummary::ParamAccess::Call &R) {
860       return std::tie(L.ParamNo, L.Callee) < std::tie(R.ParamNo, R.Callee);
861     });
862   }
863   return ParamAccesses;
864 }
865 
866 StackSafetyGlobalInfo::StackSafetyGlobalInfo() = default;
867 
StackSafetyGlobalInfo(Module * M,std::function<const StackSafetyInfo & (Function & F)> GetSSI,const ModuleSummaryIndex * Index)868 StackSafetyGlobalInfo::StackSafetyGlobalInfo(
869     Module *M, std::function<const StackSafetyInfo &(Function &F)> GetSSI,
870     const ModuleSummaryIndex *Index)
871     : M(M), GetSSI(GetSSI), Index(Index) {
872   if (StackSafetyRun)
873     getInfo();
874 }
875 
876 StackSafetyGlobalInfo::StackSafetyGlobalInfo(StackSafetyGlobalInfo &&) =
877     default;
878 
879 StackSafetyGlobalInfo &
880 StackSafetyGlobalInfo::operator=(StackSafetyGlobalInfo &&) = default;
881 
882 StackSafetyGlobalInfo::~StackSafetyGlobalInfo() = default;
883 
isSafe(const AllocaInst & AI) const884 bool StackSafetyGlobalInfo::isSafe(const AllocaInst &AI) const {
885   const auto &Info = getInfo();
886   return Info.SafeAllocas.count(&AI);
887 }
888 
print(raw_ostream & O) const889 void StackSafetyGlobalInfo::print(raw_ostream &O) const {
890   auto &SSI = getInfo().Info;
891   if (SSI.empty())
892     return;
893   const Module &M = *SSI.begin()->first->getParent();
894   for (auto &F : M.functions()) {
895     if (!F.isDeclaration()) {
896       SSI.find(&F)->second.print(O, F.getName(), &F);
897       O << "\n";
898     }
899   }
900 }
901 
dump() const902 LLVM_DUMP_METHOD void StackSafetyGlobalInfo::dump() const { print(dbgs()); }
903 
904 AnalysisKey StackSafetyAnalysis::Key;
905 
run(Function & F,FunctionAnalysisManager & AM)906 StackSafetyInfo StackSafetyAnalysis::run(Function &F,
907                                          FunctionAnalysisManager &AM) {
908   return StackSafetyInfo(&F, [&AM, &F]() -> ScalarEvolution & {
909     return AM.getResult<ScalarEvolutionAnalysis>(F);
910   });
911 }
912 
run(Function & F,FunctionAnalysisManager & AM)913 PreservedAnalyses StackSafetyPrinterPass::run(Function &F,
914                                               FunctionAnalysisManager &AM) {
915   OS << "'Stack Safety Local Analysis' for function '" << F.getName() << "'\n";
916   AM.getResult<StackSafetyAnalysis>(F).print(OS);
917   return PreservedAnalyses::all();
918 }
919 
920 char StackSafetyInfoWrapperPass::ID = 0;
921 
StackSafetyInfoWrapperPass()922 StackSafetyInfoWrapperPass::StackSafetyInfoWrapperPass() : FunctionPass(ID) {
923   initializeStackSafetyInfoWrapperPassPass(*PassRegistry::getPassRegistry());
924 }
925 
getAnalysisUsage(AnalysisUsage & AU) const926 void StackSafetyInfoWrapperPass::getAnalysisUsage(AnalysisUsage &AU) const {
927   AU.addRequiredTransitive<ScalarEvolutionWrapperPass>();
928   AU.setPreservesAll();
929 }
930 
print(raw_ostream & O,const Module * M) const931 void StackSafetyInfoWrapperPass::print(raw_ostream &O, const Module *M) const {
932   SSI.print(O);
933 }
934 
runOnFunction(Function & F)935 bool StackSafetyInfoWrapperPass::runOnFunction(Function &F) {
936   auto *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
937   SSI = {&F, [SE]() -> ScalarEvolution & { return *SE; }};
938   return false;
939 }
940 
941 AnalysisKey StackSafetyGlobalAnalysis::Key;
942 
943 StackSafetyGlobalInfo
run(Module & M,ModuleAnalysisManager & AM)944 StackSafetyGlobalAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
945   // FIXME: Lookup Module Summary.
946   FunctionAnalysisManager &FAM =
947       AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
948   return {&M,
949           [&FAM](Function &F) -> const StackSafetyInfo & {
950             return FAM.getResult<StackSafetyAnalysis>(F);
951           },
952           nullptr};
953 }
954 
run(Module & M,ModuleAnalysisManager & AM)955 PreservedAnalyses StackSafetyGlobalPrinterPass::run(Module &M,
956                                                     ModuleAnalysisManager &AM) {
957   OS << "'Stack Safety Analysis' for module '" << M.getName() << "'\n";
958   AM.getResult<StackSafetyGlobalAnalysis>(M).print(OS);
959   return PreservedAnalyses::all();
960 }
961 
962 char StackSafetyGlobalInfoWrapperPass::ID = 0;
963 
StackSafetyGlobalInfoWrapperPass()964 StackSafetyGlobalInfoWrapperPass::StackSafetyGlobalInfoWrapperPass()
965     : ModulePass(ID) {
966   initializeStackSafetyGlobalInfoWrapperPassPass(
967       *PassRegistry::getPassRegistry());
968 }
969 
970 StackSafetyGlobalInfoWrapperPass::~StackSafetyGlobalInfoWrapperPass() = default;
971 
print(raw_ostream & O,const Module * M) const972 void StackSafetyGlobalInfoWrapperPass::print(raw_ostream &O,
973                                              const Module *M) const {
974   SSGI.print(O);
975 }
976 
getAnalysisUsage(AnalysisUsage & AU) const977 void StackSafetyGlobalInfoWrapperPass::getAnalysisUsage(
978     AnalysisUsage &AU) const {
979   AU.setPreservesAll();
980   AU.addRequired<StackSafetyInfoWrapperPass>();
981 }
982 
runOnModule(Module & M)983 bool StackSafetyGlobalInfoWrapperPass::runOnModule(Module &M) {
984   const ModuleSummaryIndex *ImportSummary = nullptr;
985   if (auto *IndexWrapperPass =
986           getAnalysisIfAvailable<ImmutableModuleSummaryIndexWrapperPass>())
987     ImportSummary = IndexWrapperPass->getIndex();
988 
989   SSGI = {&M,
990           [this](Function &F) -> const StackSafetyInfo & {
991             return getAnalysis<StackSafetyInfoWrapperPass>(F).getResult();
992           },
993           ImportSummary};
994   return false;
995 }
996 
needsParamAccessSummary(const Module & M)997 bool llvm::needsParamAccessSummary(const Module &M) {
998   if (StackSafetyRun)
999     return true;
1000   for (auto &F : M.functions())
1001     if (F.hasFnAttribute(Attribute::SanitizeMemTag))
1002       return true;
1003   return false;
1004 }
1005 
generateParamAccessSummary(ModuleSummaryIndex & Index)1006 void llvm::generateParamAccessSummary(ModuleSummaryIndex &Index) {
1007   if (!Index.hasParamAccess())
1008     return;
1009   const ConstantRange FullSet(FunctionSummary::ParamAccess::RangeWidth, true);
1010 
1011   auto CountParamAccesses = [&](auto &Stat) {
1012     if (!AreStatisticsEnabled())
1013       return;
1014     for (auto &GVS : Index)
1015       for (auto &GV : GVS.second.SummaryList)
1016         if (FunctionSummary *FS = dyn_cast<FunctionSummary>(GV.get()))
1017           Stat += FS->paramAccesses().size();
1018   };
1019 
1020   CountParamAccesses(NumCombinedParamAccessesBefore);
1021 
1022   std::map<const FunctionSummary *, FunctionInfo<FunctionSummary>> Functions;
1023 
1024   // Convert the ModuleSummaryIndex to a FunctionMap
1025   for (auto &GVS : Index) {
1026     for (auto &GV : GVS.second.SummaryList) {
1027       FunctionSummary *FS = dyn_cast<FunctionSummary>(GV.get());
1028       if (!FS || FS->paramAccesses().empty())
1029         continue;
1030       if (FS->isLive() && FS->isDSOLocal()) {
1031         FunctionInfo<FunctionSummary> FI;
1032         for (auto &PS : FS->paramAccesses()) {
1033           auto &US =
1034               FI.Params
1035                   .emplace(PS.ParamNo, FunctionSummary::ParamAccess::RangeWidth)
1036                   .first->second;
1037           US.Range = PS.Use;
1038           for (auto &Call : PS.Calls) {
1039             assert(!Call.Offsets.isFullSet());
1040             FunctionSummary *S =
1041                 findCalleeFunctionSummary(Call.Callee, FS->modulePath());
1042             ++NumCombinedCalleeLookupTotal;
1043             if (!S) {
1044               ++NumCombinedCalleeLookupFailed;
1045               US.Range = FullSet;
1046               US.Calls.clear();
1047               break;
1048             }
1049             US.Calls.emplace(CallInfo<FunctionSummary>(S, Call.ParamNo),
1050                              Call.Offsets);
1051           }
1052         }
1053         Functions.emplace(FS, std::move(FI));
1054       }
1055       // Reset data for all summaries. Alive and DSO local will be set back from
1056       // of data flow results below. Anything else will not be accessed
1057       // by ThinLTO backend, so we can save on bitcode size.
1058       FS->setParamAccesses({});
1059     }
1060   }
1061   NumCombinedDataFlowNodes += Functions.size();
1062   StackSafetyDataFlowAnalysis<FunctionSummary> SSDFA(
1063       FunctionSummary::ParamAccess::RangeWidth, std::move(Functions));
1064   for (auto &KV : SSDFA.run()) {
1065     std::vector<FunctionSummary::ParamAccess> NewParams;
1066     NewParams.reserve(KV.second.Params.size());
1067     for (auto &Param : KV.second.Params) {
1068       // It's not needed as FullSet is processed the same as a missing value.
1069       if (Param.second.Range.isFullSet())
1070         continue;
1071       NewParams.emplace_back();
1072       FunctionSummary::ParamAccess &New = NewParams.back();
1073       New.ParamNo = Param.first;
1074       New.Use = Param.second.Range; // Only range is needed.
1075     }
1076     const_cast<FunctionSummary *>(KV.first)->setParamAccesses(
1077         std::move(NewParams));
1078   }
1079 
1080   CountParamAccesses(NumCombinedParamAccessesAfter);
1081 }
1082 
1083 static const char LocalPassArg[] = "stack-safety-local";
1084 static const char LocalPassName[] = "Stack Safety Local Analysis";
1085 INITIALIZE_PASS_BEGIN(StackSafetyInfoWrapperPass, LocalPassArg, LocalPassName,
1086                       false, true)
1087 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
1088 INITIALIZE_PASS_END(StackSafetyInfoWrapperPass, LocalPassArg, LocalPassName,
1089                     false, true)
1090 
1091 static const char GlobalPassName[] = "Stack Safety Analysis";
1092 INITIALIZE_PASS_BEGIN(StackSafetyGlobalInfoWrapperPass, DEBUG_TYPE,
1093                       GlobalPassName, false, true)
1094 INITIALIZE_PASS_DEPENDENCY(StackSafetyInfoWrapperPass)
1095 INITIALIZE_PASS_DEPENDENCY(ImmutableModuleSummaryIndexWrapperPass)
1096 INITIALIZE_PASS_END(StackSafetyGlobalInfoWrapperPass, DEBUG_TYPE,
1097                     GlobalPassName, false, true)
1098