[Mangler] Verify USR manglings as well.

Author: jrose-apple

include/swift/Basic/Mangler.h

@@ -14,14 +14,13 @@
 #define SWIFT_BASIC_MANGLER_H
 
 #include "swift/Demangling/ManglingUtils.h"
+#include "swift/Basic/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/raw_ostream.h"
 
-using llvm::StringRef;
-using llvm::ArrayRef;
-
 namespace swift {
 namespace Mangle {
 
@@ -39,7 +38,7 @@ class Mangler {
   friend class SubstitutionMerging;
 
   /// The storage for the mangled symbol.
-  llvm::SmallVector<char, 128> Storage;
+  llvm::SmallString<128> Storage;
 
   /// The output stream for the mangled symbol.
   llvm::raw_svector_ostream Buffer;
@@ -106,7 +105,7 @@ class Mangler {
   void finalize(llvm::raw_ostream &stream);
 
   /// Verify that demangling and remangling works.
-  void verify(const std::string &mangledName);
+  static void verify(StringRef mangledName);
 
   /// Appends a mangled identifier string.
   void appendIdentifier(StringRef ident);

lib/AST/ASTMangler.cpp

@@ -350,6 +350,9 @@ std::string ASTMangler::mangleDeclAsUSR(const ValueDecl *Decl,
   } else {
     appendEntity(Decl);
   }
+
+  // We have a custom prefix, so finalize() won't verify for us. Do it manually.
+  verify(Storage.str().drop_front(USRPrefix.size()));
   return finalize();
 }
 
@@ -360,6 +363,8 @@ std::string ASTMangler::mangleAccessorEntityAsUSR(AccessorKind kind,
   llvm::SaveAndRestore<bool> allowUnnamedRAII(AllowNamelessEntities, true);
   Buffer << USRPrefix;
   appendAccessorEntity(kind, addressorKind, decl, /*isStatic*/ false);
+  // We have a custom prefix, so finalize() won't verify for us. Do it manually.
+  verify(Storage.str().drop_front(USRPrefix.size()));
   return finalize();
 }
 

lib/Basic/Mangler.cpp

@@ -122,9 +122,14 @@ void Mangler::beginMangling() {
 /// Finish the mangling of the symbol and return the mangled name.
 std::string Mangler::finalize() {
   assert(Storage.size() && "Mangling an empty name");
-  std::string result = std::string(Storage.data(), Storage.size());
+  std::string result = Storage.str().str();
   Storage.clear();
-  verify(result);
+
+#ifndef NDEBUG
+  if (StringRef(result).startswith(MANGLING_PREFIX_STR))
+    verify(result);
+#endif
+
   return result;
 }
 
@@ -147,11 +152,19 @@ static bool treeContains(Demangle::NodePointer Nd, Demangle::Node::Kind Kind) {
   return false;
 }
 
-void Mangler::verify(const std::string &mangledName) {
+void Mangler::verify(StringRef nameStr) {
 #ifndef NDEBUG
-  StringRef nameStr = mangledName;
-  if (!nameStr.startswith(MANGLING_PREFIX_STR))
-    return;
+  SmallString<128> buffer;
+  if (!nameStr.startswith(MANGLING_PREFIX_STR) &&
+      !nameStr.startswith("_Tt") &&
+      !nameStr.startswith("_S")) {
+    // This list is the set of prefixes recognized by Demangler::demangleSymbol.
+    // It should be kept in sync.
+    assert(StringRef(MANGLING_PREFIX_STR) != "_S" && "redundant check");
+    buffer += MANGLING_PREFIX_STR;
+    buffer += nameStr;
+    nameStr = buffer.str();
+  }
 
   Demangler Dem;
   NodePointer Root = Dem.demangleSymbol(nameStr);
@@ -160,15 +173,14 @@ void Mangler::verify(const std::string &mangledName) {
     abort();
   }
   std::string Remangled = mangleNode(Root);
-  if (Remangled == mangledName)
+  if (Remangled == nameStr)
     return;
 
   // There are cases (e.g. with dependent associated types) which results in
   // different remangled names. See ASTMangler::appendAssociatedTypeName.
   // This is no problem for the compiler, but we have to be more tolerant for
   // those cases. Instead we try to re-de-mangle the remangled name.
-  nameStr = Remangled;
-  NodePointer RootOfRemangled = Dem.demangleSymbol(nameStr);
+  NodePointer RootOfRemangled = Dem.demangleSymbol(Remangled);
   std::string ReDemangled = mangleNode(RootOfRemangled);
   if (Remangled == ReDemangled)
     return;

lib/Demangling/Demangler.cpp

@@ -237,6 +237,8 @@ NodePointer Demangler::demangleSymbol(StringRef MangledName) {
       && !nextIf("_S"))
     return nullptr;
 
+  // If any other prefixes are accepted, please update Mangler::verify.
+
   NodePointer topLevel = createNode(Node::Kind::Global);
 
   parseAndPushNodes();