Reorganize media assignment to correctly place facade order optimization

6 files changed, 135 insertions, 124 deletions

diff --git a/src/WixToolset.Core.WindowsInstaller/Bind/AssignMediaCommand.cs b/src/WixToolset.Core.WindowsInstaller/Bind/AssignMediaCommand.cs
index 1d677a70..b75956b4 100644
--- a/src/WixToolset.Core.WindowsInstaller/Bind/AssignMediaCommand.cs
+++ b/src/WixToolset.Core.WindowsInstaller/Bind/AssignMediaCommand.cs
@@ -35,7 +35,7 @@ namespace WixToolset.Core.WindowsInstaller.Bind
 
         private bool FilesCompressed { get; }
 
-        public string CabinetNameTemplate { private get; set; }
+        private string CabinetNameTemplate { get; set; }
 
         /// <summary>
         /// Gets cabinets with their file rows.
@@ -43,11 +43,6 @@ namespace WixToolset.Core.WindowsInstaller.Bind
         public Dictionary<MediaTuple, IEnumerable<FileFacade>> FileFacadesByCabinetMedia { get; private set; }
 
         /// <summary>
-        /// Get media rows.
-        /// </summary>
-        public Dictionary<int, MediaTuple> MediaRows { get; private set; }
-
-        /// <summary>
         /// Get uncompressed file rows. This will contain file rows of File elements that are marked with compression=no.
         /// This contains all the files when Package element is marked with compression=no
         /// </summary>
@@ -55,17 +50,11 @@ namespace WixToolset.Core.WindowsInstaller.Bind
 
         public void Execute()
         {
-            var filesByCabinetMedia = new Dictionary<MediaTuple, List<FileFacade>>();
-
-            var mediaRows = new Dictionary<int, MediaTuple>();
-
-            var uncompressedFiles = new List<FileFacade>();
-
-            var mediaTable = this.Section.Tuples.OfType<MediaTuple>().ToList();
-            var mediaTemplateTable = this.Section.Tuples.OfType<WixMediaTemplateTuple>().ToList();
+            var mediaTuples = this.Section.Tuples.OfType<MediaTuple>().ToList();
+            var mediaTemplateTuples = this.Section.Tuples.OfType<WixMediaTemplateTuple>().ToList();
 
-            // If both tables are authored, it is an error.
-            if (mediaTemplateTable.Count > 0 && mediaTable.Count > 1)
+            // If both tuples are authored, it is an error.
+            if (mediaTemplateTuples.Count > 0 && mediaTuples.Count > 1)
             {
                 throw new WixException(ErrorMessages.MediaTableCollision(null));
             }
@@ -78,34 +67,44 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                     Cabinet = "#MergeModule.CABinet",
                 });
 
-                filesByCabinetMedia.Add(mergeModuleMediaTuple, new List<FileFacade>(this.FileFacades));
+                this.FileFacadesByCabinetMedia = new Dictionary<MediaTuple, IEnumerable<FileFacade>>
+                {
+                    { mergeModuleMediaTuple, this.FileFacades }
+                };
+
+                this.UncompressedFileFacades = Array.Empty<FileFacade>();
             }
-            else if (mediaTemplateTable.Count == 0)
+            else if (mediaTemplateTuples.Count == 0)
             {
-                this.ManuallyAssignFiles(mediaTable, this.FileFacades, filesByCabinetMedia, mediaRows, uncompressedFiles);
+                var filesByCabinetMedia = new Dictionary<MediaTuple, List<FileFacade>>();
+
+                var uncompressedFiles = new List<FileFacade>();
+
+                this.ManuallyAssignFiles(mediaTuples, filesByCabinetMedia, uncompressedFiles);
+
+                this.FileFacadesByCabinetMedia = filesByCabinetMedia.ToDictionary(kvp => kvp.Key, kvp => (IEnumerable<FileFacade>)kvp.Value);
+
+                this.UncompressedFileFacades = uncompressedFiles;
             }
             else
             {
-                this.AutoAssignFiles(mediaTable, filesByCabinetMedia, mediaRows, uncompressedFiles);
-            }
+                var filesByCabinetMedia = new Dictionary<MediaTuple, List<FileFacade>>();
 
-            this.FileFacadesByCabinetMedia = new Dictionary<MediaTuple, IEnumerable<FileFacade>>();
+                var uncompressedFiles = new List<FileFacade>();
 
-            foreach (var mediaRowWithFiles in filesByCabinetMedia)
-            {
-                this.FileFacadesByCabinetMedia.Add(mediaRowWithFiles.Key, mediaRowWithFiles.Value);
-            }
+                this.AutoAssignFiles(mediaTuples, filesByCabinetMedia, uncompressedFiles);
 
-            this.MediaRows = mediaRows;
+                this.FileFacadesByCabinetMedia = filesByCabinetMedia.ToDictionary(kvp => kvp.Key, kvp => (IEnumerable<FileFacade>)kvp.Value);
 
-            this.UncompressedFileFacades = uncompressedFiles;
+                this.UncompressedFileFacades = uncompressedFiles;
+            }
         }
 
         /// <summary>
         /// Assign files to cabinets based on MediaTemplate authoring.
         /// </summary>
         /// <param name="fileFacades">FileRowCollection</param>
-        private void AutoAssignFiles(List<MediaTuple> mediaTable, Dictionary<MediaTuple, List<FileFacade>> filesByCabinetMedia, Dictionary<int, MediaTuple> mediaRows, List<FileFacade> uncompressedFiles)
+        private void AutoAssignFiles(List<MediaTuple> mediaTable, Dictionary<MediaTuple, List<FileFacade>> filesByCabinetMedia, List<FileFacade> uncompressedFiles)
         {
             const int MaxCabIndex = 999;
 
@@ -158,6 +157,8 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                 throw new WixException(ErrorMessages.MaximumUncompressedMediaSizeTooLarge(null, maxPreCabSizeInMB));
             }
 
+            var mediaTuplesByDiskId = new Dictionary<int, MediaTuple>();
+
             foreach (var facade in this.FileFacades)
             {
                 // When building a product, if the current file is not to be compressed or if
@@ -171,44 +172,38 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                 if (currentCabIndex == MaxCabIndex)
                 {
                     // Associate current file with last cab (irrespective of the size) and cab index is not incremented anymore.
-                    var cabinetFiles = filesByCabinetMedia[currentMediaRow];
-                    facade.DiskId = currentCabIndex;
-                    cabinetFiles.Add(facade);
-                    continue;
-                }
-
-                // Update current cab size.
-                currentPreCabSize += (ulong)facade.FileSize;
-
-                if (currentPreCabSize > maxPreCabSizeInBytes)
-                {
-                    // Overflow due to current file
-                    currentMediaRow = this.AddMediaRow(mediaTemplateRow, ++currentCabIndex);
-                    mediaRows.Add(currentMediaRow.DiskId, currentMediaRow);
-                    filesByCabinetMedia.Add(currentMediaRow, new List<FileFacade>());
-
-                    var cabinetFileRows = filesByCabinetMedia[currentMediaRow];
-                    facade.DiskId = currentCabIndex;
-                    cabinetFileRows.Add(facade);
-                    // Now files larger than MaxUncompressedMediaSize will be the only file in its cabinet so as to respect MaxUncompressedMediaSize
-                    currentPreCabSize = (ulong)facade.FileSize;
                 }
                 else
                 {
-                    // File fits in the current cab.
-                    if (currentMediaRow == null)
+                    // Update current cab size.
+                    currentPreCabSize += (ulong)facade.FileSize;
+
+                    // Overflow due to current file
+                    if (currentPreCabSize > maxPreCabSizeInBytes)
                     {
-                        // Create new cab and MediaRow
-                        currentMediaRow = this.AddMediaRow(mediaTemplateRow, ++currentCabIndex);
-                        mediaRows.Add(currentMediaRow.DiskId, currentMediaRow);
+                        currentMediaRow = this.AddMediaTuple(mediaTemplateRow, ++currentCabIndex);
+                        mediaTuplesByDiskId.Add(currentMediaRow.DiskId, currentMediaRow);
                         filesByCabinetMedia.Add(currentMediaRow, new List<FileFacade>());
-                    }
 
-                    // Associate current file with current cab.
-                    var cabinetFiles = filesByCabinetMedia[currentMediaRow];
-                    facade.DiskId = currentCabIndex;
-                    cabinetFiles.Add(facade);
+                        // Now files larger than MaxUncompressedMediaSize will be the only file in its cabinet so as to respect MaxUncompressedMediaSize
+                        currentPreCabSize = (ulong)facade.FileSize;
+                    }
+                    else // file fits in the current cab.
+                    {
+                        if (currentMediaRow == null)
+                        {
+                            // Create new cab and MediaRow
+                            currentMediaRow = this.AddMediaTuple(mediaTemplateRow, ++currentCabIndex);
+                            mediaTuplesByDiskId.Add(currentMediaRow.DiskId, currentMediaRow);
+                            filesByCabinetMedia.Add(currentMediaRow, new List<FileFacade>());
+                        }
+                    }
                 }
+
+                // Associate current file with current cab.
+                var cabinetFiles = filesByCabinetMedia[currentMediaRow];
+                facade.DiskId = currentCabIndex;
+                cabinetFiles.Add(facade);
             }
 
             // If there are uncompressed files and no MediaRow, create a default one.
@@ -219,51 +214,45 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                     DiskId = 1,
                 });
 
-                mediaRows.Add(1, defaultMediaRow);
+                mediaTuplesByDiskId.Add(1, defaultMediaRow);
             }
         }
 
         /// <summary>
         /// Assign files to cabinets based on Media authoring.
         /// </summary>
-        /// <param name="mediaTable"></param>
-        /// <param name="fileFacades"></param>
-        private void ManuallyAssignFiles(List<MediaTuple> mediaTable, IEnumerable<FileFacade> fileFacades, Dictionary<MediaTuple, List<FileFacade>> filesByCabinetMedia, Dictionary<int, MediaTuple> mediaRows, List<FileFacade> uncompressedFiles)
+        private void ManuallyAssignFiles(List<MediaTuple> mediaTuples, Dictionary<MediaTuple, List<FileFacade>> filesByCabinetMedia, List<FileFacade> uncompressedFiles)
         {
-            if (mediaTable.Any())
+            var mediaTuplesByDiskId = new Dictionary<int, MediaTuple>();
+
+            if (mediaTuples.Any())
             {
-                var cabinetMediaRows = new Dictionary<string, MediaTuple>(StringComparer.OrdinalIgnoreCase);
-                foreach (var mediaRow in mediaTable)
+                var cabinetMediaTuples = new Dictionary<string, MediaTuple>(StringComparer.OrdinalIgnoreCase);
+                foreach (var mediaTuple in mediaTuples)
                 {
                     // If the Media row has a cabinet, make sure it is unique across all Media rows.
-                    if (!String.IsNullOrEmpty(mediaRow.Cabinet))
+                    if (!String.IsNullOrEmpty(mediaTuple.Cabinet))
                     {
-                        if (cabinetMediaRows.TryGetValue(mediaRow.Cabinet, out var existingRow))
+                        if (cabinetMediaTuples.TryGetValue(mediaTuple.Cabinet, out var existingRow))
                         {
-                            this.Messaging.Write(ErrorMessages.DuplicateCabinetName(mediaRow.SourceLineNumbers, mediaRow.Cabinet));
+                            this.Messaging.Write(ErrorMessages.DuplicateCabinetName(mediaTuple.SourceLineNumbers, mediaTuple.Cabinet));
                             this.Messaging.Write(ErrorMessages.DuplicateCabinetName2(existingRow.SourceLineNumbers, existingRow.Cabinet));
                         }
                         else
                         {
-                            cabinetMediaRows.Add(mediaRow.Cabinet, mediaRow);
+                            cabinetMediaTuples.Add(mediaTuple.Cabinet, mediaTuple);
                         }
-                    }
 
-                    mediaRows.Add(mediaRow.DiskId, mediaRow);
-                }
-            }
+                        filesByCabinetMedia.Add(mediaTuple, new List<FileFacade>());
+                    }
 
-            foreach (var mediaRow in mediaRows.Values)
-            {
-                if (null != mediaRow.Cabinet)
-                {
-                    filesByCabinetMedia.Add(mediaRow, new List<FileFacade>());
+                    mediaTuplesByDiskId.Add(mediaTuple.DiskId, mediaTuple);
                 }
             }
 
-            foreach (var facade in fileFacades)
+            foreach (var facade in this.FileFacades)
             {
-                if (!mediaRows.TryGetValue(facade.DiskId, out var mediaRow))
+                if (!mediaTuplesByDiskId.TryGetValue(facade.DiskId, out var mediaTuple))
                 {
                     this.Messaging.Write(ErrorMessages.MissingMedia(facade.SourceLineNumber, facade.DiskId));
                     continue;
@@ -279,7 +268,7 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                 }
                 else // file is marked compressed.
                 {
-                    if (filesByCabinetMedia.TryGetValue(mediaRow, out var cabinetFiles))
+                    if (filesByCabinetMedia.TryGetValue(mediaTuple, out var cabinetFiles))
                     {
                         cabinetFiles.Add(facade);
                     }
@@ -292,12 +281,12 @@ namespace WixToolset.Core.WindowsInstaller.Bind
         }
 
         /// <summary>
-        /// Adds a row to the media table with cab name template filled in.
+        /// Adds a tuple to the section with cab name template filled in.
         /// </summary>
         /// <param name="mediaTable"></param>
         /// <param name="cabIndex"></param>
         /// <returns></returns>
-        private MediaTuple AddMediaRow(WixMediaTemplateTuple mediaTemplateTuple, int cabIndex)
+        private MediaTuple AddMediaTuple(WixMediaTemplateTuple mediaTemplateTuple, int cabIndex)
         {
             return this.Section.AddTuple(new MediaTuple(mediaTemplateTuple.SourceLineNumbers, new Identifier(AccessModifier.Private, cabIndex))
             {
diff --git a/src/WixToolset.Core.WindowsInstaller/Bind/BindDatabaseCommand.cs b/src/WixToolset.Core.WindowsInstaller/Bind/BindDatabaseCommand.cs
index 6c2968ec..da92be69 100644
--- a/src/WixToolset.Core.WindowsInstaller/Bind/BindDatabaseCommand.cs
+++ b/src/WixToolset.Core.WindowsInstaller/Bind/BindDatabaseCommand.cs
@@ -281,19 +281,6 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                 command.Execute();
             }
 
-            // Assign files to media.
-            Dictionary<int, MediaTuple> assignedMediaRows;
-            Dictionary<MediaTuple, IEnumerable<FileFacade>> filesByCabinetMedia;
-            IEnumerable<FileFacade> uncompressedFiles;
-            {
-                var command = new AssignMediaCommand(section, this.Messaging, fileFacades, compressed);
-                command.Execute();
-
-                assignedMediaRows = command.MediaRows;
-                filesByCabinetMedia = command.FileFacadesByCabinetMedia;
-                uncompressedFiles = command.UncompressedFileFacades;
-            }
-
             // stop processing if an error previously occurred
             if (this.Messaging.EncounteredError)
             {
@@ -366,10 +353,23 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                 command.Execute();
             }
 
-            // Update file sequence.
+            // Assign files to media and update file sequences.
+            Dictionary<MediaTuple, IEnumerable<FileFacade>> filesByCabinetMedia;
+            IEnumerable<FileFacade> uncompressedFiles;
             {
-                var command = new UpdateMediaSequencesCommand(section, fileFacades);
-                command.Execute();
+                var order = new OptimizeFileFacadesOrderCommand(fileFacades);
+                order.Execute();
+
+                fileFacades = order.FileFacades;
+
+                var assign = new AssignMediaCommand(section, this.Messaging, fileFacades, compressed);
+                assign.Execute();
+
+                filesByCabinetMedia = assign.FileFacadesByCabinetMedia;
+                uncompressedFiles = assign.UncompressedFileFacades;
+
+                var update = new UpdateMediaSequencesCommand(section, fileFacades);
+                update.Execute();
             }
 
             // stop processing if an error previously occurred
diff --git a/src/WixToolset.Core.WindowsInstaller/Bind/CabinetBuilder.cs b/src/WixToolset.Core.WindowsInstaller/Bind/CabinetBuilder.cs
index 486ee67a..dce89f78 100644
--- a/src/WixToolset.Core.WindowsInstaller/Bind/CabinetBuilder.cs
+++ b/src/WixToolset.Core.WindowsInstaller/Bind/CabinetBuilder.cs
@@ -7,7 +7,6 @@ namespace WixToolset.Core.WindowsInstaller.Bind
     using System.IO;
     using System.Linq;
     using System.Threading;
-    using WixToolset.Core.Bind;
     using WixToolset.Core.Native;
     using WixToolset.Data;
     using WixToolset.Extensibility.Services;
@@ -18,12 +17,13 @@ namespace WixToolset.Core.WindowsInstaller.Bind
     /// </summary>
     internal sealed class CabinetBuilder
     {
-        private Queue cabinetWorkItems;
-        private object lockObject;
+        private readonly object lockObject = new object();
+
+        private readonly Queue cabinetWorkItems;
         private int threadCount;
 
         // Address of Binder's callback function for Cabinet Splitting
-        private IntPtr newCabNamesCallBackAddress;
+        private readonly IntPtr newCabNamesCallBackAddress;
 
         /// <summary>
         /// Instantiate a new CabinetBuilder.
@@ -38,7 +38,6 @@ namespace WixToolset.Core.WindowsInstaller.Bind
             }
 
             this.cabinetWorkItems = new Queue();
-            this.lockObject = new object();
             this.Messaging = messaging;
             this.threadCount = threadCount;
 
@@ -56,10 +55,7 @@ namespace WixToolset.Core.WindowsInstaller.Bind
         /// Enqueues a CabinetWorkItem to the queue.
         /// </summary>
         /// <param name="cabinetWorkItem">cabinet work item</param>
-        public void Enqueue(CabinetWorkItem cabinetWorkItem)
-        {
-            this.cabinetWorkItems.Enqueue(cabinetWorkItem);
-        }
+        public void Enqueue(CabinetWorkItem cabinetWorkItem) => this.cabinetWorkItems.Enqueue(cabinetWorkItem);
 
         /// <summary>
         /// Create the queued cabinets.
diff --git a/src/WixToolset.Core.WindowsInstaller/Bind/CreateCabinetsCommand.cs b/src/WixToolset.Core.WindowsInstaller/Bind/CreateCabinetsCommand.cs
index de357e53..9741fcd9 100644
--- a/src/WixToolset.Core.WindowsInstaller/Bind/CreateCabinetsCommand.cs
+++ b/src/WixToolset.Core.WindowsInstaller/Bind/CreateCabinetsCommand.cs
@@ -171,7 +171,6 @@ namespace WixToolset.Core.WindowsInstaller.Bind
             return cabbingThreadCount;
         }
 
-
         /// <summary>
         /// Creates a work item to create a cabinet.
         /// </summary>
diff --git a/src/WixToolset.Core.WindowsInstaller/Bind/OptimizeFileFacadesOrderCommand.cs b/src/WixToolset.Core.WindowsInstaller/Bind/OptimizeFileFacadesOrderCommand.cs
new file mode 100644
index 00000000..6943d345
--- /dev/null
+++ b/src/WixToolset.Core.WindowsInstaller/Bind/OptimizeFileFacadesOrderCommand.cs
@@ -0,0 +1,38 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.WindowsInstaller.Bind
+{
+    using System;
+    using System.Collections.Generic;
+    using WixToolset.Core.Bind;
+
+    internal class OptimizeFileFacadesOrderCommand
+    {
+        public OptimizeFileFacadesOrderCommand(List<FileFacade> fileFacades)
+        {
+            this.FileFacades = fileFacades;
+        }
+
+        public List<FileFacade> FileFacades { get; private set; }
+
+        public List<FileFacade> Execute()
+        {
+            this.FileFacades.Sort(FileFacadeOptimizer.Instance);
+
+            return this.FileFacades;
+        }
+
+        private class FileFacadeOptimizer : IComparer<FileFacade>
+        {
+            public static readonly FileFacadeOptimizer Instance = new FileFacadeOptimizer();
+
+            public int Compare(FileFacade x, FileFacade y)
+            {
+                // TODO: Sort these facades even smarter by directory path and component id 
+                //       and maybe file size or file extension and other creative ideas to
+                //       get optimal install speed out of MSI.
+                return String.Compare(x.ComponentRef, y.ComponentRef, StringComparison.Ordinal);
+            }
+        }
+    }
+}
diff --git a/src/WixToolset.Core.WindowsInstaller/Bind/UpdateMediaSequencesCommand.cs b/src/WixToolset.Core.WindowsInstaller/Bind/UpdateMediaSequencesCommand.cs
index 9aab7b98..bf28b279 100644
--- a/src/WixToolset.Core.WindowsInstaller/Bind/UpdateMediaSequencesCommand.cs
+++ b/src/WixToolset.Core.WindowsInstaller/Bind/UpdateMediaSequencesCommand.cs
@@ -29,9 +29,7 @@ namespace WixToolset.Core.WindowsInstaller.Bind
             {
                 var lastSequence = 0;
 
-                // Order by Component to group the files by directory.
-                var optimized = this.OptimizedFileFacades();
-                foreach (var facade in optimized)
+                foreach (var facade in this.FileFacades)
                 {
                     facade.Sequence = ++lastSequence;
                 }
@@ -43,8 +41,7 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                 var patchGroups = new Dictionary<int, List<FileFacade>>();
 
                 // sequence the non-patch-added files
-                var optimized = this.OptimizedFileFacades();
-                foreach (var facade in optimized)
+                foreach (var facade in this.FileFacades)
                 {
                     if (null == mediaTuple)
                     {
@@ -108,13 +105,5 @@ namespace WixToolset.Core.WindowsInstaller.Bind
                 }
             }
         }
-
-        private IEnumerable<FileFacade> OptimizedFileFacades()
-        {
-            // TODO: Sort these facades even smarter by directory path and component id 
-            //       and maybe file size or file extension and other creative ideas to
-            //       get optimal install speed out of MSI.
-            return this.FileFacades.OrderBy(f => f.ComponentRef);
-        }
     }
 }