1 files changed, 65 insertions, 6 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyRegisterInfo.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyRegisterInfo.cpp
index 607fd6e79c2..683b52c58d4 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyRegisterInfo.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyRegisterInfo.cpp
@@ -51,6 +51,51 @@ WebAssemblyRegisterInfo::getReservedRegs(const MachineFunction & /*MF*/) const {
   return Reserved;
 }
 
+static bool isStackifiedVReg(const WebAssemblyFunctionInfo *WFI,
+                             const MachineOperand& Op) {
+  if (Op.isReg()) {
+    unsigned Reg = Op.getReg();
+    return TargetRegisterInfo::isVirtualRegister(Reg) &&
+        WFI->isVRegStackified(Reg);
+  }
+  return false;
+}
+
+static bool canStackifyOperand(const MachineInstr& Inst) {
+  unsigned Op = Inst.getOpcode();
+  return Op != TargetOpcode::PHI &&
+      Op != TargetOpcode::INLINEASM &&
+      Op != TargetOpcode::DBG_VALUE;
+}
+
+// Determine if the FI sequence can be stackified, and if so, where the code can
+// be inserted. If stackification is possible, returns true and ajusts II to
+// point to the insertion point.
+bool findInsertPt(const WebAssemblyFunctionInfo *WFI, MachineBasicBlock &MBB,
+                  unsigned OperandNum, MachineBasicBlock::iterator &II) {
+  if (!canStackifyOperand(*II)) return false;
+
+  MachineBasicBlock::iterator InsertPt(II);
+  int StackCount = 0;
+  // Operands are popped in reverse order, so any operands after FIOperand
+  // impose a constraint
+  for (unsigned i = OperandNum; i < II->getNumOperands(); i++) {
+    if (isStackifiedVReg(WFI, II->getOperand(i))) ++StackCount;
+  }
+  // Walk backwards, tracking stack depth. When it reaches 0 we have reached the
+  // top of the subtree.
+  while (StackCount) {
+    if (InsertPt == MBB.begin()) return false;
+    --InsertPt;
+    for (const auto &def : InsertPt->defs())
+      if (isStackifiedVReg(WFI, def)) --StackCount;
+    for (const auto &use : InsertPt->explicit_uses())
+      if (isStackifiedVReg(WFI, use)) ++StackCount;
+  }
+  II = InsertPt;
+  return true;
+}
+
 void WebAssemblyRegisterInfo::eliminateFrameIndex(
     MachineBasicBlock::iterator II, int SPAdj, unsigned FIOperandNum,
     RegScavenger * /*RS*/) const {
@@ -78,20 +123,34 @@ void WebAssemblyRegisterInfo::eliminateFrameIndex(
     MI.getOperand(FIOperandNum)
         .ChangeToRegister(WebAssembly::SP32, /*IsDef=*/false);
   } else {
-    // Otherwise create an i32.add SP, offset and make it the operand.
+    // Otherwise calculate the address
     auto &MRI = MF.getRegInfo();
     const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
 
     unsigned FIRegOperand = WebAssembly::SP32;
     if (FrameOffset) {
-      FIRegOperand = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
-      BuildMI(MBB, MI, MI.getDebugLoc(), TII->get(WebAssembly::CONST_I32),
-              FIRegOperand)
+      // Create i32.add SP, offset and make it the operand. We want to stackify
+      // this sequence, but we need to preserve the LIFO expr stack ordering
+      // (i.e. we can't insert our code in between MI and any operands it
+      // pops before FIOperand).
+      auto *WFI = MF.getInfo<WebAssemblyFunctionInfo>();
+      bool CanStackifyFI = findInsertPt(WFI, MBB, FIOperandNum, II);
+
+      unsigned OffsetOp = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
+      BuildMI(MBB, *II, II->getDebugLoc(), TII->get(WebAssembly::CONST_I32),
+              OffsetOp)
           .addImm(FrameOffset);
-      BuildMI(MBB, MI, MI.getDebugLoc(), TII->get(WebAssembly::ADD_I32),
+      if (CanStackifyFI) {
+        WFI->stackifyVReg(OffsetOp);
+        FIRegOperand = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
+        WFI->stackifyVReg(FIRegOperand);
+      } else {
+        FIRegOperand = OffsetOp;
+      }
+      BuildMI(MBB, *II, II->getDebugLoc(), TII->get(WebAssembly::ADD_I32),
               FIRegOperand)
           .addReg(WebAssembly::SP32)
-          .addReg(FIRegOperand);
+          .addReg(OffsetOp);
     }
     MI.getOperand(FIOperandNum).ChangeToRegister(FIRegOperand, /*IsDef=*/false);
   }