// dnlib: See LICENSE.txt for more info using System; using System.Collections.Generic; using System.Diagnostics; using dnlib.IO; using dnlib.PE; namespace dnlib.DotNet.Writer { /// /// Stores all method body chunks /// public sealed class MethodBodyChunks : IChunk { const uint FAT_BODY_ALIGNMENT = 4; Dictionary tinyMethodsDict; Dictionary fatMethodsDict; readonly List tinyMethods; readonly List fatMethods; readonly List reusedMethods; readonly Dictionary rvaToReusedMethod; readonly bool shareBodies; FileOffset offset; RVA rva; uint length; bool setOffsetCalled; readonly bool alignFatBodies; uint savedBytes; readonly struct ReusedMethodInfo { public readonly MethodBody MethodBody; public readonly RVA RVA; public ReusedMethodInfo(MethodBody methodBody, RVA rva) { MethodBody = methodBody; RVA = rva; } } /// public FileOffset FileOffset => offset; /// public RVA RVA => rva; /// /// Gets the number of bytes saved by re-using method bodies /// public uint SavedBytes => savedBytes; internal bool CanReuseOldBodyLocation { get; set; } internal bool ReusedAllMethodBodyLocations => tinyMethods.Count == 0 && fatMethods.Count == 0; internal bool HasReusedMethods => reusedMethods.Count > 0; /// /// Constructor /// /// true if bodies can be shared public MethodBodyChunks(bool shareBodies) { this.shareBodies = shareBodies; alignFatBodies = true; if (shareBodies) { tinyMethodsDict = new Dictionary(); fatMethodsDict = new Dictionary(); } tinyMethods = new List(); fatMethods = new List(); reusedMethods = new List(); rvaToReusedMethod = new Dictionary(); } /// /// Adds a and returns the one that has been cached /// /// The method body /// The cached method body public MethodBody Add(MethodBody methodBody) => Add(methodBody, 0, 0); internal MethodBody Add(MethodBody methodBody, RVA origRva, uint origSize) { if (setOffsetCalled) throw new InvalidOperationException("SetOffset() has already been called"); if (CanReuseOldBodyLocation && origRva != 0 && origSize != 0 && methodBody.CanReuse(origRva, origSize)) { if (rvaToReusedMethod.TryGetValue((uint)origRva, out var reusedMethod)) { if (methodBody.Equals(reusedMethod)) return reusedMethod; } else { rvaToReusedMethod.Add((uint)origRva, methodBody); reusedMethods.Add(new ReusedMethodInfo(methodBody, origRva)); return methodBody; } } if (shareBodies) { var dict = methodBody.IsFat ? fatMethodsDict : tinyMethodsDict; if (dict.TryGetValue(methodBody, out var cached)) { savedBytes += (uint)methodBody.GetApproximateSizeOfMethodBody(); return cached; } dict[methodBody] = methodBody; } var list = methodBody.IsFat ? fatMethods : tinyMethods; list.Add(methodBody); return methodBody; } /// Removes the specified method body from this chunk /// The method body /// if the method body is removed public bool Remove(MethodBody methodBody) { if (methodBody is null) throw new ArgumentNullException(nameof(methodBody)); if (setOffsetCalled) throw new InvalidOperationException("SetOffset() has already been called"); if (CanReuseOldBodyLocation) throw new InvalidOperationException("Reusing old body locations is enabled. Can't remove bodies."); var list = methodBody.IsFat ? fatMethods : tinyMethods; return list.Remove(methodBody); } internal void InitializeReusedMethodBodies(Func getNewFileOffset) { foreach (var info in reusedMethods) { var offset = getNewFileOffset(info.RVA); info.MethodBody.SetOffset(offset, info.RVA); } } internal void WriteReusedMethodBodies(DataWriter writer, long destStreamBaseOffset) { foreach (var info in reusedMethods) { Debug.Assert(info.MethodBody.RVA == info.RVA); if (info.MethodBody.RVA != info.RVA) throw new InvalidOperationException(); writer.Position = destStreamBaseOffset + (uint)info.MethodBody.FileOffset; info.MethodBody.VerifyWriteTo(writer); } } /// public void SetOffset(FileOffset offset, RVA rva) { setOffsetCalled = true; this.offset = offset; this.rva = rva; tinyMethodsDict = null; fatMethodsDict = null; var rva2 = rva; foreach (var mb in tinyMethods) { mb.SetOffset(offset, rva2); uint len = mb.GetFileLength(); rva2 += len; offset += len; } foreach (var mb in fatMethods) { if (alignFatBodies) { uint padding = (uint)rva2.AlignUp(FAT_BODY_ALIGNMENT) - (uint)rva2; rva2 += padding; offset += padding; } mb.SetOffset(offset, rva2); uint len = mb.GetFileLength(); rva2 += len; offset += len; } length = (uint)rva2 - (uint)rva; } /// public uint GetFileLength() => length; /// public uint GetVirtualSize() => GetFileLength(); /// public uint CalculateAlignment() => 0; /// public void WriteTo(DataWriter writer) { var rva2 = rva; foreach (var mb in tinyMethods) { mb.VerifyWriteTo(writer); rva2 += mb.GetFileLength(); } foreach (var mb in fatMethods) { if (alignFatBodies) { int padding = (int)rva2.AlignUp(FAT_BODY_ALIGNMENT) - (int)rva2; writer.WriteZeroes(padding); rva2 += (uint)padding; } mb.VerifyWriteTo(writer); rva2 += mb.GetFileLength(); } } } }