[mips][ias] Fix N32 and N64 .cprestore directive when inside .set noat region.

Summary:
r268058 unintentionally made the retrieval of the current assembler temporary
unconditional. This was fine for the existing tests but it broke the cases
where the assembler temporary is not needed (N32/N64 or not PIC) and is
unavailable due to a '.set noat' directive.

This fixes FreeBSD's libc.

Reviewers: emaste, sdardis, seanbruno

Subscribers: dsanders, emaste, sdardis, llvm-commits

Differential Revision: http://reviews.llvm.org/D20093

llvm-svn: 269179

3 files changed, 28 insertions, 22 deletions

diff --git a/llvm/lib/Target/Mips/AsmParser/MipsAsmParser.cpp b/llvm/lib/Target/Mips/AsmParser/MipsAsmParser.cpp
index 2fa590190c4..f7fd72fd27f 100644
--- a/llvm/lib/Target/Mips/AsmParser/MipsAsmParser.cpp
+++ b/llvm/lib/Target/Mips/AsmParser/MipsAsmParser.cpp
@@ -5448,11 +5448,9 @@ bool MipsAsmParser::parseDirectiveCpRestore(SMLoc Loc) {
     return false;
   }
 
-  unsigned ATReg = getATReg(Loc);
-  if (!ATReg)
+  if (!getTargetStreamer().emitDirectiveCpRestore(
+          CpRestoreOffset, [&]() { return getATReg(Loc); }, Loc, STI))
     return true;
-
-  getTargetStreamer().emitDirectiveCpRestore(CpRestoreOffset, ATReg, Loc, STI);
   Parser.Lex(); // Consume the EndOfStatement.
   return false;
 }
diff --git a/llvm/lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp b/llvm/lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
index 23c46351b82..05c1026c25c 100644
--- a/llvm/lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
+++ b/llvm/lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
@@ -96,10 +96,11 @@ void MipsTargetStreamer::emitDirectiveSetHardFloat() {
 void MipsTargetStreamer::emitDirectiveSetDsp() { forbidModuleDirective(); }
 void MipsTargetStreamer::emitDirectiveSetNoDsp() { forbidModuleDirective(); }
 void MipsTargetStreamer::emitDirectiveCpLoad(unsigned RegNo) {}
-void MipsTargetStreamer::emitDirectiveCpRestore(int Offset, unsigned ATReg,
-                                                SMLoc IDLoc,
-                                                const MCSubtargetInfo *STI) {
+bool MipsTargetStreamer::emitDirectiveCpRestore(
+    int Offset, std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
   forbidModuleDirective();
+  return true;
 }
 void MipsTargetStreamer::emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
                                               const MCSymbol &Sym, bool IsReg) {
@@ -578,11 +579,12 @@ void MipsTargetAsmStreamer::emitDirectiveCpLoad(unsigned RegNo) {
   forbidModuleDirective();
 }
 
-void MipsTargetAsmStreamer::emitDirectiveCpRestore(int Offset, unsigned ATReg,
-                                                   SMLoc IDLoc,
-                                                   const MCSubtargetInfo *STI) {
-  MipsTargetStreamer::emitDirectiveCpRestore(Offset, ATReg, IDLoc, STI);
+bool MipsTargetAsmStreamer::emitDirectiveCpRestore(
+    int Offset, std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
+  MipsTargetStreamer::emitDirectiveCpRestore(Offset, GetATReg, IDLoc, STI);
   OS << "\t.cprestore\t" << Offset << "\n";
+  return true;
 }
 
 void MipsTargetAsmStreamer::emitDirectiveCpsetup(unsigned RegNo,
@@ -1034,10 +1036,10 @@ void MipsTargetELFStreamer::emitDirectiveCpLoad(unsigned RegNo) {
   forbidModuleDirective();
 }
 
-void MipsTargetELFStreamer::emitDirectiveCpRestore(int Offset, unsigned ATReg,
-                                                   SMLoc IDLoc,
-                                                   const MCSubtargetInfo *STI) {
-  MipsTargetStreamer::emitDirectiveCpRestore(Offset, ATReg, IDLoc, STI);
+bool MipsTargetELFStreamer::emitDirectiveCpRestore(
+    int Offset, std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
+  MipsTargetStreamer::emitDirectiveCpRestore(Offset, GetATReg, IDLoc, STI);
   // .cprestore offset
   // When PIC mode is enabled and the O32 ABI is used, this directive expands
   // to:
@@ -1047,11 +1049,16 @@ void MipsTargetELFStreamer::emitDirectiveCpRestore(int Offset, unsigned ATReg,
   // 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;
+    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,
                          STI);
+  return true;
 }
 
 void MipsTargetELFStreamer::emitDirectiveCpsetup(unsigned RegNo,
diff --git a/llvm/lib/Target/Mips/MipsTargetStreamer.h b/llvm/lib/Target/Mips/MipsTargetStreamer.h
index 2b39da6ade5..26f38829413 100644
--- a/llvm/lib/Target/Mips/MipsTargetStreamer.h
+++ b/llvm/lib/Target/Mips/MipsTargetStreamer.h
@@ -81,8 +81,9 @@ public:
 
   // PIC support
   virtual void emitDirectiveCpLoad(unsigned RegNo);
-  virtual void emitDirectiveCpRestore(int Offset, unsigned ATReg, SMLoc IDLoc,
-                                      const MCSubtargetInfo *STI);
+  virtual bool emitDirectiveCpRestore(int Offset,
+                                      std::function<unsigned()> GetATReg,
+                                      SMLoc IDLoc, const MCSubtargetInfo *STI);
   virtual void emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
                                     const MCSymbol &Sym, bool IsReg);
   virtual void emitDirectiveCpreturn(unsigned SaveLocation,
@@ -235,8 +236,8 @@ public:
 
   // PIC support
   void emitDirectiveCpLoad(unsigned RegNo) override;
-  void emitDirectiveCpRestore(int Offset, unsigned ATReg, SMLoc IDLoc,
-                              const MCSubtargetInfo *STI) override;
+  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,
@@ -290,8 +291,8 @@ public:
 
   // PIC support
   void emitDirectiveCpLoad(unsigned RegNo) override;
-  void emitDirectiveCpRestore(int Offset, unsigned ATReg, SMLoc IDLoc,
-                              const MCSubtargetInfo *STI) override;
+  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,