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