[stack-safety] Refactoring StackSafetyAnalysis to be accsessible from other passes

[gilang]

Summary

Refactoring StackSafetyAnalysis to expose information about params and allocas outside of its compilation unit. Previously, the analysis does not expose params and allocas information externally and only provides print method for the user. In order to use the information programmatically from other passes, the data structure has to be exposed in the header files under llvm namespace, instead of anonymous namespace.

No modification of the original algorithm, therefore, it can utilize the original test scripts for testing.

[vlad.tsyrklevich]

Hi gilang, out of curiosity what is your intended use of the StackSafetyAnalysis? Is it necessary to expose all of the information exposed here?

[gilang]

In D63703#1556296, @vlad.tsyrklevich wrote:

Hi gilang, out of curiosity what is your intended use of the StackSafetyAnalysis? Is it necessary to expose all of the information exposed here?

Hi, thanks for the response.

We are currently planning to utilize StackSafetyAnalysis for memory protection schemes that continue our prior work on ARM pointer authentication (https://arxiv.org/abs/1811.09189).

Since the code seems to be dormant around 6 months, and I haven’t found recent discussion regarding this, I thought that I can suggest some refactoring so other users can use this analysis in other passes, similar with the rest of LLVM analyses suites (e.g., alias analysis). This is also to avoid changing the StackSafetyAnalysis code for every uses by different users, which will be redundant. However, if some other approach there is already an internal development to improve its API that the public currently does not know of, this changes can be dropped.

Right now we plan to use the exposed objects to identify stack-allocated variable which has out-of-bound access and use the information in our memory protection scheme.

PS: Apologize for double posting via e-mail. Seems that the email reply is not automatically added to the comment here. I replied via e-mail to also add a team member without Phabricator account to the loop. :)

[eugenis]

llvm/lib/Analysis/StackSafetyAnalysis.cpp

127

I'm not sure I understand this comment. Are the overloads for the new types shadowing the common overloads? Would it help to move them out of the anonymous namespace?

Btw, please upload change with full context, it's hard to review otherwise.
See https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

[gilang]

Updating the patch with full context. Forgetting to add -U argument.

[gilang]

llvm/lib/Analysis/StackSafetyAnalysis.cpp

127

Hi. Apologize for the inconvenience. I have updated the patch to include the full context.

So, without adding the using statement inside the block, the code will not compile due to overloading failure. Putting the operator overload declaration outside the anonymous namespace also does not help.

Clang used: 6.0.0-1ubuntu2

[eugenis]

If that works for your use case, perhaps a better interface would be a query function that returns offset range or even bool is_safe given an AllocaInst* ?

llvm/lib/Analysis/StackSafetyAnalysis.cpp

127

I'm not sure where the problem is, but there are no such "using ::operator<<" declarations anywhere else in llvm, so there must be a way around it.

[gilang]

In D63703#1561326, @eugenis wrote:

If that works for your use case, perhaps a better interface would be a query function that returns offset range or even bool is_safe given an AllocaInst* ?

I initially thought the same as it is more useful and more similar with the pattern used in alias analysis, but since the internals uses SmallVector to store the information, it would be inefficient if the user also need to iterate every AllocaInst* in their own pass, since the querying will takes O(n) complexity. I believe refactoring the SmallVector to, for instance, DenseMap, will require more code changes.

My first thought was to minimize the change to the original code so exposing the iterator to the structure was the faster choice, as the user can also access other information in the struct such as size. However, I will try to implement it as you suggested.

[gilang]

Revising the proposal to create a query function instead, by refactoring some internals to use DenseMap instead of SmallVector and provide isSafe(AllocaInst*) and getSizeOf(AllocaInst*) function.

[eugenis]

llvm/lib/Analysis/StackSafetyAnalysis.cpp

149

This could be very expensive - AllocaInfo is large, and DenseMap starts with 64 buckets. I wonder if SmallMapVector would work better here.

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_STACKSAFETYANALYSIS_H
 #define LLVM_ANALYSIS_STACKSAFETYANALYSIS_H
 
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
+#include "llvm/IR/Instructions.h"
 
 namespace llvm {
 
 /// Interface to access stack safety analysis results for single function.
 class StackSafetyInfo {
 public:
   struct FunctionInfo;
 
 private:
   std::unique_ptr<FunctionInfo> Info;
 
 public:
   StackSafetyInfo();
   StackSafetyInfo(FunctionInfo &&Info);
   StackSafetyInfo(StackSafetyInfo &&);
   StackSafetyInfo &operator=(StackSafetyInfo &&);
   ~StackSafetyInfo();
 
   // TODO: Add useful for client methods.
   void print(raw_ostream &O) const;
+  bool isSafe(const AllocaInst *AI) const;
+  uint64_t getSizeOf(const AllocaInst *AI) const;
 };
 
 /// StackSafetyInfo wrapper for the new pass manager.
 class StackSafetyAnalysis : public AnalysisInfoMixin<StackSafetyAnalysis> {
   friend AnalysisInfoMixin<StackSafetyAnalysis>;
   static AnalysisKey Key;
 
 public:

llvm/lib/Analysis/StackSafetyAnalysis.cpp

explicit ParamInfo(unsigned PointerSize, const Argument *Arg)
       : Arg(Arg), Use(PointerSize) {}
 
   StringRef getName() const { return Arg ? Arg->getName() : "<N/A>"; }
 };
 
 raw_ostream &operator<<(raw_ostream &OS, const ParamInfo &P) {
   return OS << P.getName() << "[]: " << P.Use;
 }
 
 /// Calculate the allocation size of a given alloca. Returns 0 if the
 /// size can not be statically determined.
 uint64_t getStaticAllocaAllocationSize(const AllocaInst *AI) {
   const DataLayout &DL = AI->getModule()->getDataLayout();
   uint64_t Size = DL.getTypeAllocSize(AI->getAllocatedType());
   if (AI->isArrayAllocation()) {
     auto C = dyn_cast<ConstantInt>(AI->getArraySize());
     if (!C)
       return 0;
     Size *= C->getZExtValue();
   }
   return Size;
 }
 
 } // end anonymous namespace
 
 /// Describes uses of allocas and parameters inside of a single function.
 struct StackSafetyInfo::FunctionInfo {
   // May be a Function or a GlobalAlias
   const GlobalValue *GV = nullptr;
   // Informations about allocas uses.
-  SmallVector<AllocaInfo, 4> Allocas;
+  DenseMap<const AllocaInst*, AllocaInfo> Allocas;
   // Informations about parameters uses.
   SmallVector<ParamInfo, 4> Params;
   // TODO: describe return value as depending on one or more of its arguments.
 
   // StackSafetyDataFlowAnalysis counter stored here for faster access.
   int UpdateCount = 0;
 
   FunctionInfo(const StackSafetyInfo &SSI) : FunctionInfo(*SSI.Info) {}
void print(raw_ostream &O) const {
     // StackSafetyDataFlowAnalysis. Calls and parameters are irrelevant then.
     O << "  @" << getName() << (IsDSOLocal() ? "" : " dso_preemptable")
       << (IsInterposable() ? " interposable" : "") << "\n";
     O << "    args uses:\n";
     for (auto &P : Params)
       O << "      " << P << "\n";
     O << "    allocas uses:\n";
     for (auto &AS : Allocas)
-      O << "      " << AS << "\n";
+      O << "      " << AS.second << "\n";
   }
 
 private:
   FunctionInfo(const FunctionInfo &) = default;
 };
 
 StackSafetyInfo::FunctionInfo::FunctionInfo(const GlobalAlias *A) : GV(A) {
   unsigned PointerSize = A->getParent()->getDataLayout().getPointerSizeInBits();
StackSafetyInfo StackSafetyLocalAnalysis::run() {
   StackSafetyInfo::FunctionInfo Info(&F);
   assert(!F.isDeclaration() &&
          "Can't run StackSafety on a function declaration");
 
   LLVM_DEBUG(dbgs() << "[StackSafety] " << F.getName() << "\n");
 
   for (auto &I : instructions(F)) {
     if (auto AI = dyn_cast<AllocaInst>(&I)) {
-      Info.Allocas.emplace_back(PointerSize, AI,
-                                getStaticAllocaAllocationSize(AI));
-      AllocaInfo &AS = Info.Allocas.back();
+      auto inserted = Info.Allocas.insert(std::make_pair(
+          AI, AllocaInfo{PointerSize, AI, getStaticAllocaAllocationSize(AI)}));
+      // Obtain the AllocaInfo
+      AllocaInfo &AS = inserted.first->second; 
       analyzeAllUses(AI, AS.Use);
     }
   }
 
   for (const Argument &A : make_range(F.arg_begin(), F.arg_end())) {
     Info.Params.emplace_back(PointerSize, &A);
     ParamInfo &PS = Info.Params.back();
     analyzeAllUses(&A, PS.Use);
bool StackSafetyDataFlowAnalysis::updateOneUse(UseInfo &US,
   return Changed;
 }
 
 void StackSafetyDataFlowAnalysis::updateOneNode(
     const GlobalValue *Callee, StackSafetyInfo::FunctionInfo &FS) {
   bool UpdateToFullSet = FS.UpdateCount > StackSafetyMaxIterations;
   bool Changed = false;
   for (auto &AS : FS.Allocas)
-    Changed |= updateOneUse(AS.Use, UpdateToFullSet);
+    Changed |= updateOneUse(AS.second.Use, UpdateToFullSet);
   for (auto &PS : FS.Params)
     Changed |= updateOneUse(PS.Use, UpdateToFullSet);
 
   if (Changed) {
     LLVM_DEBUG(dbgs() << "=== update [" << FS.UpdateCount
                       << (UpdateToFullSet ? ", full-set" : "") << "] "
                       << FS.getName() << "\n");
     // Callers of this function may need updating.
     for (auto &CallerID : Callers[Callee])
       WorkList.insert(CallerID);
 
     ++FS.UpdateCount;
   }
 }
 
 void StackSafetyDataFlowAnalysis::runDataFlow() {
   Callers.clear();
   WorkList.clear();
 
   SmallVector<const GlobalValue *, 16> Callees;
   for (auto &F : Functions) {
     Callees.clear();
     StackSafetyInfo::FunctionInfo &FS = F.second;
     for (auto &AS : FS.Allocas)
-      for (auto &CS : AS.Use.Calls)
+      for (auto &CS : AS.second.Use.Calls)
         Callees.push_back(CS.Callee);
     for (auto &PS : FS.Params)
       for (auto &CS : PS.Use.Calls)
         Callees.push_back(CS.Callee);
 
     llvm::sort(Callees);
     Callees.erase(std::unique(Callees.begin(), Callees.end()), Callees.end());
 

 
 StackSafetyInfo::StackSafetyInfo(FunctionInfo &&Info)
     : Info(new FunctionInfo(std::move(Info))) {}
 
 StackSafetyInfo::~StackSafetyInfo() = default;
 
 void StackSafetyInfo::print(raw_ostream &O) const { Info->print(O); }
 
+bool StackSafetyInfo::isSafe(const AllocaInst *AI) const {
+  auto iter = Info->Allocas.find(AI);
+  assert(iter != Info->Allocas.end() && "The specified AllocaInst is not a part of this function");
+
+  auto &UI = iter->second.Use;
+  if (UI.Range.isFullSet()) // FullSet indicates that the analysis cannot prove the safety of the alloca
+    return false;
+  else if (UI.Range.getLower().getLimitedValue() < 0)
+    return false;
+  else if (UI.Range.getUpper().getLimitedValue() > iter->second.Size)
+    return false;
+  return true;
+}
+
+uint64_t StackSafetyInfo::getSizeOf(const AllocaInst *AI) const {
+  auto iter = Info->Allocas.find(AI);
+  assert(iter != Info->Allocas.end() &&
+         "The specified AllocaInst is not a part of this function");
+
+  return iter->second.Size;
+}
+
 AnalysisKey StackSafetyAnalysis::Key;
 
 StackSafetyInfo StackSafetyAnalysis::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   StackSafetyLocalAnalysis SSLA(F, AM.getResult<ScalarEvolutionAnalysis>(F));
   return SSLA.run();
 }