Allow demangler's node allocator to fail, and bail out of the entire

demangling process when it does.

Use this to support a "lookup" query for the mangling canonicalizer that
does not create new nodes. This could also be used to implement
demangling with a fixed-size temporary storage buffer.

Reviewers: erik.pilkington

Subscribers: llvm-commits

Differential Revision: https://reviews.llvm.org/D51003

llvm-svn: 340670

4 files changed, 109 insertions, 28 deletions

diff --git a/llvm/include/llvm/Demangle/ItaniumDemangle.h b/llvm/include/llvm/Demangle/ItaniumDemangle.h
index 82233abba0c..81d8508ccce 100644
--- a/llvm/include/llvm/Demangle/ItaniumDemangle.h
+++ b/llvm/include/llvm/Demangle/ItaniumDemangle.h
@@ -2372,7 +2372,10 @@ template<typename Alloc> Node *Db<Alloc>::parseLocalName(NameState *State) {
 
   if (consumeIf('s')) {
     First = parse_discriminator(First, Last);
-    return make<LocalName>(Encoding, make<NameType>("string literal"));
+    auto *StringLitName = make<NameType>("string literal");
+    if (!StringLitName)
+      return nullptr;
+    return make<LocalName>(Encoding, StringLitName);
   }
 
   if (consumeIf('d')) {
@@ -2786,6 +2789,8 @@ Node *Db<Alloc>::parseCtorDtorName(Node *&SoFar, NameState *State) {
     case SpecialSubKind::ostream:
     case SpecialSubKind::iostream:
       SoFar = make<ExpandedSpecialSubstitution>(SSK);
+      if (!SoFar)
+        return nullptr;
     default:
       break;
     }
@@ -2847,13 +2852,18 @@ template<typename Alloc> Node *Db<Alloc>::parseNestedName(NameState *State) {
 
   Node *SoFar = nullptr;
   auto PushComponent = [&](Node *Comp) {
+    if (!Comp) return false;
     if (SoFar) SoFar = make<NestedName>(SoFar, Comp);
     else       SoFar = Comp;
     if (State) State->EndsWithTemplateArgs = false;
+    return SoFar != nullptr;
   };
 
-  if (consumeIf("St"))
+  if (consumeIf("St")) {
     SoFar = make<NameType>("std");
+    if (!SoFar)
+      return nullptr;
+  }
 
   while (!consumeIf('E')) {
     consumeIf('L'); // extension
@@ -2867,10 +2877,8 @@ template<typename Alloc> Node *Db<Alloc>::parseNestedName(NameState *State) {
 
     //          ::= <template-param>
     if (look() == 'T') {
-      Node *TP = parseTemplateParam();
-      if (TP == nullptr)
+      if (!PushComponent(parseTemplateParam()))
         return nullptr;
-      PushComponent(TP);
       Subs.push_back(SoFar);
       continue;
     }
@@ -2881,6 +2889,8 @@ template<typename Alloc> Node *Db<Alloc>::parseNestedName(NameState *State) {
       if (TA == nullptr || SoFar == nullptr)
         return nullptr;
       SoFar = make<NameWithTemplateArgs>(SoFar, TA);
+      if (!SoFar)
+        return nullptr;
       if (State) State->EndsWithTemplateArgs = true;
       Subs.push_back(SoFar);
       continue;
@@ -2888,10 +2898,8 @@ template<typename Alloc> Node *Db<Alloc>::parseNestedName(NameState *State) {
 
     //          ::= <decltype>
     if (look() == 'D' && (look(1) == 't' || look(1) == 'T')) {
-      Node *DT = parseDecltype();
-      if (DT == nullptr)
+      if (!PushComponent(parseDecltype()))
         return nullptr;
-      PushComponent(DT);
       Subs.push_back(SoFar);
       continue;
     }
@@ -2899,9 +2907,8 @@ template<typename Alloc> Node *Db<Alloc>::parseNestedName(NameState *State) {
     //          ::= <substitution>
     if (look() == 'S' && look(1) != 't') {
       Node *S = parseSubstitution();
-      if (S == nullptr)
+      if (!PushComponent(S))
         return nullptr;
-      PushComponent(S);
       if (SoFar != S)
         Subs.push_back(S);
       continue;
@@ -2911,10 +2918,8 @@ template<typename Alloc> Node *Db<Alloc>::parseNestedName(NameState *State) {
     if (look() == 'C' || (look() == 'D' && look(1) != 'C')) {
       if (SoFar == nullptr)
         return nullptr;
-      Node *CtorDtor = parseCtorDtorName(SoFar, State);
-      if (CtorDtor == nullptr)
+      if (!PushComponent(parseCtorDtorName(SoFar, State)))
         return nullptr;
-      PushComponent(CtorDtor);
       SoFar = parseAbiTags(SoFar);
       if (SoFar == nullptr)
         return nullptr;
@@ -2923,10 +2928,8 @@ template<typename Alloc> Node *Db<Alloc>::parseNestedName(NameState *State) {
     }
 
     //          ::= <prefix> <unqualified-name>
-    Node *N = parseUnqualifiedName(State);
-    if (N == nullptr)
+    if (!PushComponent(parseUnqualifiedName(State)))
       return nullptr;
-    PushComponent(N);
     Subs.push_back(SoFar);
   }
 
@@ -3039,6 +3042,8 @@ template<typename Alloc> Node *Db<Alloc>::parseUnresolvedName() {
       if (TA == nullptr)
         return nullptr;
       SoFar = make<NameWithTemplateArgs>(SoFar, TA);
+      if (!SoFar)
+        return nullptr;
     }
 
     while (!consumeIf('E')) {
@@ -3046,6 +3051,8 @@ template<typename Alloc> Node *Db<Alloc>::parseUnresolvedName() {
       if (Qual == nullptr)
         return nullptr;
       SoFar = make<QualifiedName>(SoFar, Qual);
+      if (!SoFar)
+        return nullptr;
     }
 
     Node *Base = parseBaseUnresolvedName();
@@ -3078,6 +3085,8 @@ template<typename Alloc> Node *Db<Alloc>::parseUnresolvedName() {
         SoFar = make<GlobalQualifiedName>(Qual);
       else
         SoFar = Qual;
+      if (!SoFar)
+        return nullptr;
     } while (!consumeIf('E'));
   }
   //      sr <unresolved-type>                 <base-unresolved-name>
@@ -3092,6 +3101,8 @@ template<typename Alloc> Node *Db<Alloc>::parseUnresolvedName() {
       if (TA == nullptr)
         return nullptr;
       SoFar = make<NameWithTemplateArgs>(SoFar, TA);
+      if (!SoFar)
+        return nullptr;
     }
   }
 
@@ -3111,6 +3122,8 @@ template<typename Alloc> Node *Db<Alloc>::parseAbiTags(Node *N) {
     if (SN.empty())
       return nullptr;
     N = make<AbiTagAttr>(N, SN);
+    if (!N)
+      return nullptr;
   }
   return N;
 }
@@ -3163,11 +3176,15 @@ template<typename Alloc> Node *Db<Alloc>::parseFunctionType() {
   Node *ExceptionSpec = nullptr;
   if (consumeIf("Do")) {
     ExceptionSpec = make<NameType>("noexcept");
+    if (!ExceptionSpec)
+      return nullptr;
   } else if (consumeIf("DO")) {
     Node *E = parseExpr();
     if (E == nullptr || !consumeIf('E'))
       return nullptr;
     ExceptionSpec = make<NoexceptSpec>(E);
+    if (!ExceptionSpec)
+      return nullptr;
   } else if (consumeIf("Dw")) {
     size_t SpecsBegin = Names.size();
     while (!consumeIf('E')) {
@@ -3178,6 +3195,8 @@ template<typename Alloc> Node *Db<Alloc>::parseFunctionType() {
     }
     ExceptionSpec =
       make<DynamicExceptionSpec>(popTrailingNodeArray(SpecsBegin));
+    if (!ExceptionSpec)
+      return nullptr;
   }
 
   consumeIf("Dx"); // transaction safe
@@ -4517,9 +4536,10 @@ template<typename Alloc> Node *Db<Alloc>::parseExpr() {
           return nullptr;
         Names.push_back(Arg);
       }
-      return make<EnclosingExpr>(
-          "sizeof... (", make<NodeArrayNode>(popTrailingNodeArray(ArgsBegin)),
-          ")");
+      auto *Pack = make<NodeArrayNode>(popTrailingNodeArray(ArgsBegin));
+      if (!Pack)
+        return nullptr;
+      return make<EnclosingExpr>("sizeof... (", Pack, ")");
     }
     }
     return nullptr;
@@ -4771,6 +4791,8 @@ template<typename Alloc> Node *Db<Alloc>::parseEncoding() {
       Names.push_back(Arg);
     }
     Attrs = make<EnableIfAttr>(popTrailingNodeArray(BeforeArgs));
+    if (!Attrs)
+      return nullptr;
   }
 
   Node *ReturnType = nullptr;
@@ -4915,6 +4937,8 @@ template<typename Alloc> Node *Db<Alloc>::parseSubstitution() {
     default:
       return nullptr;
     }
+    if (!SpecialSub)
+      return nullptr;
     // Itanium C++ ABI 5.1.2: If a name that would use a built-in <substitution>
     // has ABI tags, the tags are appended to the substitution; the result is a
     // substitutable component.
@@ -4968,6 +4992,8 @@ template<typename Alloc> Node *Db<Alloc>::parseTemplateParam() {
   // operator types), then we should only look it up in the right context.
   if (PermitForwardTemplateReferences) {
     Node *ForwardRef = make<ForwardTemplateReference>(Index);
+    if (!ForwardRef)
+      return nullptr;
     assert(ForwardRef->getKind() == Node::KForwardTemplateReference);
     ForwardTemplateRefs.push_back(
         static_cast<ForwardTemplateReference *>(ForwardRef));
@@ -5047,6 +5073,8 @@ Node *Db<Alloc>::parseTemplateArgs(bool TagTemplates) {
       if (Arg->getKind() == Node::KTemplateArgumentPack) {
         TableEntry = make<ParameterPack>(
             static_cast<TemplateArgumentPack*>(TableEntry)->getElements());
+        if (!TableEntry)
+          return nullptr;
       }
       TemplateParams.push_back(TableEntry);
     } else {
diff --git a/llvm/include/llvm/Support/ItaniumManglingCanonicalizer.h b/llvm/include/llvm/Support/ItaniumManglingCanonicalizer.h
index 81d57eb9b93..34eb9f7deaa 100644
--- a/llvm/include/llvm/Support/ItaniumManglingCanonicalizer.h
+++ b/llvm/include/llvm/Support/ItaniumManglingCanonicalizer.h
@@ -76,8 +76,14 @@ public:
   ///
   /// Returns Key() if (and only if) the mangling is not a valid Itanium C++
   /// ABI mangling.
+  ///
+  /// The string denoted by Mangling must live as long as the canonicalizer.
   Key canonicalize(StringRef Mangling);
 
+  /// Find a canonical key for the specified mangling, if one has already been
+  /// formed. Otherwise returns Key().
+  Key lookup(StringRef Mangling);
+
 private:
   struct Impl;
   Impl *P;
diff --git a/llvm/lib/Support/ItaniumManglingCanonicalizer.cpp b/llvm/lib/Support/ItaniumManglingCanonicalizer.cpp
index 9d78ddc0b8b..f7b8b98ed16 100644
--- a/llvm/lib/Support/ItaniumManglingCanonicalizer.cpp
+++ b/llvm/lib/Support/ItaniumManglingCanonicalizer.cpp
@@ -105,7 +105,7 @@ public:
   void reset() {}
 
   template<typename T, typename ...Args>
-  std::pair<Node*, bool> getOrCreateNode(Args &&...As) {
+  std::pair<Node*, bool> getOrCreateNode(bool CreateNewNodes, Args &&...As) {
     llvm::FoldingSetNodeID ID;
     profileCtor(ID, NodeKind<T>::Kind, As...);
 
@@ -113,6 +113,9 @@ public:
     if (NodeHeader *Existing = Nodes.FindNodeOrInsertPos(ID, InsertPos))
       return {static_cast<T*>(Existing->getNode()), false};
 
+    if (!CreateNewNodes)
+      return {nullptr, true};
+
     static_assert(alignof(T) <= alignof(NodeHeader),
                   "underaligned node header for specific node kind");
     void *Storage =
@@ -125,7 +128,7 @@ public:
 
   template<typename T, typename... Args>
   Node *makeNode(Args &&...As) {
-    return getOrCreateNode<T>(std::forward<Args>(As)...).first;
+    return getOrCreateNode<T>(true, std::forward<Args>(As)...).first;
   }
 
   void *allocateNodeArray(size_t sz) {
@@ -138,7 +141,8 @@ public:
 // of creation.
 template<>
 std::pair<Node *, bool>
-FoldingNodeAllocator::getOrCreateNode<ForwardTemplateReference>(size_t &Index) {
+FoldingNodeAllocator::getOrCreateNode<ForwardTemplateReference>(bool,
+                                                                size_t &Index) {
   return {new (RawAlloc.Allocate(sizeof(ForwardTemplateReference),
                                  alignof(ForwardTemplateReference)))
               ForwardTemplateReference(Index),
@@ -149,15 +153,16 @@ class CanonicalizerAllocator : public FoldingNodeAllocator {
   Node *MostRecentlyCreated = nullptr;
   Node *TrackedNode = nullptr;
   bool TrackedNodeIsUsed = false;
+  bool CreateNewNodes = true;
   llvm::SmallDenseMap<Node*, Node*, 32> Remappings;
 
   template<typename T, typename ...Args> Node *makeNodeSimple(Args &&...As) {
     std::pair<Node *, bool> Result =
-        getOrCreateNode<T>(std::forward<Args>(As)...);
+        getOrCreateNode<T>(CreateNewNodes, std::forward<Args>(As)...);
     if (Result.second) {
       // Node is new. Make a note of that.
       MostRecentlyCreated = Result.first;
-    } else {
+    } else if (Result.first) {
       // Node is pre-existing; check if it's in our remapping table.
       if (auto *N = Remappings.lookup(Result.first)) {
         Result.first = N;
@@ -185,6 +190,8 @@ public:
 
   void reset() { MostRecentlyCreated = nullptr; }
 
+  void setCreateNewNodes(bool CNN) { CreateNewNodes = CNN; }
+
   void addRemapping(Node *A, Node *B) {
     // Note, we don't need to check whether B is also remapped, because if it
     // was we would have already remapped it when building it.
@@ -230,6 +237,7 @@ ItaniumManglingCanonicalizer::EquivalenceError
 ItaniumManglingCanonicalizer::addEquivalence(FragmentKind Kind, StringRef First,
                                              StringRef Second) {
   auto &Alloc = P->Demangler.ASTAllocator;
+  Alloc.setCreateNewNodes(true);
 
   auto Parse = [&](StringRef Str) {
     P->Demangler.reset(Str.begin(), Str.end());
@@ -302,6 +310,14 @@ ItaniumManglingCanonicalizer::addEquivalence(FragmentKind Kind, StringRef First,
 
 ItaniumManglingCanonicalizer::Key
 ItaniumManglingCanonicalizer::canonicalize(StringRef Mangling) {
+  P->Demangler.ASTAllocator.setCreateNewNodes(true);
+  P->Demangler.reset(Mangling.begin(), Mangling.end());
+  return reinterpret_cast<Key>(P->Demangler.parse());
+}
+
+ItaniumManglingCanonicalizer::Key
+ItaniumManglingCanonicalizer::lookup(StringRef Mangling) {
+  P->Demangler.ASTAllocator.setCreateNewNodes(false);
   P->Demangler.reset(Mangling.begin(), Mangling.end());
   return reinterpret_cast<Key>(P->Demangler.parse());
 }
diff --git a/llvm/unittests/Support/ItaniumManglingCanonicalizerTest.cpp b/llvm/unittests/Support/ItaniumManglingCanonicalizerTest.cpp
index 740cc0319f0..8807aa1abee 100644
--- a/llvm/unittests/Support/ItaniumManglingCanonicalizerTest.cpp
+++ b/llvm/unittests/Support/ItaniumManglingCanonicalizerTest.cpp
@@ -7,10 +7,13 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include <cstdlib>
 #include "llvm/Support/ItaniumManglingCanonicalizer.h"
+
 #include "gtest/gtest.h"
 
+#include <cstdlib>
+#include <map>
+
 using EquivalenceError = llvm::ItaniumManglingCanonicalizer::EquivalenceError;
 using FragmentKind = llvm::ItaniumManglingCanonicalizer::FragmentKind;
 
@@ -222,7 +225,9 @@ static std::initializer_list<Testcase> Testcases = {
     },
     {}
   },
+};
 
+static std::initializer_list<Testcase> ForwardTemplateReferenceTestcases = {
   // ForwardTemplateReference does not support canonicalization.
   // FIXME: We should consider ways of fixing this, perhaps by eliminating
   // the ForwardTemplateReference node with a tree transformation.
@@ -245,7 +250,8 @@ static std::initializer_list<Testcase> Testcases = {
   },
 };
 
-TEST(ItaniumManglingCanonicalizerTest, TestTestcases) {
+template<bool CanonicalizeFirst>
+static void testTestcases(std::initializer_list<Testcase> Testcases) {
   for (const auto &Testcase : Testcases) {
     llvm::ItaniumManglingCanonicalizer Canonicalizer;
     for (const auto &Equiv : Testcase.Equivalences) {
@@ -257,11 +263,21 @@ TEST(ItaniumManglingCanonicalizerTest, TestTestcases) {
     }
 
     using CanonKey = llvm::ItaniumManglingCanonicalizer::Key;
+
+    std::map<const EquivalenceClass*, CanonKey> Keys;
+    if (CanonicalizeFirst)
+      for (const auto &Class : Testcase.Classes)
+        Keys.insert({&Class, Canonicalizer.canonicalize(*Class.begin())});
+
     std::map<CanonKey, llvm::StringRef> Found;
     for (const auto &Class : Testcase.Classes) {
-      CanonKey ClassKey = {};
+      CanonKey ClassKey = Keys[&Class];
       for (llvm::StringRef Str : Class) {
-        CanonKey ThisKey = Canonicalizer.canonicalize(Str);
+        // Force a copy to be made when calling lookup to test that it doesn't
+        // retain any part of the provided string.
+        CanonKey ThisKey = CanonicalizeFirst
+                               ? Canonicalizer.lookup(std::string(Str))
+                               : Canonicalizer.canonicalize(Str);
         EXPECT_NE(ThisKey, CanonKey()) << "couldn't canonicalize " << Str;
         if (ClassKey) {
           EXPECT_EQ(ThisKey, ClassKey)
@@ -276,6 +292,21 @@ TEST(ItaniumManglingCanonicalizerTest, TestTestcases) {
   }
 }
 
+TEST(ItaniumManglingCanonicalizerTest, TestCanonicalize) {
+  testTestcases<false>(Testcases);
+}
+
+TEST(ItaniumManglingCanonicalizerTest, TestLookup) {
+  testTestcases<true>(Testcases);
+}
+
+TEST(ItaniumManglingCanonicalizerTest, TestForwardTemplateReference) {
+  // lookup(...) after canonicalization (intentionally) returns different
+  // values for this testcase.
+  testTestcases<false>(ForwardTemplateReferenceTestcases);
+}
+
+
 TEST(ItaniumManglingCanonicalizerTest, TestInvalidManglings) {
   llvm::ItaniumManglingCanonicalizer Canonicalizer;
   EXPECT_EQ(Canonicalizer.addEquivalence(FragmentKind::Type, "", "1X"),