From 6958f34c3256bd7e3adb7fd7caff8ce12e6462d9 Mon Sep 17 00:00:00 2001
From: halgari <tbaldridge@gmail.com>
Date: Sat, 24 Sep 2022 16:13:29 -0500
Subject: [PATCH] Several fixes to the compiler performance code

---
 Wabbajack.Common/AsyncParallelExtensions.cs | 41 +++++++++++++++++++++
 Wabbajack.Common/IEnumerableExtensions.cs   | 21 ++++++++++-
 Wabbajack.Compiler/MO2Compiler.cs           |  2 +-
 Wabbajack.Paths.IO/KnownFolders.cs          | 10 ++++-
 Wabbajack.Paths/AbsolutePath.cs             |  4 +-
 5 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/Wabbajack.Common/AsyncParallelExtensions.cs b/Wabbajack.Common/AsyncParallelExtensions.cs
index 260b366b..71ae5032 100644
--- a/Wabbajack.Common/AsyncParallelExtensions.cs
+++ b/Wabbajack.Common/AsyncParallelExtensions.cs
@@ -62,6 +62,47 @@ public static class AsyncParallelExtensions
         }
     }
     
+    /// <summary>
+    /// Faster version of PMapAll for when the function invocation will take a very small amount of time
+    /// batches all the inputs into N groups and executes them all on one task, where N is the number of
+    /// threads supported by the limiter
+    /// </summary>
+    /// <param name="coll"></param>
+    /// <param name="limiter"></param>
+    /// <param name="mapFn"></param>
+    /// <typeparam name="TIn"></typeparam>
+    /// <typeparam name="TJob"></typeparam>
+    /// <typeparam name="TOut"></typeparam>
+    /// <returns></returns>
+    public static async IAsyncEnumerable<TOut> PMapAllBatched<TIn, TJob, TOut>(this IEnumerable<TIn> coll,
+        IResource<TJob> limiter, Func<TIn, Task<TOut>> mapFn)
+    {
+        var asList = coll.ToList();
+        
+        var tasks = new List<Task<List<TOut>>>();
+
+        tasks.AddRange(Enumerable.Range(0, limiter.MaxTasks).Select(i => Task.Run(async () =>
+        {
+            using var job = await limiter.Begin(limiter.Name, asList.Count / limiter.MaxTasks, CancellationToken.None);
+            var list = new List<TOut>();
+            for (var idx = i; idx < asList.Count; idx += limiter.MaxTasks)
+            {
+                job.ReportNoWait(1);
+                list.Add(await mapFn(asList[idx]));
+            }
+
+            return list;
+        })));
+
+        foreach (var result in tasks)
+        {
+            foreach (var itm in (await result))
+            {
+                yield return itm;
+            }
+        }
+    }
+    
     public static async IAsyncEnumerable<TOut> PKeepAll<TIn, TJob, TOut>(this IEnumerable<TIn> coll,
         IResource<TJob> limiter, Func<TIn, Task<TOut>> mapFn)
     where TOut : class
diff --git a/Wabbajack.Common/IEnumerableExtensions.cs b/Wabbajack.Common/IEnumerableExtensions.cs
index 26922597..ee8141fa 100644
--- a/Wabbajack.Common/IEnumerableExtensions.cs
+++ b/Wabbajack.Common/IEnumerableExtensions.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Reflection.Metadata;
 using System.Threading.Tasks;
 
 namespace Wabbajack.Common;
@@ -11,8 +12,9 @@ public static class IEnumerableExtensions
     {
         foreach (var i in coll) f(i);
     }
-    
+
     #region Shuffle
+
     /// https://stackoverflow.com/questions/5807128/an-extension-method-on-ienumerable-needed-for-shuffling
 
     public static IEnumerable<T> Shuffle<T>(this IEnumerable<T> source, Random rng)
@@ -32,6 +34,7 @@ public static class IEnumerableExtensions
             buffer[j] = buffer[i];
         }
     }
+
     #endregion
 
 
@@ -55,6 +58,22 @@ public static class IEnumerableExtensions
         return data;
     }
 
+    public static IEnumerable<IEnumerable<T>> Partition<T>(this IEnumerable<T> coll, int size)
+    {
+        var asList = coll.ToList();
+
+        IEnumerable<T> SkipEnumerable(IList<T> list, int offset, int size)
+        {
+            for (var i = offset; i < list.Count; i += size)
+            {
+                yield return list[i];
+            }
+        }
+
+        return Enumerable.Range(0, size).Select(offset => SkipEnumerable(asList, offset, size));
+    }
+
+
     public static IEnumerable<T> OnEach<T>(this IEnumerable<T> coll, Action<T> fn)
     {
         foreach (var itm in coll)
diff --git a/Wabbajack.Compiler/MO2Compiler.cs b/Wabbajack.Compiler/MO2Compiler.cs
index 614c507d..db9dba88 100644
--- a/Wabbajack.Compiler/MO2Compiler.cs
+++ b/Wabbajack.Compiler/MO2Compiler.cs
@@ -141,7 +141,7 @@ public class MO2Compiler : ACompiler
         var stack = MakeStack();
 
         NextStep("Compiling", "Running Compilation Stack", AllFiles.Count);
-        var results = await AllFiles.PMapAll(CompilerLimiter, f =>
+        var results = await AllFiles.PMapAllBatched(CompilerLimiter, f =>
         {
             UpdateProgress(1);
             return RunStack(stack, f);
diff --git a/Wabbajack.Paths.IO/KnownFolders.cs b/Wabbajack.Paths.IO/KnownFolders.cs
index 59c44222..b242a7d0 100644
--- a/Wabbajack.Paths.IO/KnownFolders.cs
+++ b/Wabbajack.Paths.IO/KnownFolders.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Diagnostics;
 using System.IO;
+using System.Linq;
 using System.Reflection;
 
 namespace Wabbajack.Paths.IO;
@@ -13,8 +14,15 @@ public static class KnownFolders
         {
             var result = Process.GetCurrentProcess().MainModule?.FileName?.ToAbsolutePath() ?? default;
 
+            if (result != default &&
+                result.PathParts.Any(p => p.Equals("TestRunner", StringComparison.CurrentCultureIgnoreCase)))
+            {
+                return Assembly.GetExecutingAssembly().Location.ToAbsolutePath().Parent;
+            }
             
-            if ((result != default && result.Depth > 1 && result.FileName == "dotnet".ToRelativePath()) || Assembly.GetEntryAssembly() != null)
+            
+            if ((result != default && result.Depth > 1 && result.FileName == "dotnet".ToRelativePath())
+                || Assembly.GetEntryAssembly() != null)
             {
                 result = Assembly.GetEntryAssembly()!.Location.ToAbsolutePath();
             }
diff --git a/Wabbajack.Paths/AbsolutePath.cs b/Wabbajack.Paths/AbsolutePath.cs
index 7f849cff..9d9faa04 100644
--- a/Wabbajack.Paths/AbsolutePath.cs
+++ b/Wabbajack.Paths/AbsolutePath.cs
@@ -18,9 +18,11 @@ public struct AbsolutePath : IPath, IComparable<AbsolutePath>, IEquatable<Absolu
     public PathFormat PathFormat { get; }
 
     private int _hashCode = 0;
-    
+
     internal readonly string[] Parts;
 
+    public string[] PathParts => Parts == default ? Array.Empty<string>() : Parts;
+
     public Extension Extension => Extension.FromPath(Parts[^1]);
     public RelativePath FileName => new(Parts[^1..]);