3 files changed, 128 insertions, 16 deletions

diff --git a/llvm/lib/Demangle/ItaniumDemangle.cpp b/llvm/lib/Demangle/ItaniumDemangle.cpp
index 31cac48dcf3..10ec84815e9 100644
--- a/llvm/lib/Demangle/ItaniumDemangle.cpp
+++ b/llvm/lib/Demangle/ItaniumDemangle.cpp
@@ -4925,32 +4925,24 @@ bool initializeOutputStream(char *Buf, size_t *N, OutputStream &S,
 
 }  // unnamed namespace
 
-enum {
-  unknown_error = -4,
-  invalid_args = -3,
-  invalid_mangled_name = -2,
-  memory_alloc_failure = -1,
-  success = 0,
-};
-
 char *llvm::itaniumDemangle(const char *MangledName, char *Buf,
                             size_t *N, int *Status) {
   if (MangledName == nullptr || (Buf != nullptr && N == nullptr)) {
     if (Status)
-      *Status = invalid_args;
+      *Status = demangle_invalid_args;
     return nullptr;
   }
 
-  int InternalStatus = success;
+  int InternalStatus = demangle_success;
   Db Parser(MangledName, MangledName + std::strlen(MangledName));
   OutputStream S;
 
   Node *AST = Parser.parse();
 
   if (AST == nullptr)
-    InternalStatus = invalid_mangled_name;
+    InternalStatus = demangle_invalid_mangled_name;
   else if (initializeOutputStream(Buf, N, S, 1024))
-    InternalStatus = memory_alloc_failure;
+    InternalStatus = demangle_memory_alloc_failure;
   else {
     assert(Parser.ForwardTemplateRefs.empty());
     AST->print(S);
@@ -4962,7 +4954,7 @@ char *llvm::itaniumDemangle(const char *MangledName, char *Buf,
 
   if (Status)
     *Status = InternalStatus;
-  return InternalStatus == success ? Buf : nullptr;
+  return InternalStatus == demangle_success ? Buf : nullptr;
 }
 
 namespace llvm {
diff --git a/llvm/lib/Demangle/StringView.h b/llvm/lib/Demangle/StringView.h
index 6485e61630f..ca27b41d33b 100644
--- a/llvm/lib/Demangle/StringView.h
+++ b/llvm/lib/Demangle/StringView.h
@@ -14,6 +14,7 @@
 #define LLVM_DEMANGLE_STRINGVIEW_H
 
 #include <algorithm>
+#include <cassert>
 
 class StringView {
   const char *First;
@@ -24,9 +25,16 @@ public:
   StringView(const char (&Str)[N]) : First(Str), Last(Str + N - 1) {}
   StringView(const char *First_, const char *Last_)
       : First(First_), Last(Last_) {}
+  StringView(const char *First_, size_t Len)
+      : First(First_), Last(First_ + Len) {}
+  StringView(const char *Str) : First(Str), Last(Str + strlen(Str)) {}
   StringView() : First(nullptr), Last(nullptr) {}
 
-  StringView substr(size_t From, size_t To) {
+  StringView substr(size_t From) const {
+    return StringView(begin() + From, size() - From);
+  }
+
+  StringView substr(size_t From, size_t To) const {
     if (To >= size())
       To = size() - 1;
     if (From >= size())
@@ -34,12 +42,38 @@ public:
     return StringView(First + From, First + To);
   }
 
-  StringView dropFront(size_t N) const {
+  StringView dropFront(size_t N = 1) const {
     if (N >= size())
       N = size() - 1;
     return StringView(First + N, Last);
   }
 
+  char front() const {
+    assert(!empty());
+    return *begin();
+  }
+
+  char popFront() {
+    assert(!empty());
+    return *First++;
+  }
+
+  bool consumeFront(char C) {
+    if (!startsWith(C))
+      return false;
+    *this = dropFront(1);
+    return true;
+  }
+
+  bool consumeFront(StringView S) {
+    if (!startsWith(S))
+      return false;
+    *this = dropFront(S.size());
+    return true;
+  }
+
+  bool startsWith(char C) const { return !empty() && *begin() == C; }
+
   bool startsWith(StringView Str) const {
     if (Str.size() > size())
       return false;
diff --git a/llvm/lib/Demangle/Utility.h b/llvm/lib/Demangle/Utility.h
index 62c0e7430a5..f764663436d 100644
--- a/llvm/lib/Demangle/Utility.h
+++ b/llvm/lib/Demangle/Utility.h
@@ -12,6 +12,8 @@
 #ifndef LLVM_DEMANGLE_UTILITY_H
 #define LLVM_DEMANGLE_UTILITY_H
 
+#include "StringView.h"
+
 #include <cstdlib>
 #include <cstring>
 #include <limits>
@@ -33,6 +35,27 @@ class OutputStream {
     }
   }
 
+  void writeUnsigned(uint64_t N, bool isNeg = false) {
+    // Handle special case...
+    if (N == 0) {
+      *this << '0';
+      return;
+    }
+
+    char Temp[21];
+    char *TempPtr = std::end(Temp);
+
+    while (N) {
+      *--TempPtr = '0' + char(N % 10);
+      N /= 10;
+    }
+
+    // Add negative sign...
+    if (isNeg)
+      *--TempPtr = '-';
+    this->operator<<(StringView(TempPtr, std::end(Temp)));
+  }
+
 public:
   OutputStream(char *StartBuf, size_t Size)
       : Buffer(StartBuf), CurrentPosition(0), BufferCapacity(Size) {}
@@ -43,6 +66,20 @@ public:
     BufferCapacity = BufferCapacity_;
   }
 
+  /// Create an OutputStream from a buffer and a size.  If either of these are
+  /// null a buffer is allocated.
+  static OutputStream create(char *StartBuf, size_t *Size, size_t AllocSize) {
+    OutputStream Result;
+
+    if (!StartBuf || !Size) {
+      StartBuf = static_cast<char *>(std::malloc(AllocSize));
+      Size = &AllocSize;
+    }
+
+    Result.reset(StartBuf, *Size);
+    return Result;
+  }
+
   /// If a ParameterPackExpansion (or similar type) is encountered, the offset
   /// into the pack that we're currently printing.
   unsigned CurrentPackIndex = std::numeric_limits<unsigned>::max();
@@ -64,6 +101,39 @@ public:
     return *this;
   }
 
+  OutputStream &operator<<(StringView R) { return (*this += R); }
+
+  OutputStream &operator<<(char C) { return (*this += C); }
+
+  OutputStream &operator<<(long long N) {
+    if (N < 0)
+      writeUnsigned(static_cast<unsigned long long>(-N), true);
+    else
+      writeUnsigned(static_cast<unsigned long long>(N));
+    return *this;
+  }
+
+  OutputStream &operator<<(unsigned long long N) {
+    writeUnsigned(N, false);
+    return *this;
+  }
+
+  OutputStream &operator<<(long N) {
+    return this->operator<<(static_cast<long long>(N));
+  }
+
+  OutputStream &operator<<(unsigned long N) {
+    return this->operator<<(static_cast<unsigned long long>(N));
+  }
+
+  OutputStream &operator<<(int N) {
+    return this->operator<<(static_cast<long long>(N));
+  }
+
+  OutputStream &operator<<(unsigned int N) {
+    return this->operator<<(static_cast<unsigned long long>(N));
+  }
+
   size_t getCurrentPosition() const { return CurrentPosition; }
   void setCurrentPosition(size_t NewPos) { CurrentPosition = NewPos; }
 
@@ -81,13 +151,29 @@ public:
 template <class T> class SwapAndRestore {
   T &Restore;
   T OriginalValue;
+  bool ShouldRestore = true;
 
 public:
+  SwapAndRestore(T &Restore_) : SwapAndRestore(Restore_, Restore_) {}
+
   SwapAndRestore(T &Restore_, T NewVal)
       : Restore(Restore_), OriginalValue(Restore) {
     Restore = std::move(NewVal);
   }
-  ~SwapAndRestore() { Restore = std::move(OriginalValue); }
+  ~SwapAndRestore() {
+    if (ShouldRestore)
+      Restore = std::move(OriginalValue);
+  }
+
+  void shouldRestore(bool ShouldRestore_) { ShouldRestore = ShouldRestore_; }
+
+  void restoreNow(bool Force) {
+    if (!Force && !ShouldRestore)
+      return;
+
+    Restore = std::move(OriginalValue);
+    ShouldRestore = false;
+  }
 
   SwapAndRestore(const SwapAndRestore &) = delete;
   SwapAndRestore &operator=(const SwapAndRestore &) = delete;