Index: llvm/trunk/lib/CodeGen/SafeStack.cpp
===================================================================
--- llvm/trunk/lib/CodeGen/SafeStack.cpp
+++ llvm/trunk/lib/CodeGen/SafeStack.cpp
@@ -144,7 +144,8 @@
   Value *moveStaticAllocasToUnsafeStack(IRBuilder<> &IRB, Function &F,
                                         ArrayRef<AllocaInst *> StaticAllocas,
                                         ArrayRef<Argument *> ByValArguments,
-                                        ArrayRef<ReturnInst *> Returns);
+                                        ArrayRef<ReturnInst *> Returns,
+                                        Instruction *BasePointer);
 
   /// \brief Generate code to restore the stack after all stack restore points
   /// in \p StackRestorePoints.
@@ -431,6 +432,8 @@
 SafeStack::createStackRestorePoints(IRBuilder<> &IRB, Function &F,
                                     ArrayRef<Instruction *> StackRestorePoints,
                                     Value *StaticTop, bool NeedDynamicTop) {
+  assert(StaticTop && "The stack top isn't set.");
+
   if (StackRestorePoints.empty())
     return nullptr;
 
@@ -441,19 +444,13 @@
   // runtime itself.
 
   AllocaInst *DynamicTop = nullptr;
-  if (NeedDynamicTop)
+  if (NeedDynamicTop) {
     // If we also have dynamic alloca's, the stack pointer value changes
     // throughout the function. For now we store it in an alloca.
     DynamicTop = IRB.CreateAlloca(StackPtrTy, /*ArraySize=*/nullptr,
                                   "unsafe_stack_dynamic_ptr");
-
-  if (!StaticTop)
-    // We need the original unsafe stack pointer value, even if there are
-    // no unsafe static allocas.
-    StaticTop = IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
-
-  if (NeedDynamicTop)
     IRB.CreateStore(StaticTop, DynamicTop);
+  }
 
   // Restore current stack pointer after longjmp/exception catch.
   for (Instruction *I : StackRestorePoints) {
@@ -467,29 +464,18 @@
   return DynamicTop;
 }
 
+/// We explicitly compute and set the unsafe stack layout for all unsafe
+/// static alloca instructions. We save the unsafe "base pointer" in the
+/// prologue into a local variable and restore it in the epilogue.
 Value *SafeStack::moveStaticAllocasToUnsafeStack(
     IRBuilder<> &IRB, Function &F, ArrayRef<AllocaInst *> StaticAllocas,
-    ArrayRef<Argument *> ByValArguments, ArrayRef<ReturnInst *> Returns) {
+    ArrayRef<Argument *> ByValArguments, ArrayRef<ReturnInst *> Returns,
+    Instruction *BasePointer) {
   if (StaticAllocas.empty() && ByValArguments.empty())
-    return nullptr;
+    return BasePointer;
 
   DIBuilder DIB(*F.getParent());
 
-  // We explicitly compute and set the unsafe stack layout for all unsafe
-  // static alloca instructions. We save the unsafe "base pointer" in the
-  // prologue into a local variable and restore it in the epilogue.
-
-  // Load the current stack pointer (we'll also use it as a base pointer).
-  // FIXME: use a dedicated register for it ?
-  Instruction *BasePointer =
-      IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
-  assert(BasePointer->getType() == StackPtrTy);
-
-  for (ReturnInst *RI : Returns) {
-    IRB.SetInsertPoint(RI);
-    IRB.CreateStore(BasePointer, UnsafeStackPtr);
-  }
-
   // Compute maximum alignment among static objects on the unsafe stack.
   unsigned MaxAlignment = 0;
   for (Argument *Arg : ByValArguments) {
@@ -726,9 +712,16 @@
   IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt());
   UnsafeStackPtr = getOrCreateUnsafeStackPtr(IRB, F);
 
+  // Load the current stack pointer (we'll also use it as a base pointer).
+  // FIXME: use a dedicated register for it ?
+  Instruction *BasePointer =
+    IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
+  assert(BasePointer->getType() == StackPtrTy);
+
   // The top of the unsafe stack after all unsafe static allocas are allocated.
   Value *StaticTop = moveStaticAllocasToUnsafeStack(IRB, F, StaticAllocas,
-                                                    ByValArguments, Returns);
+                                                    ByValArguments, Returns,
+                                                    BasePointer);
 
   // Safe stack object that stores the current unsafe stack top. It is updated
   // as unsafe dynamic (non-constant-sized) allocas are allocated and freed.
@@ -743,6 +736,12 @@
   moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop,
                                   DynamicAllocas);
 
+  // Restore the unsafe stack pointer before each return.
+  for (ReturnInst *RI : Returns) {
+    IRB.SetInsertPoint(RI);
+    IRB.CreateStore(BasePointer, UnsafeStackPtr);
+  }
+
   DEBUG(dbgs() << "[SafeStack]     safestack applied\n");
   return true;
 }
Index: llvm/trunk/test/Transforms/SafeStack/ARM/setjmp.ll
===================================================================
--- llvm/trunk/test/Transforms/SafeStack/ARM/setjmp.ll
+++ llvm/trunk/test/Transforms/SafeStack/ARM/setjmp.ll
@@ -6,8 +6,8 @@
 define void @f(i32 %b) safestack {
 entry:
 ; CHECK: %[[SPA:.*]] = call i8** @__safestack_pointer_address()
-; CHECK: %[[USDP:.*]] = alloca i8*
 ; CHECK: %[[USP:.*]] = load i8*, i8** %[[SPA]]
+; CHECK: %[[USDP:.*]] = alloca i8*
 ; CHECK: store i8* %[[USP]], i8** %[[USDP]]
 ; CHECK: call i32 @setjmp
 
@@ -26,6 +26,8 @@
   br label %if.end
 
 if.end:
+; CHECK: store i8* %[[USP:.*]], i8** %[[SPA:.*]]
+
   ret void
 }
 
Index: llvm/trunk/test/Transforms/SafeStack/dynamic-alloca.ll
===================================================================
--- llvm/trunk/test/Transforms/SafeStack/dynamic-alloca.ll
+++ llvm/trunk/test/Transforms/SafeStack/dynamic-alloca.ll
@@ -8,7 +8,7 @@
 ; Requires protector.
 define void @foo(i32 %n) nounwind uwtable safestack {
 entry:
-  ; CHECK: __safestack_unsafe_stack_ptr
+  ; CHECK: %[[SP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
   %n.addr = alloca i32, align 4
   %a = alloca i32*, align 8
   store i32 %n, i32* %n.addr, align 4
@@ -17,5 +17,6 @@
   %1 = alloca i8, i64 %conv
   %2 = bitcast i8* %1 to i32*
   store i32* %2, i32** %a, align 8
+  ; CHECK: store i8* %[[SP:.*]], i8** @__safestack_unsafe_stack_ptr
   ret void
 }
Index: llvm/trunk/test/Transforms/SafeStack/setjmp2.ll
===================================================================
--- llvm/trunk/test/Transforms/SafeStack/setjmp2.ll
+++ llvm/trunk/test/Transforms/SafeStack/setjmp2.ll
@@ -12,8 +12,8 @@
 ; CHECK: @foo(i32 %[[ARG:.*]])
 define i32 @foo(i32 %size) nounwind uwtable safestack {
 entry:
-  ; CHECK: %[[DYNPTR:.*]] = alloca i8*
-  ; CHECK-NEXT: %[[SP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+  ; CHECK: %[[SP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+  ; CHECK-NEXT: %[[DYNPTR:.*]] = alloca i8*
   ; CHECK-NEXT: store i8* %[[SP]], i8** %[[DYNPTR]]
 
   ; CHECK-NEXT: %[[ZEXT:.*]] = zext i32 %[[ARG]] to i64
@@ -35,6 +35,7 @@
 
   ; CHECK: call void @funcall(i32* %[[ALLOCA]])
   call void @funcall(i32* %a)
+  ; CHECK-NEXT: store i8* %[[SP:.*]], i8** @__safestack_unsafe_stack_ptr
   ret i32 0
 }