Use a raw_stream to implement the mangler.

This is a bit more convenient for some callers, but more importantly, it is
easier to implement correctly. Doing this removes the patching of already
printed data that was used for fastcall, fixing a crash with private fastcall
symbols.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@200367 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 89 insertions, 67 deletions

diff --git a/include/llvm/IR/Mangler.h b/include/llvm/IR/Mangler.h
index d4a9559819..9229dd8446 100644
--- a/include/llvm/IR/Mangler.h
+++ b/include/llvm/IR/Mangler.h
@@ -15,6 +15,7 @@
 #define LLVM_TARGET_MANGLER_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
 
@@ -47,14 +48,16 @@ private:
 public:
   Mangler(const DataLayout *DL) : DL(DL), NextAnonGlobalID(1) {}
 
-  /// getNameWithPrefix - Fill OutName with the name of the appropriate prefix
-  /// and the specified global variable's name.  If the global variable doesn't
-  /// have a name, this fills in a unique name for the global.
+  /// Print the appropriate prefix and the specified global variable's name.
+  /// If the global variable doesn't have a name, this fills in a unique name
+  /// for the global.
+  void getNameWithPrefix(raw_ostream &OS, const GlobalValue *GV);
   void getNameWithPrefix(SmallVectorImpl<char> &OutName, const GlobalValue *GV);
 
-  /// getNameWithPrefix - Fill OutName with the name of the appropriate prefix
-  /// and the specified name as the global variable name.  GVName must not be
-  /// empty.
+  /// Print the appropriate prefix and the specified name as the global variable
+  /// name. GVName must not be empty.
+  void getNameWithPrefix(raw_ostream &OS, const Twine &GVName,
+                         ManglerPrefixTy PrefixTy = Mangler::Default);
   void getNameWithPrefix(SmallVectorImpl<char> &OutName, const Twine &GVName,
                          ManglerPrefixTy PrefixTy = Mangler::Default);
 };
diff --git a/lib/IR/Mangler.cpp b/lib/IR/Mangler.cpp
index c067f0f392..5a099738d5 100644
--- a/lib/IR/Mangler.cpp
+++ b/lib/IR/Mangler.cpp
@@ -20,41 +20,46 @@
 #include "llvm/Support/raw_ostream.h"
 using namespace llvm;
 
-/// getNameWithPrefix - Fill OutName with the name of the appropriate prefix
-/// and the specified name as the global variable name.  GVName must not be
-/// empty.
-void Mangler::getNameWithPrefix(SmallVectorImpl<char> &OutName,
-                                const Twine &GVName, ManglerPrefixTy PrefixTy) {
+static void getNameWithPrefixx(raw_ostream &OS, const Twine &GVName,
+                              Mangler::ManglerPrefixTy PrefixTy,
+                              const DataLayout &DL, bool UseAt) {
   SmallString<256> TmpData;
   StringRef Name = GVName.toStringRef(TmpData);
   assert(!Name.empty() && "getNameWithPrefix requires non-empty name");
 
-  // If the global name is not led with \1, add the appropriate prefixes.
-  if (Name[0] == '\1') {
-    Name = Name.substr(1);
-  } else {
-    if (PrefixTy == Mangler::Private) {
-      const char *Prefix = DL->getPrivateGlobalPrefix();
-      OutName.append(Prefix, Prefix+strlen(Prefix));
-    } else if (PrefixTy == Mangler::LinkerPrivate) {
-      const char *Prefix = DL->getLinkerPrivateGlobalPrefix();
-      OutName.append(Prefix, Prefix+strlen(Prefix));
-    }
+  if (PrefixTy == Mangler::Private)
+    OS << DL.getPrivateGlobalPrefix();
+  else if (PrefixTy == Mangler::LinkerPrivate)
+    OS << DL.getLinkerPrivateGlobalPrefix();
 
-    char Prefix = DL->getGlobalPrefix();
+  if (UseAt) {
+    OS << '@';
+  } else {
+    char Prefix = DL.getGlobalPrefix();
     if (Prefix != '\0')
-      OutName.push_back(Prefix);
+      OS << Prefix;
   }
 
   // If this is a simple string that doesn't need escaping, just append it.
-  OutName.append(Name.begin(), Name.end());
+  OS << Name;
+}
+
+void Mangler::getNameWithPrefix(raw_ostream &OS,
+                                const Twine &GVName, ManglerPrefixTy PrefixTy) {
+  return getNameWithPrefixx(OS, GVName, PrefixTy, *DL, false);
+}
+
+void Mangler::getNameWithPrefix(SmallVectorImpl<char> &OutName,
+                                const Twine &GVName, ManglerPrefixTy PrefixTy) {
+  raw_svector_ostream OS(OutName);
+  return getNameWithPrefix(OS, GVName, PrefixTy);
 }
 
 /// AddFastCallStdCallSuffix - Microsoft fastcall and stdcall functions require
 /// a suffix on their name indicating the number of words of arguments they
 /// take.
-static void AddFastCallStdCallSuffix(SmallVectorImpl<char> &OutName,
-                                     const Function *F, const DataLayout &TD) {
+static void AddFastCallStdCallSuffix(raw_ostream &OS, const Function *F,
+                                     const DataLayout &TD) {
   // Calculate arguments size total.
   unsigned ArgWords = 0;
   for (Function::const_arg_iterator AI = F->arg_begin(), AE = F->arg_end();
@@ -66,62 +71,69 @@ static void AddFastCallStdCallSuffix(SmallVectorImpl<char> &OutName,
     // Size should be aligned to DWORD boundary
     ArgWords += ((TD.getTypeAllocSize(Ty) + 3)/4)*4;
   }
-  
-  raw_svector_ostream(OutName) << '@' << ArgWords;
-}
 
+  OS << '@' << ArgWords;
+}
 
-/// getNameWithPrefix - Fill OutName with the name of the appropriate prefix
-/// and the specified global variable's name.  If the global variable doesn't
-/// have a name, this fills in a unique name for the global.
-void Mangler::getNameWithPrefix(SmallVectorImpl<char> &OutName,
-                                const GlobalValue *GV) {
+void Mangler::getNameWithPrefix(raw_ostream &OS, const GlobalValue *GV) {
   ManglerPrefixTy PrefixTy = Mangler::Default;
   if (GV->hasPrivateLinkage())
     PrefixTy = Mangler::Private;
   else if (GV->hasLinkerPrivateLinkage() || GV->hasLinkerPrivateWeakLinkage())
     PrefixTy = Mangler::LinkerPrivate;
 
-  size_t NameBegin = OutName.size();
-
-  // If this global has a name, handle it simply.
-  if (GV->hasName()) {
-    StringRef Name = GV->getName();
-    getNameWithPrefix(OutName, Name, PrefixTy);
-    // No need to do anything else if the global has the special "do not mangle"
-    // flag in the name.
-    if (Name[0] == 1)
-      return;
-  } else {
+  if (!GV->hasName()) {
     // Get the ID for the global, assigning a new one if we haven't got one
     // already.
     unsigned &ID = AnonGlobalIDs[GV];
-    if (ID == 0) ID = NextAnonGlobalID++;
-  
+    if (ID == 0)
+      ID = NextAnonGlobalID++;
+
     // Must mangle the global into a unique ID.
-    getNameWithPrefix(OutName, "__unnamed_" + Twine(ID), PrefixTy);
+    getNameWithPrefix(OS, "__unnamed_" + Twine(ID), PrefixTy);
+    return;
   }
 
-  // If we are supposed to add a microsoft-style suffix for stdcall/fastcall,
-  // add it.
+  StringRef Name = GV->getName();
+
+  // No need to do anything special if the global has the special "do not
+  // mangle" flag in the name.
+  if (Name[0] == '\1') {
+    OS << Name.substr(1);
+    return;
+  }
+
+  bool UseAt = false;
+  const Function *MSFunc = NULL;
+  CallingConv::ID CC;
   if (DL->hasMicrosoftFastStdCallMangling()) {
-    if (const Function *F = dyn_cast<Function>(GV)) {
-      CallingConv::ID CC = F->getCallingConv();
-    
+    if ((MSFunc = dyn_cast<Function>(GV))) {
+      CC = MSFunc->getCallingConv();
       // fastcall functions need to start with @ instead of _.
-      if (CC == CallingConv::X86_FastCall) {
-        assert(OutName[NameBegin] == '_' && DL->getGlobalPrefix() == '_');
-        OutName[NameBegin] = '@';
-      }
-    
-      // fastcall and stdcall functions usually need @42 at the end to specify
-      // the argument info.
-      FunctionType *FT = F->getFunctionType();
-      if ((CC == CallingConv::X86_FastCall || CC == CallingConv::X86_StdCall) &&
-          // "Pure" variadic functions do not receive @0 suffix.
-          (!FT->isVarArg() || FT->getNumParams() == 0 ||
-           (FT->getNumParams() == 1 && F->hasStructRetAttr())))
-        AddFastCallStdCallSuffix(OutName, F, *DL);
+      if (CC == CallingConv::X86_FastCall)
+        UseAt = true;
     }
   }
+
+  getNameWithPrefixx(OS, Name, PrefixTy, *DL, UseAt);
+
+  if (!MSFunc)
+    return;
+
+  // If we are supposed to add a microsoft-style suffix for stdcall/fastcall,
+  // add it.
+  // fastcall and stdcall functions usually need @42 at the end to specify
+  // the argument info.
+  FunctionType *FT = MSFunc->getFunctionType();
+  if ((CC == CallingConv::X86_FastCall || CC == CallingConv::X86_StdCall) &&
+      // "Pure" variadic functions do not receive @0 suffix.
+      (!FT->isVarArg() || FT->getNumParams() == 0 ||
+       (FT->getNumParams() == 1 && MSFunc->hasStructRetAttr())))
+    AddFastCallStdCallSuffix(OS, MSFunc, *DL);
+}
+
+void Mangler::getNameWithPrefix(SmallVectorImpl<char> &OutName,
+                                const GlobalValue *GV) {
+  raw_svector_ostream OS(OutName);
+  getNameWithPrefix(OS, GV);
 }
diff --git a/test/CodeGen/X86/fastcall-correct-mangling.ll b/test/CodeGen/X86/fastcall-correct-mangling.ll
index 15cffa4fbd..00dc44e75e 100644
--- a/test/CodeGen/X86/fastcall-correct-mangling.ll
+++ b/test/CodeGen/X86/fastcall-correct-mangling.ll
@@ -24,3 +24,10 @@ define x86_fastcallcc i32 @"\01DoNotMangle"(i32 %a) {
 entry:
   ret i32 %a
 }
+
+define private x86_fastcallcc void @dontCrash() {
+; The name is fairly arbitrary since it is private. Just don't crash.
+; CHECK32-LABEL: {{^}}L@dontCrash@0:
+; CHECK64-LABEL: {{^}}.LdontCrash:
+  ret void
+}