Addressed feedback

Author: ElektroKill

src/DotNet/Pdb/PdbCustomDebugInfo.cs

@@ -1144,15 +1144,17 @@ protected virtual void InitializeDocuments() =>
 			Interlocked.CompareExchange(ref documents, new List<PdbDocument>(), null);
 	}
 
-	class PdbTypeDefinitionDocumentsDebugInfoMD : PdbTypeDefinitionDocumentsDebugInfo {
+	sealed class PdbTypeDefinitionDocumentsDebugInfoMD : PdbTypeDefinitionDocumentsDebugInfo {
 		readonly ModuleDef readerModule;
 		readonly IList<MDToken> documentTokens;
 
 		protected override void InitializeDocuments() {
 			var list = new List<PdbDocument>(documentTokens.Count);
-			for (var i = 0; i < documentTokens.Count; i++) {
-				if (readerModule.PdbState.tokenToDocument.TryGetValue(documentTokens[i], out var document))
-					list.Add(document);
+			if (readerModule.PdbState is not null) {
+				for (var i = 0; i < documentTokens.Count; i++) {
+					if (readerModule.PdbState.tokenToDocument.TryGetValue(documentTokens[i], out var document))
+						list.Add(document);
+				}
 			}
 			Interlocked.CompareExchange(ref documents, list, null);
 		}

src/DotNet/Pdb/PdbDocument.cs

@@ -53,6 +53,11 @@ public sealed class PdbDocument : IHasCustomDebugInformation {
 		public IList<PdbCustomDebugInfo> CustomDebugInfos => customDebugInfos;
 		IList<PdbCustomDebugInfo> customDebugInfos;
 
+		/// <summary>
+		/// Gets the Metadata token of the document if available.
+		/// </summary>
+		public MDToken? MDToken { get; internal set; }
+
 		/// <summary>
 		/// Default constructor
 		/// </summary>
@@ -86,6 +91,7 @@ internal void Initialize(SymbolDocument symDoc) {
 			customDebugInfos = new List<PdbCustomDebugInfo>();
 			foreach (var cdi in symDoc.CustomDebugInfos)
 				customDebugInfos.Add(cdi);
+			MDToken = symDoc.MDToken;
 		}
 
 		/// <summary>

src/DotNet/Pdb/PdbState.cs

@@ -124,6 +124,8 @@ PdbDocument Add_NoLock(PdbDocument doc) {
 			if (docDict.TryGetValue(doc, out var orig))
 				return orig;
 			docDict.Add(doc, doc);
+			if (doc.MDToken.HasValue)
+				tokenToDocument.Add(doc.MDToken.Value, doc);
 			return doc;
 		}
 
@@ -148,6 +150,8 @@ public bool Remove(PdbDocument doc) {
 #if THREAD_SAFE
 			theLock.EnterWriteLock(); try {
 #endif
+			if (doc.MDToken.HasValue)
+				tokenToDocument.Remove(doc.MDToken.Value);
 			return docDict.Remove(doc);
 #if THREAD_SAFE
 			} finally { theLock.ExitWriteLock(); }
@@ -189,6 +193,7 @@ public List<PdbDocument> RemoveAllDocuments(bool returnDocs) {
 			theLock.EnterWriteLock(); try {
 #endif
 			var docs = returnDocs ? new List<PdbDocument>(docDict.Values) : null;
+			tokenToDocument.Clear();
 			docDict.Clear();
 			return docs;
 #if THREAD_SAFE