diff --git a/Wabbajack.Lib/CompilationSteps/IncludePatches.cs b/Wabbajack.Lib/CompilationSteps/IncludePatches.cs index eeaddaf0..9f7fe5da 100644 --- a/Wabbajack.Lib/CompilationSteps/IncludePatches.cs +++ b/Wabbajack.Lib/CompilationSteps/IncludePatches.cs @@ -6,6 +6,7 @@ using Alphaleonis.Win32.Filesystem; using F23.StringSimilarity; using Newtonsoft.Json; using Wabbajack.Common; +using Wabbajack.Lib.Downloaders; using Wabbajack.VirtualFileSystem; namespace Wabbajack.Lib.CompilationSteps @@ -104,7 +105,7 @@ namespace Wabbajack.Lib.CompilationSteps if (patches.All(p => p.Item1)) { - var (_, bytes, file) = patches.OrderBy(f => f.data!.Length).First(); + var (_, bytes, file) = PickPatch(_mo2Compiler, patches); e.FromHash = file.Hash; e.ArchiveHashPath = file.MakeRelativePaths(); e.PatchID = await _compiler.IncludeFile(bytes!); @@ -125,6 +126,29 @@ namespace Wabbajack.Lib.CompilationSteps return e; } + public static (bool, byte[], VirtualFile) PickPatch(MO2Compiler mo2Compiler, IEnumerable<(bool foundHash, byte[]? data, VirtualFile file)> patches) + { + var ordered = patches + .Select(f => (f.foundHash, f.data!, f.file)) + .OrderBy(f => f.Item2.Length) + .ToArray(); + + var primaryChoice = ordered.FirstOrDefault(itm => + { + var baseHash = itm.file.TopParent.Hash; + + // If this file doesn't come from a game use it + if (!mo2Compiler.GamesWithHashes.TryGetValue(baseHash, out var games)) + return true; + + // Otherwise skip files that are not from the primary game + return games.Contains(mo2Compiler.CompilingGame.Game); + }); + + // If we didn't find a file from an archive or the primary game, use a secondary game file. + return primaryChoice != default ? primaryChoice : ordered.FirstOrDefault(); + } + private AbsolutePath ModForFile(AbsolutePath file) { return file.RelativeTo(((MO2Compiler)_compiler).MO2ModsFolder).TopParent diff --git a/Wabbajack.Lib/MO2Compiler.cs b/Wabbajack.Lib/MO2Compiler.cs index 76ff79f2..708a2f2e 100644 --- a/Wabbajack.Lib/MO2Compiler.cs +++ b/Wabbajack.Lib/MO2Compiler.cs @@ -210,6 +210,7 @@ namespace Wabbajack.Lib { var files = await ClientAPI.GetExistingGameFiles(Queue, ag); Utils.Log($"Including {files.Length} stock game files from {ag} as download sources"); + GameHashes[ag] = files.Select(f => f.Hash).ToHashSet(); IndexedArchives.AddRange(files.Select(f => { @@ -228,6 +229,10 @@ namespace Wabbajack.Lib Utils.Error(e, "Unable to find existing game files, skipping."); } } + + GamesWithHashes = GameHashes.SelectMany(g => g.Value.Select(h => (g, h))) + .GroupBy(gh => gh.h) + .ToDictionary(gh => gh.Key, gh => gh.Select(p => p.g.Key).ToArray()); } IndexedArchives = IndexedArchives.DistinctBy(a => a.File.AbsoluteName).ToList(); @@ -364,6 +369,9 @@ namespace Wabbajack.Lib } + public Dictionary> GameHashes { get; set; } = new Dictionary>(); + public Dictionary GamesWithHashes { get; set; } = new Dictionary(); + public bool UseGamePaths { get; set; } = true; private async Task CleanInvalidArchivesAndFillState() @@ -508,7 +516,7 @@ namespace Wabbajack.Lib if (patches.All(p => p.Item1)) { - var (_, bytes, file) = patches.OrderBy(f => f.data!.Length).First(); + var (_, bytes, file) = IncludePatches.PickPatch(this, patches); pfa.FromFile = file; pfa.FromHash = file.Hash; pfa.ArchiveHashPath = file.MakeRelativePaths();