Add some helper functions to the demangle utility classes.

These are all methods that, while not currently used in the
Itanium demangler, are generally useful enough that it's
likely the itanium demangler could find a use for them.  More
importantly, they are all necessary for the Microsoft demangler
which is up and coming in a subsequent patch.  Rather than
combine these into a single monolithic patch, I think it makes
sense to commit this utility code first since it is very simple,
this way it won't detract from the substance of the MS demangler
patch.

llvm-svn: 337316

4 files changed, 136 insertions, 22 deletions

diff --git a/llvm/include/llvm/Demangle/Demangle.h b/llvm/include/llvm/Demangle/Demangle.h
index 8203d539fdc..fde77d224f7 100644
--- a/llvm/include/llvm/Demangle/Demangle.h
+++ b/llvm/include/llvm/Demangle/Demangle.h
@@ -16,12 +16,14 @@ namespace llvm {
 /// The mangled_name is demangled into buf and returned. If the buffer is not
 /// large enough, realloc is used to expand it.
 ///
-/// The *status will be set to
-///   unknown_error: -4
-///   invalid_args:  -3
-///   invalid_mangled_name: -2
-///   memory_alloc_failure: -1
-///   success: 0
+/// The *status will be set to a value from the enumeration
+enum : int {
+  demangle_unknown_error = -4,
+  demangle_invalid_args = -3,
+  demangle_invalid_mangled_name = -2,
+  demangle_memory_alloc_failure = -1,
+  demangle_success = 0,
+};
 
 char *itaniumDemangle(const char *mangled_name, char *buf, size_t *n,
                       int *status);
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;