[safestack] Inline safestack pointer access when possible.

Summary:
This adds an -mllvm flag that forces the use of a runtime function call to
get the unsafe stack pointer, the same that is currently used on non-x86, non-aarch64 android.
The call may be inlined.

Reviewers: pcc

Subscribers: aemerson, kristof.beyls, hiraditya, llvm-commits

Differential Revision: https://reviews.llvm.org/D37405

llvm-svn: 323259
diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index 51233be..841a5ec 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -24,6 +24,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
@@ -61,6 +62,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
 #include <cassert>
@@ -88,6 +90,13 @@
 
 } // namespace llvm
 
+/// Use __safestack_pointer_address even if the platform has a faster way of
+/// access safe stack pointer.
+static cl::opt<bool>
+    SafeStackUsePointerAddress("safestack-use-pointer-address",
+                                  cl::init(false), cl::Hidden);
+
+
 namespace {
 
 /// Rewrite an SCEV expression for a memory access address to an expression that
@@ -191,6 +200,9 @@
   bool IsAccessSafe(Value *Addr, uint64_t Size, const Value *AllocaPtr,
                     uint64_t AllocaSize);
 
+  bool ShouldInlinePointerAddress(CallSite &CS);
+  void TryInlinePointerAddress();
+
 public:
   SafeStack(Function &F, const TargetLoweringBase &TL, const DataLayout &DL,
             ScalarEvolution &SE)
@@ -695,6 +707,35 @@
   }
 }
 
+bool SafeStack::ShouldInlinePointerAddress(CallSite &CS) {
+  Function *Callee = CS.getCalledFunction();
+  if (CS.hasFnAttr(Attribute::AlwaysInline) && isInlineViable(*Callee))
+    return true;
+  if (Callee->isInterposable() || Callee->hasFnAttribute(Attribute::NoInline) ||
+      CS.isNoInline())
+    return false;
+  return true;
+}
+
+void SafeStack::TryInlinePointerAddress() {
+  if (!isa<CallInst>(UnsafeStackPtr))
+    return;
+
+  if(F.hasFnAttribute(Attribute::OptimizeNone))
+    return;
+
+  CallSite CS(UnsafeStackPtr);
+  Function *Callee = CS.getCalledFunction();
+  if (!Callee || Callee->isDeclaration())
+    return;
+
+  if (!ShouldInlinePointerAddress(CS))
+    return;
+
+  InlineFunctionInfo IFI;
+  InlineFunction(CS, IFI);
+}
+
 bool SafeStack::run() {
   assert(F.hasFnAttribute(Attribute::SafeStack) &&
          "Can't run SafeStack on a function without the attribute");
@@ -731,7 +772,13 @@
     ++NumUnsafeStackRestorePointsFunctions;
 
   IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt());
-  UnsafeStackPtr = TL.getSafeStackPointerLocation(IRB);
+  if (SafeStackUsePointerAddress) {
+    Value *Fn = F.getParent()->getOrInsertFunction(
+        "__safestack_pointer_address", StackPtrTy->getPointerTo(0));
+    UnsafeStackPtr = IRB.CreateCall(Fn);
+  } else {
+    UnsafeStackPtr = TL.getSafeStackPointerLocation(IRB);
+  }
 
   // Load the current stack pointer (we'll also use it as a base pointer).
   // FIXME: use a dedicated register for it ?
@@ -779,6 +826,8 @@
     IRB.CreateStore(BasePointer, UnsafeStackPtr);
   }
 
+  TryInlinePointerAddress();
+
   DEBUG(dbgs() << "[SafeStack]     safestack applied\n");
   return true;
 }