2 files changed, 133 insertions, 0 deletions

diff --git a/llvm/lib/DebugInfo/PDB/Native/TpiHashing.cpp b/llvm/lib/DebugInfo/PDB/Native/TpiHashing.cpp
index 77a2d57a836..18708826ffc 100644
--- a/llvm/lib/DebugInfo/PDB/Native/TpiHashing.cpp
+++ b/llvm/lib/DebugInfo/PDB/Native/TpiHashing.cpp
@@ -50,6 +50,32 @@ static Expected<uint32_t> getHashForUdt(const CVType &Rec) {
 }
 
 template <typename T>
+static Expected<TagRecordHash> getTagRecordHashForUdt(const CVType &Rec) {
+  T Deserialized;
+  if (auto E = TypeDeserializer::deserializeAs(const_cast<CVType &>(Rec),
+                                               Deserialized))
+    return std::move(E);
+
+  ClassOptions Opts = Deserialized.getOptions();
+
+  bool ForwardRef = bool(Opts & ClassOptions::ForwardReference);
+
+  uint32_t ThisRecordHash = getHashForUdt(Deserialized, Rec.data());
+
+  // If we don't have a forward ref we can't compute the hash of it from the
+  // full record because it requires hashing the entire buffer.
+  if (!ForwardRef)
+    return TagRecordHash{std::move(Deserialized), ThisRecordHash, 0};
+
+  bool Scoped = bool(Opts & ClassOptions::Scoped);
+
+  StringRef NameToHash =
+      Scoped ? Deserialized.getUniqueName() : Deserialized.getName();
+  uint32_t FullHash = hashStringV1(NameToHash);
+  return TagRecordHash{std::move(Deserialized), FullHash, ThisRecordHash};
+}
+
+template <typename T>
 static Expected<uint32_t> getSourceLineHash(const CVType &Rec) {
   T Deserialized;
   if (auto E = TypeDeserializer::deserializeAs(const_cast<CVType &>(Rec),
@@ -60,6 +86,23 @@ static Expected<uint32_t> getSourceLineHash(const CVType &Rec) {
   return hashStringV1(StringRef(Buf, 4));
 }
 
+Expected<TagRecordHash> llvm::pdb::hashTagRecord(const codeview::CVType &Type) {
+  switch (Type.kind()) {
+  case LF_CLASS:
+  case LF_STRUCTURE:
+  case LF_INTERFACE:
+    return getTagRecordHashForUdt<ClassRecord>(Type);
+  case LF_UNION:
+    return getTagRecordHashForUdt<UnionRecord>(Type);
+  case LF_ENUM:
+    return getTagRecordHashForUdt<EnumRecord>(Type);
+  default:
+    assert(false && "Type is not a tag record!");
+  }
+  return make_error<StringError>("Invalid record type",
+                                 inconvertibleErrorCode());
+}
+
 Expected<uint32_t> llvm::pdb::hashTypeRecord(const CVType &Rec) {
   switch (Rec.kind()) {
   case LF_CLASS:
diff --git a/llvm/lib/DebugInfo/PDB/Native/TpiStream.cpp b/llvm/lib/DebugInfo/PDB/Native/TpiStream.cpp
index 0680b673380..de0f888137f 100644
--- a/llvm/lib/DebugInfo/PDB/Native/TpiStream.cpp
+++ b/llvm/lib/DebugInfo/PDB/Native/TpiStream.cpp
@@ -11,8 +11,11 @@
 
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/DebugInfo/CodeView/LazyRandomTypeCollection.h"
+#include "llvm/DebugInfo/CodeView/TypeDeserializer.h"
+#include "llvm/DebugInfo/CodeView/TypeIndexDiscovery.h"
 #include "llvm/DebugInfo/CodeView/TypeRecord.h"
 #include "llvm/DebugInfo/MSF/MappedBlockStream.h"
+#include "llvm/DebugInfo/PDB/Native/Hash.h"
 #include "llvm/DebugInfo/PDB/Native/PDBFile.h"
 #include "llvm/DebugInfo/PDB/Native/RawConstants.h"
 #include "llvm/DebugInfo/PDB/Native/RawError.h"
@@ -140,6 +143,93 @@ uint16_t TpiStream::getTypeHashStreamAuxIndex() const {
 uint32_t TpiStream::getNumHashBuckets() const { return Header->NumHashBuckets; }
 uint32_t TpiStream::getHashKeySize() const { return Header->HashKeySize; }
 
+void TpiStream::buildHashMap() {
+  if (!HashMap.empty())
+    return;
+  if (HashValues.empty())
+    return;
+
+  HashMap.resize(Header->NumHashBuckets);
+
+  TypeIndex TIB{Header->TypeIndexBegin};
+  TypeIndex TIE{Header->TypeIndexEnd};
+  while (TIB < TIE) {
+    uint32_t HV = HashValues[TIB.toArrayIndex()];
+    HashMap[HV].push_back(TIB++);
+  }
+}
+
+bool TpiStream::supportsTypeLookup() const { return !HashMap.empty(); }
+
+template <typename RecordT> static ClassOptions getUdtOptions(CVType CVT) {
+  RecordT Record;
+  if (auto EC = TypeDeserializer::deserializeAs<RecordT>(CVT, Record)) {
+    consumeError(std::move(EC));
+    return ClassOptions::None;
+  }
+  return Record.getOptions();
+}
+
+static bool isUdtForwardRef(CVType CVT) {
+  ClassOptions UdtOptions = ClassOptions::None;
+  switch (CVT.kind()) {
+  case LF_STRUCTURE:
+  case LF_CLASS:
+  case LF_INTERFACE:
+    UdtOptions = getUdtOptions<ClassRecord>(std::move(CVT));
+    break;
+  case LF_ENUM:
+    UdtOptions = getUdtOptions<EnumRecord>(std::move(CVT));
+    break;
+  case LF_UNION:
+    UdtOptions = getUdtOptions<UnionRecord>(std::move(CVT));
+    break;
+  default:
+    return false;
+  }
+  return (UdtOptions & ClassOptions::ForwardReference) != ClassOptions::None;
+}
+
+Expected<TypeIndex>
+TpiStream::findFullDeclForForwardRef(TypeIndex ForwardRefTI) const {
+  CVType F = Types->getType(ForwardRefTI);
+  if (!isUdtForwardRef(F))
+    return ForwardRefTI;
+
+  Expected<TagRecordHash> ForwardTRH = hashTagRecord(F);
+  if (!ForwardTRH)
+    return ForwardTRH.takeError();
+
+  TagRecordHash Copy = std::move(*ForwardTRH);
+  uint32_t BucketIdx = ForwardTRH->FullRecordHash % Header->NumHashBuckets;
+
+  for (TypeIndex TI : HashMap[BucketIdx]) {
+    CVType CVT = Types->getType(TI);
+    if (CVT.kind() != F.kind())
+      continue;
+
+    Expected<TagRecordHash> FullTRH = hashTagRecord(CVT);
+    if (!FullTRH)
+      return FullTRH.takeError();
+    if (ForwardTRH->FullRecordHash != FullTRH->FullRecordHash)
+      continue;
+    TagRecord &ForwardTR = ForwardTRH->getRecord();
+    TagRecord &FullTR = FullTRH->getRecord();
+
+    if (!ForwardTR.hasUniqueName()) {
+      if (ForwardTR.getName() == FullTR.getName())
+        return TI;
+      continue;
+    }
+
+    if (!FullTR.hasUniqueName())
+      continue;
+    if (ForwardTR.getUniqueName() == FullTR.getUniqueName())
+      return TI;
+  }
+  return ForwardRefTI;
+}
+
 BinarySubstreamRef TpiStream::getTypeRecordsSubstream() const {
   return TypeRecordsSubstream;
 }