Merge pull request #330 from Washi1337/hotfix/preserve-order-nested-t…

…yperefs

Hotfix: Preserve unusual order of nested type refs

src/AsmResolver.DotNet/Builder/DotNetDirectoryBuffer.CodedIndices.cs

@@ -26,16 +26,16 @@ private void AddCustomAttribute(MetadataToken ownerToken, CustomAttribute attrib
             table.Add(attribute, row);
         }
 
-        private uint AddResolutionScope(IResolutionScope? scope)
+        private uint AddResolutionScope(IResolutionScope? scope, bool allowDuplicates, bool preserveRid)
         {
             if (!AssertIsImported(scope))
                 return 0;
 
             var token = scope.MetadataToken.Table switch
             {
-                TableIndex.AssemblyRef => GetAssemblyReferenceToken(scope as AssemblyReference),
-                TableIndex.TypeRef => GetTypeReferenceToken(scope as TypeReference),
-                TableIndex.ModuleRef => GetModuleReferenceToken(scope as ModuleReference),
+                TableIndex.AssemblyRef => AddAssemblyReference(scope as AssemblyReference, allowDuplicates, preserveRid),
+                TableIndex.TypeRef => AddTypeReference(scope as TypeReference, allowDuplicates, preserveRid),
+                TableIndex.ModuleRef => AddModuleReference(scope as ModuleReference, allowDuplicates, preserveRid),
                 TableIndex.Module => 0,
                 _ => throw new ArgumentOutOfRangeException(nameof(scope))
             };

src/AsmResolver.DotNet/Builder/DotNetDirectoryBuffer.TokenProvider.cs

@@ -10,7 +10,10 @@ public partial class DotNetDirectoryBuffer : IMetadataTokenProvider
         public uint GetUserStringIndex(string value) => Metadata.UserStringsStream.GetStringIndex(value);
 
         /// <inheritdoc />
-        public MetadataToken GetTypeReferenceToken(TypeReference? type) => AddTypeReference(type, false);
+        public MetadataToken GetTypeReferenceToken(TypeReference? type)
+        {
+            return AddTypeReference(type, false, false);
+        }
 
         /// <summary>
         /// Adds a type reference to the buffer.
@@ -20,19 +23,25 @@ public partial class DotNetDirectoryBuffer : IMetadataTokenProvider
         /// <c>true</c> if the row is always to be added to the end of the buffer, <c>false</c> if a duplicated row
         /// is supposed to be removed and the token of the original should be returned instead.
         /// </param>
+        /// <param name="preserveRid">
+        /// <c>true</c> if the metadata token of the type should be preserved, <c>false</c> otherwise.
+        /// </param>
         /// <returns>The newly assigned metadata token.</returns>
-        public MetadataToken AddTypeReference(TypeReference? type, bool allowDuplicates)
+        public MetadataToken AddTypeReference(TypeReference? type, bool allowDuplicates, bool preserveRid)
         {
             if (!AssertIsImported(type))
                 return MetadataToken.Zero;
 
             var table = Metadata.TablesStream.GetDistinctTable<TypeReferenceRow>(TableIndex.TypeRef);
             var row = new TypeReferenceRow(
-                AddResolutionScope(type.Scope),
+                AddResolutionScope(type.Scope, allowDuplicates, preserveRid),
                 Metadata.StringsStream.GetStringIndex(type.Name),
                 Metadata.StringsStream.GetStringIndex(type.Namespace));
 
-            var token = table.Add(row, allowDuplicates);
+            var token = preserveRid
+                ? table.Insert(type.MetadataToken.Rid, row, allowDuplicates)
+                : table.Add(row, allowDuplicates);
+
             _tokenMapping.Register(type, token);
             AddCustomAttributes(token, type);
             return token;
@@ -164,7 +173,7 @@ public MetadataToken AddStandAloneSignature(StandAloneSignature? signature, bool
         /// <inheritdoc />
         public MetadataToken GetAssemblyReferenceToken(AssemblyReference? assembly)
         {
-            return AddAssemblyReference(assembly, false);
+            return AddAssemblyReference(assembly, false, false);
         }
 
         /// <summary>
@@ -175,8 +184,11 @@ public MetadataToken GetAssemblyReferenceToken(AssemblyReference? assembly)
         /// <c>true</c> if the row is always to be added to the end of the buffer, <c>false</c> if a duplicated row
         /// is supposed to be removed and the token of the original should be returned instead.
         /// </param>
+        /// <param name="preserveRid">
+        /// <c>true</c> if the metadata token of the assembly should be preserved, <c>false</c> otherwise.
+        /// </param>
         /// <returns>The newly assigned metadata token.</returns>
-        public MetadataToken AddAssemblyReference(AssemblyReference? assembly, bool allowDuplicates)
+        public MetadataToken AddAssemblyReference(AssemblyReference? assembly, bool allowDuplicates, bool preserveRid)
         {
             if (assembly is null || !AssertIsImported(assembly))
                 return MetadataToken.Zero;
@@ -193,7 +205,10 @@ public MetadataToken AddAssemblyReference(AssemblyReference? assembly, bool allo
                 Metadata.StringsStream.GetStringIndex(assembly.Culture),
                 Metadata.BlobStream.GetBlobIndex(assembly.HashValue));
 
-            var token = table.Add(row, allowDuplicates);
+            var token = preserveRid
+                ? table.Insert(assembly.MetadataToken.Rid, row, allowDuplicates)
+                : table.Add(row, allowDuplicates);
+
             AddCustomAttributes(token, assembly);
             return token;
         }
@@ -205,7 +220,7 @@ public MetadataToken AddAssemblyReference(AssemblyReference? assembly, bool allo
         /// <returns>The new metadata token assigned to the module reference.</returns>
         public MetadataToken GetModuleReferenceToken(ModuleReference? reference)
         {
-            return AddModuleReference(reference, false);
+            return AddModuleReference(reference, false, false);
         }
 
         /// <summary>
@@ -216,16 +231,22 @@ public MetadataToken GetModuleReferenceToken(ModuleReference? reference)
         /// <c>true</c> if the row is always to be added to the end of the buffer, <c>false</c> if a duplicated row
         /// is supposed to be removed and the token of the original should be returned instead.
         /// </param>
+        /// <param name="preserveRid">
+        /// <c>true</c> if the metadata token of the module should be preserved, <c>false</c> otherwise.
+        /// </param>
         /// <returns>The newly assigned metadata token.</returns>
-        public MetadataToken AddModuleReference(ModuleReference? reference, bool allowDuplicates)
+        public MetadataToken AddModuleReference(ModuleReference? reference, bool allowDuplicates, bool preserveRid)
         {
             if (!AssertIsImported(reference))
                 return MetadataToken.Zero;
 
             var table = Metadata.TablesStream.GetDistinctTable<ModuleReferenceRow>(TableIndex.ModuleRef);
 
             var row = new ModuleReferenceRow(Metadata.StringsStream.GetStringIndex(reference.Name));
-            var token = table.Add(row, allowDuplicates);
+            var token = preserveRid
+                ? table.Insert(reference.MetadataToken.Rid, row, allowDuplicates)
+                : table.Add(row, allowDuplicates);
+
             AddCustomAttributes(token, reference);
             return token;
         }

src/AsmResolver.DotNet/Builder/DotNetDirectoryFactory.cs

@@ -212,19 +212,19 @@ private void ImportBasicTablesIfSpecified(ModuleDefinition module, DotNetDirecto
             if ((MetadataBuilderFlags & MetadataBuilderFlags.PreserveAssemblyReferenceIndices) != 0)
             {
                 ImportTables<AssemblyReference>(module, TableIndex.AssemblyRef,
-                    r => buffer.AddAssemblyReference(r, true));
+                    r => buffer.AddAssemblyReference(r, true, true));
             }
 
             if ((MetadataBuilderFlags & MetadataBuilderFlags.PreserveModuleReferenceIndices) != 0)
             {
                 ImportTables<ModuleReference>(module, TableIndex.ModuleRef,
-                    r => buffer.AddModuleReference(r, true));
+                    r => buffer.AddModuleReference(r, true, true));
             }
 
             if ((MetadataBuilderFlags & MetadataBuilderFlags.PreserveTypeReferenceIndices) != 0)
             {
                 ImportTables<TypeReference>(module, TableIndex.TypeRef,
-                    r => buffer.AddTypeReference(r, true));
+                    r => buffer.AddTypeReference(r, true, true));
             }
         }
 

src/AsmResolver.DotNet/Builder/Metadata/Tables/DistinctMetadataTableBuffer.cs

@@ -50,6 +50,33 @@ public TRow this[uint rid]
         /// <inheritdoc />
         public MetadataToken Add(in TRow row) => Add(row, false);
 
+        /// <inheritdoc />
+        public MetadataToken Insert(uint rid, in TRow row) => Insert(rid, row, false);
+
+        /// <summary>
+        /// Inserts a row into the metadata table at the provided row identifier.
+        /// </summary>
+        /// <param name="rid">The row identifier.</param>
+        /// <param name="row">The row to add.</param>
+        /// <param name="allowDuplicates">
+        /// <c>true</c> if the row is always to be added to the end of the buffer, <c>false</c> if a duplicated row
+        /// is supposed to be removed and the token of the original should be returned instead.</param>
+        /// <returns>The metadata token that this row was assigned to.</returns>
+        public MetadataToken Insert(uint rid, in TRow row, bool allowDuplicates)
+        {
+            if (!_entries.TryGetValue(row, out var token))
+            {
+                token = _underlyingBuffer.Insert(rid, in row);
+                _entries.Add(row, token);
+            }
+            else if (allowDuplicates)
+            {
+                token = _underlyingBuffer.Insert(rid, in row);
+            }
+
+            return token;
+        }
+
         /// <summary>
         /// Adds a row to the metadata table buffer.
         /// </summary>

src/AsmResolver.DotNet/Builder/Metadata/Tables/IMetadataTableBuffer.cs

@@ -56,5 +56,13 @@ TRow this[uint rid]
         /// <param name="row">The row to add.</param>
         /// <returns>The metadata token that this row was assigned to.</returns>
         MetadataToken Add(in TRow row);
+
+        /// <summary>
+        /// Inserts a row into the metadata table at the provided row identifier.
+        /// </summary>
+        /// <param name="rid">The row identifier.</param>
+        /// <param name="row">The row to add.</param>
+        /// <returns>The metadata token that this row was assigned to.</returns>
+        MetadataToken Insert(uint rid, in TRow row);
     }
 }

src/AsmResolver.DotNet/Builder/Metadata/Tables/UnsortedMetadataTableBuffer.cs

@@ -15,7 +15,9 @@ public class UnsortedMetadataTableBuffer<TRow> : IMetadataTableBuffer<TRow>
         where TRow : struct, IMetadataRow
     {
         private readonly RefList<TRow> _entries = new();
+        private readonly BitList _available = new();
         private readonly MetadataTable<TRow> _table;
+        private uint _currentRid;
 
         /// <summary>
         /// Creates a new unsorted metadata table buffer.
@@ -33,7 +35,11 @@ public UnsortedMetadataTableBuffer(MetadataTable<TRow> table)
         public virtual TRow this[uint rid]
         {
             get => _entries[(int) (rid - 1)];
-            set => _entries[(int) (rid - 1)] = value;
+            set
+            {
+                _entries[(int) (rid - 1)] = value;
+                _available[(int) (rid - 1)] = false;
+            }
         }
 
         /// <inheritdoc />
@@ -42,8 +48,45 @@ public virtual TRow this[uint rid]
         /// <inheritdoc />
         public virtual MetadataToken Add(in TRow row)
         {
-            _entries.Add(row);
-            return new MetadataToken(_table.TableIndex, (uint) _entries.Count);
+            // Move over unavailable slots.
+            while (_currentRid < _available.Count && !_available[(int) _currentRid])
+                _currentRid++;
+
+            // If we moved over all entries, we're adding to the end.
+            if (_currentRid == _entries.Count)
+            {
+                _currentRid++;
+            }
+
+            return Insert(_currentRid++, row);
+        }
+
+        /// <inheritdoc />
+        public MetadataToken Insert(uint rid, in TRow row)
+        {
+            EnsureRowsAllocated(rid);
+
+            var token = new MetadataToken(_table.TableIndex, rid);
+
+            if (!_available[(int) (rid - 1)])
+            {
+                if (EqualityComparer<TRow>.Default.Equals(row, _entries[(int) (rid - 1)]))
+                    return token;
+
+                throw new InvalidOperationException($"Token 0x{token.ToString()} is already in use.");
+            }
+
+            this[rid] = row;
+            return token;
+        }
+
+        private void EnsureRowsAllocated(uint rid)
+        {
+            while (_entries.Count < rid)
+            {
+                _entries.Add(default);
+                _available.Add(true);
+            }
         }
 
         /// <inheritdoc />

src/AsmResolver.DotNet/Builder/PEImageBuildResult.cs

@@ -34,7 +34,7 @@ public IPEImage? ConstructedImage
         /// Gets a value indicating whether the image was constructed successfully or not.
         /// </summary>
         [MemberNotNullWhen(false, nameof(ConstructedImage))]
-        public bool HasFailed => ConstructedImage is null;
+        public bool HasFailed => DiagnosticBag.IsFatal;
 
         /// <summary>
         /// Gets the bag containing the diagnostics that were collected during the construction of the image.