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.

[rsmith]

Summary

This could also be used to implement demangling with a fixed-size
temporary storage buffer.

Depends on https://reviews.llvm.org/D50935

[erik.pilkington]

It seems like we would need to add a lot more nullptr checks in order to really support this then this patch does. One potential workaround is to add a subclass of Node in ItaniumManglingCanonicalizer.cpp that serves as the canonical 'sentinel' node. We could just hand that out whenever we would otherwise allocate anything, and ignore the results of parse() if we ever did. I think that would probably maybe work.

[rsmith]

In D51003#1208827, @erik.pilkington wrote:

It seems like we would need to add a lot more nullptr checks in order to really support this then this patch does. One potential workaround is to add a subclass of Node in ItaniumManglingCanonicalizer.cpp that serves as the canonical 'sentinel' node. We could just hand that out whenever we would otherwise allocate anything, and ignore the results of parse() if we ever did. I think that would probably maybe work.

This patch fixes up every place where we create a node and do not immediately return it; all parsing functions returning a Node* are already assumed to potentially fail, so their callers already perform a null check. So I think this is the full extent of what we need to do to support this.

(If we wanted to actually make the demangler be able to run with a fixed size buffer, there are two other things we need to address: NodeArray allocation and the explicit malloc/free calls made by PODSmallVector. Those both seem feasible, if we ever want to actually go there.)

[erik.pilkington]

In D51003#1210421, @rsmith wrote:

In D51003#1208827, @erik.pilkington wrote:

It seems like we would need to add a lot more nullptr checks in order to really support this then this patch does. One potential workaround is to add a subclass of Node in ItaniumManglingCanonicalizer.cpp that serves as the canonical 'sentinel' node. We could just hand that out whenever we would otherwise allocate anything, and ignore the results of parse() if we ever did. I think that would probably maybe work.

This patch fixes up every place where we create a node and do not immediately return it; all parsing functions returning a Node* are already assumed to potentially fail, so their callers already perform a null check. So I think this is the full extent of what we need to do to support this.

Oh, right, my mistake. LGTM after some inline comments.

(If we wanted to actually make the demangler be able to run with a fixed size buffer, there are two other things we need to address: NodeArray allocation and the explicit malloc/free calls made by PODSmallVector. Those both seem feasible, if we ever want to actually go there.)

Ya, I think that would be a good thing to do. It'd also allow us to more gracefully handle failed allocation (right new we just std::terminate()), which would be nice.

include/llvm/Demangle/ItaniumDemangle.h

2788 ↗

(On Diff #161584)

Missed a spot here too!

2852 ↗

(On Diff #161584)

Might be nice to inline the nullptr check on Comp into this function too.

2888 ↗

(On Diff #161584)

Missed a spot!

4773 ↗

(On Diff #161584)

You should check here too. It won't result in a crash, but it'll cause lookup to return a non-null key when it shouldn't.

[rsmith]

include/llvm/Demangle/ItaniumDemangle.h

2905–2915 ↗

(On Diff #161584)

Incidentally: this doesn't look quite right to me. A <substitution> can only appear at the start of a <prefix>, right? (We should presumably parse this before the loop, like we do for St.)

llvm/trunk/include/llvm/Demangle/ItaniumDemangle.h

template<typename Alloc> Node *Db<Alloc>::parseLocalName(NameState *State) {
   if (!consumeIf('Z'))
     return nullptr;
   Node *Encoding = parseEncoding();
   if (Encoding == nullptr || !consumeIf('E'))
     return nullptr;
 
   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')) {
     parseNumber(true);
     if (!consumeIf('_'))
       return nullptr;
     Node *N = parseName(State);
     if (N == nullptr)
Node *Db<Alloc>::parseCtorDtorName(Node *&SoFar, NameState *State) {
   if (SoFar->getKind() == Node::KSpecialSubstitution) {
     auto SSK = static_cast<SpecialSubstitution *>(SoFar)->SSK;
     switch (SSK) {
     case SpecialSubKind::string:
     case SpecialSubKind::istream:
     case SpecialSubKind::ostream:
     case SpecialSubKind::iostream:
       SoFar = make<ExpandedSpecialSubstitution>(SSK);
+      if (!SoFar)
+        return nullptr;
     default:
       break;
     }
   }
 
   if (consumeIf('C')) {
     bool IsInherited = consumeIf('I');
     if (look() != '1' && look() != '2' && look() != '3' && look() != '5')
if (consumeIf('O')) {
     if (State) State->ReferenceQualifier = FrefQualRValue;
   } else if (consumeIf('R')) {
     if (State) State->ReferenceQualifier = FrefQualLValue;
   } else
     if (State) State->ReferenceQualifier = FrefQualNone;
 
   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
 
     // <data-member-prefix> := <member source-name> [<template-args>] M
     if (consumeIf('M')) {
       if (SoFar == nullptr)
         return nullptr;
       continue;
     }
 
     //          ::= <template-param>
     if (look() == 'T') {
-      Node *TP = parseTemplateParam();
-      if (TP == nullptr)
+      if (!PushComponent(parseTemplateParam()))
         return nullptr;
-      PushComponent(TP);
       Subs.push_back(SoFar);
       continue;
     }
 
     //          ::= <template-prefix> <template-args>
     if (look() == 'I') {
       Node *TA = parseTemplateArgs(State != nullptr);
       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;
     }
 
     //          ::= <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;
     }
 
     //          ::= <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;
     }
 
     // Parse an <unqualified-name> thats actually a <ctor-dtor-name>.
     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;
       Subs.push_back(SoFar);
       continue;
     }
 
     //          ::= <prefix> <unqualified-name>
-    Node *N = parseUnqualifiedName(State);
-    if (N == nullptr)
+    if (!PushComponent(parseUnqualifiedName(State)))
       return nullptr;
-    PushComponent(N);
     Subs.push_back(SoFar);
   }
 
   if (SoFar == nullptr || Subs.empty())
     return nullptr;
 
   Subs.pop_back();
   return SoFar;
if (consumeIf("srN")) {
     if (SoFar == nullptr)
       return nullptr;
 
     if (look() == 'I') {
       Node *TA = parseTemplateArgs();
       if (TA == nullptr)
         return nullptr;
       SoFar = make<NameWithTemplateArgs>(SoFar, TA);
+      if (!SoFar)
+        return nullptr;
     }
 
     while (!consumeIf('E')) {
       Node *Qual = parseSimpleId();
       if (Qual == nullptr)
         return nullptr;
       SoFar = make<QualifiedName>(SoFar, Qual);
+      if (!SoFar)
+        return nullptr;
     }
 
     Node *Base = parseBaseUnresolvedName();
     if (Base == nullptr)
       return nullptr;
     return make<QualifiedName>(SoFar, Base);
   }
 
do {
       if (Qual == nullptr)
         return nullptr;
       if (SoFar)
         SoFar = make<QualifiedName>(SoFar, Qual);
       else if (Global)
         SoFar = make<GlobalQualifiedName>(Qual);
       else
         SoFar = Qual;
+      if (!SoFar)
+        return nullptr;
     } while (!consumeIf('E'));
   }
   //      sr <unresolved-type>                 <base-unresolved-name>
   //      sr <unresolved-type> <template-args> <base-unresolved-name>
   else {
     SoFar = parseUnresolvedType();
     if (SoFar == nullptr)
       return nullptr;
 
     if (look() == 'I') {
       Node *TA = parseTemplateArgs();
       if (TA == nullptr)
         return nullptr;
       SoFar = make<NameWithTemplateArgs>(SoFar, TA);
+      if (!SoFar)
+        return nullptr;
     }
   }
 
   assert(SoFar != nullptr);
 
   Node *Base = parseBaseUnresolvedName();
   if (Base == nullptr)
     return nullptr;
   return make<QualifiedName>(SoFar, Base);
 }
 
 // <abi-tags> ::= <abi-tag> [<abi-tags>]
 // <abi-tag> ::= B <source-name>
 template<typename Alloc> Node *Db<Alloc>::parseAbiTags(Node *N) {
   while (consumeIf('B')) {
     StringView SN = parseBareSourceName();
     if (SN.empty())
       return nullptr;
     N = make<AbiTagAttr>(N, SN);
+    if (!N)
+      return nullptr;
   }
   return N;
 }
 
 // <number> ::= [n] <non-negative decimal integer>
 template<typename Alloc>
 StringView Db<Alloc>::parseNumber(bool AllowNegative) {
   const char *Tmp = First;

 // <ref-qualifier> ::= R                   # & ref-qualifier
 // <ref-qualifier> ::= O                   # && ref-qualifier
 template<typename Alloc> Node *Db<Alloc>::parseFunctionType() {
   Qualifiers CVQuals = parseCVQualifiers();
 
   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')) {
       Node *T = parseType();
       if (T == nullptr)
         return nullptr;
       Names.push_back(T);
     }
     ExceptionSpec =
       make<DynamicExceptionSpec>(popTrailingNodeArray(SpecsBegin));
+    if (!ExceptionSpec)
+      return nullptr;
   }
 
   consumeIf("Dx"); // transaction safe
 
   if (!consumeIf('F'))
     return nullptr;
   consumeIf('Y'); // extern "C"
   Node *ReturnType = parseType();
case 'P': {
       First += 2;
       size_t ArgsBegin = Names.size();
       while (!consumeIf('E')) {
         Node *Arg = parseTemplateArg();
         if (Arg == nullptr)
           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;
   case 't':
     switch (First[1]) {
     case 'e': {
       First += 2;
       Node *Ex = parseExpr();
if (consumeIf("Ua9enable_ifI")) {
     size_t BeforeArgs = Names.size();
     while (!consumeIf('E')) {
       Node *Arg = parseTemplateArg();
       if (Arg == nullptr)
         return nullptr;
       Names.push_back(Arg);
     }
     Attrs = make<EnableIfAttr>(popTrailingNodeArray(BeforeArgs));
+    if (!Attrs)
+      return nullptr;
   }
 
   Node *ReturnType = nullptr;
   if (!NameInfo.CtorDtorConversion && NameInfo.EndsWithTemplateArgs) {
     ReturnType = parseType();
     if (ReturnType == nullptr)
       return nullptr;
   }
case 'o':
       break;
     case 'd':
       ++First;
       SpecialSub = make<SpecialSubstitution>(SpecialSubKind::iostream);
       break;
     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.
     Node *WithTags = parseAbiTags(SpecialSub);
     if (WithTags != SpecialSub) {
       Subs.push_back(WithTags);
       SpecialSub = WithTags;
     }
template<typename Alloc> Node *Db<Alloc>::parseTemplateParam() {
   if (ParsingLambdaParams)
     return make<NameType>("auto");
 
   // If we're in a context where this <template-param> refers to a
   // <template-arg> further ahead in the mangled name (currently just conversion
   // 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));
     return ForwardRef;
   }
 
   if (Index >= TemplateParams.size())
     return nullptr;
if (TagTemplates) {
       TemplateParams = std::move(OldParams);
       if (Arg == nullptr)
         return nullptr;
       Names.push_back(Arg);
       Node *TableEntry = Arg;
       if (Arg->getKind() == Node::KTemplateArgumentPack) {
         TableEntry = make<ParameterPack>(
             static_cast<TemplateArgumentPack*>(TableEntry)->getElements());
+        if (!TableEntry)
+          return nullptr;
       }
       TemplateParams.push_back(TableEntry);
     } else {
       Node *Arg = parseTemplateArg();
       if (Arg == nullptr)
         return nullptr;
       Names.push_back(Arg);
     }

llvm/trunk/include/llvm/Support/ItaniumManglingCanonicalizer.h

public:
   using Key = uintptr_t;
 
   /// Form a canonical key for the specified mangling. They key will be the
   /// same for all equivalent manglings, and different for any two
   /// non-equivalent manglings, but is otherwise unspecified.
   ///
   /// 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;
 };
 } // namespace llvm
 
 #endif // LLVM_SUPPORT_ITANIUMMANGLINGCANONICALIZER_H

llvm/trunk/lib/Support/ItaniumManglingCanonicalizer.cpp

class FoldingNodeAllocator {
 
   BumpPtrAllocator RawAlloc;
   llvm::FoldingSet<NodeHeader> Nodes;
 
 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...);
 
     void *InsertPos;
     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 =
         RawAlloc.Allocate(sizeof(NodeHeader) + sizeof(T), alignof(NodeHeader));
     NodeHeader *New = new (Storage) NodeHeader;
     T *Result = new (New->getNode()) T(std::forward<Args>(As)...);
     Nodes.InsertNode(New, InsertPos);
     return {Result, true};
   }
 
   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) {
     return RawAlloc.Allocate(sizeof(Node *) * sz, alignof(Node *));
   }
 };
 
 // FIXME: Don't canonicalize forward template references for now, because they
 // contain state (the resolved template node) that's not known at their point
 // 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),
           true};
 }
 
 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;
         assert(Remappings.find(Result.first) == Remappings.end() &&
                "should never need multiple remap steps");
       }
       if (Result.first == TrackedNode)
         TrackedNodeIsUsed = true;

 
 public:
   template<typename T, typename ...Args> Node *makeNode(Args &&...As) {
     return MakeNodeImpl<T>{*this}.make(std::forward<Args>(As)...);
   }
 
   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.
     Remappings.insert(std::make_pair(A, B));
   }
 
   bool isMostRecentlyCreated(Node *N) const { return MostRecentlyCreated == N; }
 

 
 ItaniumManglingCanonicalizer::ItaniumManglingCanonicalizer() : P(new Impl) {}
 ItaniumManglingCanonicalizer::~ItaniumManglingCanonicalizer() { delete P; }
 
 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());
     Node *N = nullptr;
     switch (Kind) {
       // A <name>, with minor extensions to allow arbitrary namespace and
       // template names that can't easily be written as <name>s.
     case FragmentKind::Name:
ItaniumManglingCanonicalizer::addEquivalence(FragmentKind Kind, StringRef First,
   else
     return EquivalenceError::ManglingAlreadyUsed;
 
   return EquivalenceError::Success;
 }
 
 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());
 }

llvm/trunk/unittests/Support/ItaniumManglingCanonicalizerTest.cpp

 //===-------------- ItaniumManglingCanonicalizerTest.cpp ------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is dual licensed under the MIT and the University of Illinois Open
 // Source Licenses. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
-#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;
 
 struct Equivalence {
   FragmentKind Kind;
   llvm::StringRef First;
   llvm::StringRef Second;
 };
static std::initializer_list<Testcase> Testcases = {
   {
     {
       {FragmentKind::Name, "3std", "St"},
       {FragmentKind::Name, "1X", "1Y"},
       {FragmentKind::Name, "N1X1ZE", "N1Y1ZE"},
     },
     {}
   },
+};
 
+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.
   {
     {
       // X::operator T() <with T = A> == Y::operator T() <with T = A>
       {FragmentKind::Encoding, "N1XcvT_I1AEEv", "N1YcvT_I1AEEv"},
       // A == B
       {FragmentKind::Name, "1A", "1B"},
     },
     {
       // All combinations result in unique equivalence classes.
       {"_ZN1XcvT_I1AEEv"},
       {"_ZN1XcvT_I1BEEv"},
       {"_ZN1YcvT_I1AEEv"},
       {"_ZN1YcvT_I1BEEv"},
       // Even giving the same string twice gives a new class.
       {"_ZN1XcvT_I1AEEv"},
     }
   },
 };
 
-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) {
       auto Result =
           Canonicalizer.addEquivalence(Equiv.Kind, Equiv.First, Equiv.Second);
       EXPECT_EQ(Result, EquivalenceError::Success)
           << "couldn't add equivalence between " << Equiv.First << " and "
           << Equiv.Second;
     }
 
     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)
               << Str << " not in the same class as " << *Class.begin();
         } else {
           ClassKey = ThisKey;
         }
       }
       EXPECT_TRUE(Found.insert({ClassKey, *Class.begin()}).second)
           << *Class.begin() << " is in the same class as " << Found[ClassKey];
     }
   }
 }
 
+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"),
             EquivalenceError::InvalidFirstMangling);
   EXPECT_EQ(Canonicalizer.addEquivalence(FragmentKind::Type, "1X", "1ab"),
             EquivalenceError::InvalidSecondMangling);
   EXPECT_EQ(Canonicalizer.canonicalize("_Z3fooE"),
             llvm::ItaniumManglingCanonicalizer::Key());