From 4f6481dc811250c55cda9f647b1e00331eec71cc Mon Sep 17 00:00:00 2001
From: Chandler Carruth <chandlerc@gmail.com>
Date: Wed, 25 Jul 2018 09:18:48 +0000
Subject: [PATCH] [x86/SLH] Sink the return hardening into the main block-walk
 + hardening code.

This consolidates all our hardening calls, and simplifies the code
a bit. It seems much more clear to handle all of these together.

No functionality changed here.

llvm-svn: 337895
---
 .../X86/X86SpeculativeLoadHardening.cpp       | 43 ++++++++-----------
 1 file changed, 17 insertions(+), 26 deletions(-)

diff --git a/llvm/lib/Target/X86/X86SpeculativeLoadHardening.cpp b/llvm/lib/Target/X86/X86SpeculativeLoadHardening.cpp
index 51a44dd97c45..be8cae945c1e 100644
--- a/llvm/lib/Target/X86/X86SpeculativeLoadHardening.cpp
+++ b/llvm/lib/Target/X86/X86SpeculativeLoadHardening.cpp
@@ -541,24 +541,6 @@ bool X86SpeculativeLoadHardeningPass::runOnMachineFunction(
       PS->SSA.RewriteUse(Op);
     }
 
-  // If we are hardening interprocedurally, find each returning block and
-  // protect the caller from being returned to through misspeculation.
-  if (HardenInterprocedurally)
-    for (MachineBasicBlock &MBB : MF) {
-      if (MBB.empty())
-        continue;
-
-      MachineInstr &MI = MBB.back();
-
-      // We only care about returns that are not also calls. For calls, that
-      // happen to also be returns (tail calls) we will have already handled
-      // them as calls.
-      if (!MI.isReturn() || MI.isCall())
-        continue;
-
-      hardenReturnInstr(MI);
-    }
-
   LLVM_DEBUG(dbgs() << "Final speculative load hardened function:\n"; MF.dump();
              dbgs() << "\n"; MF.verify(this));
   return true;
@@ -1588,16 +1570,24 @@ void X86SpeculativeLoadHardeningPass::tracePredStateThroughBlocksAndHarden(
           hardenIndirectCallOrJumpInstr(MI, AddrRegToHardenedReg);
       }
 
-      // After we finish processing the instruction and doing any hardening
-      // necessary for it, we need to handle transferring the predicate state
-      // into a call and recovering it after the call returns (if it returns).
-      if (!MI.isCall())
+      // After we finish hardening loads we handle interprocedural hardening if
+      // enabled and relevant for this instruction.
+      if (!HardenInterprocedurally)
+        continue;
+      if (!MI.isCall() && !MI.isReturn())
         continue;
 
-      // If we're not hardening interprocedurally, we can just skip calls.
-      if (!HardenInterprocedurally)
+      // If this is a direct return (IE, not a tail call) just directly harden
+      // it.
+      if (MI.isReturn() && !MI.isCall()) {
+        hardenReturnInstr(MI);
         continue;
+      }
 
+      // Otherwise we have a call. We need to handle transferring the predicate
+      // state into a call and recovering it after the call returns unless this
+      // is a tail call.
+      assert(MI.isCall() && "Should only reach here for calls!");
       auto InsertPt = MI.getIterator();
       DebugLoc Loc = MI.getDebugLoc();
 
@@ -1607,11 +1597,12 @@ void X86SpeculativeLoadHardeningPass::tracePredStateThroughBlocksAndHarden(
       unsigned StateReg = PS->SSA.GetValueAtEndOfBlock(&MBB);
       mergePredStateIntoSP(MBB, InsertPt, Loc, StateReg);
 
-      // If this call is also a return (because it is a tail call) we're done.
+      // If this call is also a return, it is a tail call and we don't need
+      // anything else to handle it so just continue.
       if (MI.isReturn())
         continue;
 
-      // Otherwise we need to step past the call and recover the predicate
+      // We need to step past the call and recover the predicate
       // state from SP after the return, and make this new state available.
       ++InsertPt;
       unsigned NewStateReg = extractPredStateFromSP(MBB, InsertPt, Loc);
-- 
2.34.1