[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

99

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

 // Stack Safety Analysis detects allocas and arguments with safe access.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_STACKSAFETYANALYSIS_H
 #define LLVM_ANALYSIS_STACKSAFETYANALYSIS_H
 
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ConstantRange.h"
 #include "llvm/Pass.h"
 
 namespace llvm {
+// Forward declaring
+class AllocaInst;
+class Argument;
 
 /// Interface to access stack safety analysis results for single function.
 class StackSafetyInfo {
 public:
   struct FunctionInfo;
 
+  /// Describes use of address in as a function call argument.
+  struct PassAsArgInfo {
+    /// Function being called.
+    const GlobalValue *Callee = nullptr;
+    /// Index of argument which pass address.
+    size_t ParamNo = 0;
+    // Offset range of address from base address (alloca or calling function
+    // argument).
+    // Range should never set to empty-set, that is an invalid access range
+    // that can cause empty-set to be propagated with ConstantRange::add
+    ConstantRange Offset;
+    PassAsArgInfo(const GlobalValue *Callee, size_t ParamNo,
+                  ConstantRange Offset)
+        : Callee(Callee), ParamNo(ParamNo), Offset(Offset) {}
+
+    StringRef getName() const { return Callee->getName(); }
+  };
+
+  /// Describe uses of address (alloca or parameter) inside of the function.
+  struct UseInfo {
+    // Access range if the address (alloca or parameters).
+    // It is allowed to be empty-set when there are no known accesses.
+    ConstantRange Range;
+
+    // List of calls which pass address as an argument.
+    SmallVector<PassAsArgInfo, 4> Calls;
+
+    explicit UseInfo(unsigned PointerSize) : Range{PointerSize, false} {}
+
+    void updateRange(ConstantRange R) { Range = Range.unionWith(R); }
+  };
+
+  struct AllocaInfo {
+    const AllocaInst *AI = nullptr;
+    uint64_t Size = 0;
+    UseInfo Use;
+
+    AllocaInfo(unsigned PointerSize, const AllocaInst *AI, uint64_t Size)
+        : AI(AI), Size(Size), Use(PointerSize) {}
+
+    StringRef getName() const;
+  };
+
+  struct ParamInfo {
+    const Argument *Arg = nullptr;
+    UseInfo Use;
+
+    explicit ParamInfo(unsigned PointerSize, const Argument *Arg)
+        : Arg(Arg), Use(PointerSize) {}
+
+    StringRef getName() const;
+  };
+
+  using AllocaIterator = 
+      SmallVector<StackSafetyInfo::AllocaInfo, 4>::const_iterator;
+  using ParamsIterator =
+      SmallVector<StackSafetyInfo::ParamInfo, 4>::const_iterator;
+
 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;
+
+  iterator_range<AllocaIterator> allocas() const;
+  iterator_range<ParamsIterator> params() 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

const SCEV *visitUnknown(const SCEVUnknown *Expr) {
     // This can be inttoptr(AllocaPtr) and SCEV would not unwrap
     // it for us.
     if (Expr->getValue() == AllocaPtr)
       return SE.getZero(Expr->getType());
     return Expr;
   }
 };
 
-/// Describes use of address in as a function call argument.
-struct PassAsArgInfo {
-  /// Function being called.
-  const GlobalValue *Callee = nullptr;
-  /// Index of argument which pass address.
-  size_t ParamNo = 0;
-  // Offset range of address from base address (alloca or calling function
-  // argument).
-  // Range should never set to empty-set, that is an invalid access range
-  // that can cause empty-set to be propagated with ConstantRange::add
-  ConstantRange Offset;
-  PassAsArgInfo(const GlobalValue *Callee, size_t ParamNo, ConstantRange Offset)
-      : Callee(Callee), ParamNo(ParamNo), Offset(Offset) {}
-
-  StringRef getName() const { return Callee->getName(); }
-};
-
-raw_ostream &operator<<(raw_ostream &OS, const PassAsArgInfo &P) {
+raw_ostream &operator<<(raw_ostream &OS,
+                        const StackSafetyInfo::PassAsArgInfo &P) {
   return OS << "@" << P.getName() << "(arg" << P.ParamNo << ", " << P.Offset
             << ")";
 }
 
-/// Describe uses of address (alloca or parameter) inside of the function.
-struct UseInfo {
-  // Access range if the address (alloca or parameters).
-  // It is allowed to be empty-set when there are no known accesses.
-  ConstantRange Range;
-
-  // List of calls which pass address as an argument.
-  SmallVector<PassAsArgInfo, 4> Calls;
-
-  explicit UseInfo(unsigned PointerSize) : Range{PointerSize, false} {}
-
-  void updateRange(ConstantRange R) { Range = Range.unionWith(R); }
-};
-
-raw_ostream &operator<<(raw_ostream &OS, const UseInfo &U) {
+raw_ostream &operator<<(raw_ostream &OS, const StackSafetyInfo::UseInfo &U) {
   OS << U.Range;
   for (auto &Call : U.Calls)
     OS << ", " << Call;
   return OS;
 }
 
-struct AllocaInfo {
-  const AllocaInst *AI = nullptr;
-  uint64_t Size = 0;
-  UseInfo Use;
-
-  AllocaInfo(unsigned PointerSize, const AllocaInst *AI, uint64_t Size)
-      : AI(AI), Size(Size), Use(PointerSize) {}
-
-  StringRef getName() const { return AI->getName(); }
-};
-
-raw_ostream &operator<<(raw_ostream &OS, const AllocaInfo &A) {
+raw_ostream &operator<<(raw_ostream &OS, const StackSafetyInfo::AllocaInfo &A) {
   return OS << A.getName() << "[" << A.Size << "]: " << A.Use;
 }
 
-struct ParamInfo {
-  const Argument *Arg = nullptr;
-  UseInfo Use;
-
-  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) {
+raw_ostream &operator<<(raw_ostream &OS, const StackSafetyInfo::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;
+  SmallVector<StackSafetyInfo::AllocaInfo, 4> Allocas;
   // Informations about parameters uses.
-  SmallVector<ParamInfo, 4> Params;
+  SmallVector<StackSafetyInfo::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) {}
 
   explicit FunctionInfo(const Function *F) : GV(F){};
   // Creates FunctionInfo that forwards all the parameters to the aliasee.
   explicit FunctionInfo(const GlobalAlias *A);
 
   FunctionInfo(FunctionInfo &&) = default;
 
   bool IsDSOLocal() const { return GV->isDSOLocal(); };
 
   bool IsInterposable() const { return GV->isInterposable(); };
 
   StringRef getName() const { return GV->getName(); }
 
   void print(raw_ostream &O) const {
     // TODO: Consider different printout format after
     // StackSafetyDataFlowAnalysis. Calls and parameters are irrelevant then.
+
+    // After refactoring, the operator overload that is inside the anonymous
+    // namespace is no longer accessible directly when being passed with the
+    // refactored ParamInfo and AllocaInfo objects
+    using ::operator<<;
+
     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";
StackSafetyInfo::FunctionInfo::FunctionInfo(const GlobalAlias *A) : GV(A) {
   // 'Forward' all parameters to this alias to the aliasee
   for (unsigned ArgNo = 0; ArgNo < Type->getNumParams(); ArgNo++) {
     Params.emplace_back(PointerSize, nullptr);
     UseInfo &US = Params.back().Use;
     US.Calls.emplace_back(Aliasee, ArgNo, ConstantRange(APInt(PointerSize, 0)));
   }
 }
 
+iterator_range<StackSafetyInfo::AllocaIterator> StackSafetyInfo::allocas() const {
+  assert(Info != nullptr && "FunctionInfo is null");
+  return make_range(Info->Allocas.begin(), Info->Allocas.end());
+}
+iterator_range<StackSafetyInfo::ParamsIterator> StackSafetyInfo::params() const {
+  assert(Info != nullptr && "FunctionInfo is null");
+  return make_range(Info->Params.begin(), Info->Params.end());
+}
+
+StringRef StackSafetyInfo::AllocaInfo::getName() const { return AI->getName(); }
+
+StringRef StackSafetyInfo::ParamInfo::getName() const {
+  return Arg ? Arg->getName() : "<N/A>";
+}
+
 namespace {
 
 class StackSafetyLocalAnalysis {
   const Function &F;
   const DataLayout &DL;
   ScalarEvolution &SE;
   unsigned PointerSize = 0;
 
   const ConstantRange UnknownRange;
 
   ConstantRange offsetFromAlloca(Value *Addr, const Value *AllocaPtr);
   ConstantRange getAccessRange(Value *Addr, const Value *AllocaPtr,
                                uint64_t AccessSize);
   ConstantRange getMemIntrinsicAccessRange(const MemIntrinsic *MI, const Use &U,
                                            const Value *AllocaPtr);
 
-  bool analyzeAllUses(const Value *Ptr, UseInfo &AS);
+  bool analyzeAllUses(const Value *Ptr, StackSafetyInfo::UseInfo &AS);
 
   ConstantRange getRange(uint64_t Lower, uint64_t Upper) const {
     return ConstantRange(APInt(PointerSize, Lower), APInt(PointerSize, Upper));
   }
 
 public:
   StackSafetyLocalAnalysis(const Function &F, ScalarEvolution &SE)
       : F(F), DL(F.getParent()->getDataLayout()), SE(SE),
ConstantRange StackSafetyLocalAnalysis::getMemIntrinsicAccessRange(
   if (!Len)
     return UnknownRange;
   ConstantRange AccessRange = getAccessRange(U, AllocaPtr, Len->getZExtValue());
   return AccessRange;
 }
 
 /// The function analyzes all local uses of Ptr (alloca or argument) and
 /// calculates local access range and all function calls where it was used.
-bool StackSafetyLocalAnalysis::analyzeAllUses(const Value *Ptr, UseInfo &US) {
+bool StackSafetyLocalAnalysis::analyzeAllUses(const Value *Ptr,
+                                              StackSafetyInfo::UseInfo &US) {
   SmallPtrSet<const Value *, 16> Visited;
   SmallVector<const Value *, 8> WorkList;
   WorkList.push_back(Ptr);
 
   // A DFS search through all uses of the alloca in bitcasts/PHI/GEPs/etc.
   while (!WorkList.empty()) {
     const Value *V = WorkList.pop_back_val();
     for (const Use &UI : V->uses()) {
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();
+      StackSafetyInfo::AllocaInfo &AS = Info.Allocas.back();
       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();
+    StackSafetyInfo::ParamInfo &PS = Info.Params.back();
     analyzeAllUses(&A, PS.Use);
   }
 
   LLVM_DEBUG(dbgs() << "[StackSafety] done\n");
   LLVM_DEBUG(Info.print(dbgs()));
   return StackSafetyInfo(std::move(Info));
 }
 
 class StackSafetyDataFlowAnalysis {
   using FunctionMap =
       std::map<const GlobalValue *, StackSafetyInfo::FunctionInfo>;
 
   FunctionMap Functions;
   // Callee-to-Caller multimap.
   DenseMap<const GlobalValue *, SmallVector<const GlobalValue *, 4>> Callers;
   SetVector<const GlobalValue *> WorkList;
 
   unsigned PointerSize = 0;
   const ConstantRange UnknownRange;
 
   ConstantRange getArgumentAccessRange(const GlobalValue *Callee,
                                        unsigned ParamNo) const;
-  bool updateOneUse(UseInfo &US, bool UpdateToFullSet);
+  bool updateOneUse(StackSafetyInfo::UseInfo &US, bool UpdateToFullSet);
   void updateOneNode(const GlobalValue *Callee,
                      StackSafetyInfo::FunctionInfo &FS);
   void updateOneNode(const GlobalValue *Callee) {
     updateOneNode(Callee, Functions.find(Callee)->second);
   }
   void updateAllNodes() {
     for (auto &F : Functions)
       updateOneNode(F.first, F.second);
StackSafetyDataFlowAnalysis::getArgumentAccessRange(const GlobalValue *Callee,
   // unit.
   if (!FS.IsDSOLocal() || FS.IsInterposable())
     return UnknownRange;
   if (ParamNo >= FS.Params.size()) // possibly vararg
     return UnknownRange;
   return FS.Params[ParamNo].Use.Range;
 }
 
-bool StackSafetyDataFlowAnalysis::updateOneUse(UseInfo &US,
+bool StackSafetyDataFlowAnalysis::updateOneUse(StackSafetyInfo::UseInfo &US,
                                                bool UpdateToFullSet) {
   bool Changed = false;
   for (auto &CS : US.Calls) {
     assert(!CS.Offset.isEmptySet() &&
            "Param range can't be empty-set, invalid offset range");
 
     ConstantRange CalleeRange = getArgumentAccessRange(CS.Callee, CS.ParamNo);
     CalleeRange = CalleeRange.add(CS.Offset);