[mips][ias] Fix O32 .cprestore directive when inside .set noat region and offset is in range.

[dsanders]

Summary

This expands on r269179 to fix an additional case that was not covered by our
tests. The assembler temporary is not needed when the .cprestore offset fits
inside a simm16 and it is not an error to use it inside a '.set noat' in this
case.

[sdardis]

LGTM.

lib/Target/Mips/AsmParser/MipsAsmParser.cpp

 void MipsAsmParser::expandStoreInst(MCInst &Inst, SMLoc IDLoc, MCStreamer &Out,
                                     const MCSubtargetInfo *STI,
                                     bool IsImmOpnd) {
   MipsTargetStreamer &TOut = getTargetStreamer();
 
   unsigned SrcReg = Inst.getOperand(0).getReg();
   unsigned BaseReg = Inst.getOperand(1).getReg();
 
-  unsigned ATReg = getATReg(IDLoc);
-  if (!ATReg)
-    return;
-
   if (IsImmOpnd) {
     TOut.emitStoreWithImmOffset(Inst.getOpcode(), SrcReg, BaseReg,
-                                Inst.getOperand(2).getImm(), ATReg, IDLoc, STI);
+                                Inst.getOperand(2).getImm(),
+                                [&]() { return getATReg(IDLoc); }, IDLoc, STI);
     return;
   }
 
+  unsigned ATReg = getATReg(IDLoc);
+  if (!ATReg)
+    return;
+
   const MCExpr *ExprOffset = Inst.getOperand(2).getExpr();
   MCOperand LoOperand = MCOperand::createExpr(
       MipsMCExpr::create(MipsMCExpr::MEK_LO, ExprOffset, getContext()));
   MCOperand HiOperand = MCOperand::createExpr(
       MipsMCExpr::create(MipsMCExpr::MEK_HI, ExprOffset, getContext()));
   TOut.emitStoreWithSymOffset(Inst.getOpcode(), SrcReg, BaseReg, HiOperand,
                               LoOperand, ATReg, IDLoc, STI);
 }
void MipsAsmParser::createCpRestoreMemOp(bool IsLoad, int StackOffset,
   MipsTargetStreamer &TOut = getTargetStreamer();
 
   if (IsLoad) {
     TOut.emitLoadWithImmOffset(Mips::LW, Mips::GP, Mips::SP, StackOffset,
                                Mips::GP, IDLoc, STI);
     return;
   }
 
-  unsigned ATReg = getATReg(IDLoc);
-  if (!ATReg)
-    return;
-
-  TOut.emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, StackOffset, ATReg,
-                              IDLoc, STI);
+  TOut.emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, StackOffset,
+                              [&]() { return getATReg(IDLoc); }, IDLoc, STI);
 }
 
 unsigned MipsAsmParser::checkTargetMatchPredicate(MCInst &Inst) {
   // As described by the Mips32r2 spec, the registers Rd and Rs for
   // jalr.hb must be different.
   // It also applies for registers Rt and Rs of microMIPSr6 jalrc.hb instruction
   // and registers Rd and Base for microMIPS lwp instruction
   unsigned Opcode = Inst.getOpcode();

lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp

void MipsTargetStreamer::emitGPRestore(int Offset, SMLoc IDLoc,
                                        const MCSubtargetInfo *STI) {
   emitLoadWithImmOffset(Mips::LW, Mips::GP, Mips::SP, Offset, Mips::GP, IDLoc,
                         STI);
 }
 
 /// Emit a store instruction with an immediate offset.
 void MipsTargetStreamer::emitStoreWithImmOffset(
     unsigned Opcode, unsigned SrcReg, unsigned BaseReg, int64_t Offset,
-    unsigned ATReg, SMLoc IDLoc, const MCSubtargetInfo *STI) {
+    std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
   if (isInt<16>(Offset)) {
     emitRRI(Opcode, SrcReg, BaseReg, Offset, IDLoc, STI);
     return;
   }
 
   // sw $8, offset($8) => lui $at, %hi(offset)
   //                      add $at, $at, $8
   //                      sw $8, %lo(offset)($at)
 
+  unsigned ATReg = GetATReg();
+  if (!ATReg)
+    return;
+
   unsigned LoOffset = Offset & 0x0000ffff;
   unsigned HiOffset = (Offset & 0xffff0000) >> 16;
 
   // If msb of LoOffset is 1(negative number) we must increment HiOffset
   // to account for the sign-extension of the low part.
   if (LoOffset & 0x8000)
     HiOffset++;
 
bool MipsTargetELFStreamer::emitDirectiveCpRestore(
   //    sw $gp, offset($sp)
   // and adds a corresponding LW after every JAL.
 
   // Note that .cprestore is ignored if used with the N32 and N64 ABIs or if it
   // is used in non-PIC mode.
   if (!Pic || (getABI().IsN32() || getABI().IsN64()))
     return true;
 
-  unsigned ATReg = GetATReg();
-  if (!ATReg)
-    return false;
-
   // Store the $gp on the stack.
-  emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, Offset, ATReg, IDLoc,
+  emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, Offset, GetATReg, IDLoc,
                          STI);
   return true;
 }
 
 void MipsTargetELFStreamer::emitDirectiveCpsetup(unsigned RegNo,
                                                  int RegOrOffset,
                                                  const MCSymbol &Sym,
                                                  bool IsReg) {

lib/Target/Mips/MipsTargetStreamer.h

void emitRRI(unsigned Opcode, unsigned Reg0, unsigned Reg1, int16_t Imm,
                SMLoc IDLoc, const MCSubtargetInfo *STI);
   void emitAddu(unsigned DstReg, unsigned SrcReg, unsigned TrgReg, bool Is64Bit,
                 const MCSubtargetInfo *STI);
   void emitDSLL(unsigned DstReg, unsigned SrcReg, int16_t ShiftAmount,
                 SMLoc IDLoc, const MCSubtargetInfo *STI);
   void emitEmptyDelaySlot(bool hasShortDelaySlot, SMLoc IDLoc,
                           const MCSubtargetInfo *STI);
   void emitNop(SMLoc IDLoc, const MCSubtargetInfo *STI);
+
+  /// Emit a store instruction with an offset. If the offset is out of range
+  /// then it will be synthesized using the assembler temporary.
+  ///
+  /// GetATReg() is a callback that can be used to obtain the current assembler
+  /// temporary and is only called when the assembler temporary is required. It
+  /// must handle the case where no assembler temporary is available (typically
+  /// by reporting an error).
   void emitStoreWithImmOffset(unsigned Opcode, unsigned SrcReg,
-                              unsigned BaseReg, int64_t Offset, unsigned ATReg,
-                              SMLoc IDLoc, const MCSubtargetInfo *STI);
+                              unsigned BaseReg, int64_t Offset,
+                              std::function<unsigned()> GetATReg, SMLoc IDLoc,
+                              const MCSubtargetInfo *STI);
   void emitStoreWithSymOffset(unsigned Opcode, unsigned SrcReg,
                               unsigned BaseReg, MCOperand &HiOperand,
                               MCOperand &LoOperand, unsigned ATReg, SMLoc IDLoc,
                               const MCSubtargetInfo *STI);
   void emitLoadWithImmOffset(unsigned Opcode, unsigned DstReg, unsigned BaseReg,
                              int64_t Offset, unsigned TmpReg, SMLoc IDLoc,
                              const MCSubtargetInfo *STI);
   void emitLoadWithSymOffset(unsigned Opcode, unsigned DstReg, unsigned BaseReg,
public:
   void emitDirectiveSetNoDsp() override;
   void emitDirectiveSetPop() override;
   void emitDirectiveSetPush() override;
   void emitDirectiveSetSoftFloat() override;
   void emitDirectiveSetHardFloat() override;
 
   // PIC support
   void emitDirectiveCpLoad(unsigned RegNo) override;
+
+  /// Emit a .cprestore directive.  If the offset is out of range then it will
+  /// be synthesized using the assembler temporary.
+  ///
+  /// GetATReg() is a callback that can be used to obtain the current assembler
+  /// temporary and is only called when the assembler temporary is required. It
+  /// must handle the case where no assembler temporary is available (typically
+  /// by reporting an error).
   bool emitDirectiveCpRestore(int Offset, std::function<unsigned()> GetATReg,
                               SMLoc IDLoc, const MCSubtargetInfo *STI) override;
   void emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
                             const MCSymbol &Sym, bool IsReg) override;
   void emitDirectiveCpreturn(unsigned SaveLocation,
                              bool SaveLocationIsRegister) override;
 
   // FP abiflags directives

test/MC/Mips/cprestore-noreorder-noat.s

# RUN:   FileCheck %s -check-prefix=NO-STORE
   .ent foo
 foo:
   .frame  $sp, 0, $ra
   .set noreorder
   .set noat
 
   .cpload $25
   .cprestore 8
-# O32: :[[@LINE-1]]:3: error: pseudo-instruction requires $at, which is not available
+# O32-NOT: error: pseudo-instruction requires $at, which is not available
 # N32-NOT: error: pseudo-instruction requires $at, which is not available
 # N64-NOT: error: pseudo-instruction requires $at, which is not available
 # NO-STORE-NOT: sw  $gp, 8($sp)
 
   jal $25
   jal $4, $25
   jal foo
 
   .end foo
+
+  .ent bar
+bar:
+  .frame  $sp, 0, $ra
+  .set noreorder
+  .set noat
+
+  .cpload $25
+  .cprestore 65536
+# O32: :[[@LINE-1]]:3: error: pseudo-instruction requires $at, which is not available
+# N32-NOT: error: pseudo-instruction requires $at, which is not available
+# N64-NOT: error: pseudo-instruction requires $at, which is not available
+# NO-STORE-NOT: sw $gp,
+
+  jal $25
+  jal $4, $25
+  jal bar
+
+  .end bar