4 files changed, 120 insertions, 60 deletions

diff --git a/src/WixToolset.Core.Burn/Bind/BindBundleCommand.cs b/src/WixToolset.Core.Burn/Bind/BindBundleCommand.cs
index 83cc0a67..b88251af 100644
--- a/src/WixToolset.Core.Burn/Bind/BindBundleCommand.cs
+++ b/src/WixToolset.Core.Burn/Bind/BindBundleCommand.cs
@@ -143,27 +143,6 @@ namespace WixToolset.Core.Burn
             var payloadSymbols = section.Symbols.OfType<WixBundlePayloadSymbol>().ToDictionary(t => t.Id.Id);
             var packagesPayloads = RecalculatePackagesPayloads(payloadSymbols, wixGroupSymbols);
 
-            // Update explicitly authored payloads with their parent container
-            // to make it easier to gather the payloads later.
-            foreach (var groupSymbol in wixGroupSymbols)
-            {
-                if (ComplexReferenceChildType.Payload == groupSymbol.ChildType)
-                {
-                    var payloadSymbol = payloadSymbols[groupSymbol.ChildId];
-
-                    if (ComplexReferenceParentType.Container == groupSymbol.ParentType)
-                    {
-                        // TODO: v3 didn't warn if we overwrote the payload's container.
-                        // Should we warn now?
-                        payloadSymbol.ContainerRef = groupSymbol.ParentId;
-                    }
-                    else if (ComplexReferenceParentType.Layout == groupSymbol.ParentType)
-                    {
-                        payloadSymbol.LayoutOnly = true;
-                    }
-                }
-            }
-
             var layoutDirectory = Path.GetDirectoryName(this.OutputPath);
 
             // Process the explicitly authored payloads.
diff --git a/src/WixToolset.Core/Link/FlattenAndProcessBundleTablesCommand.cs b/src/WixToolset.Core/Link/FlattenAndProcessBundleTablesCommand.cs
new file mode 100644
index 00000000..0fa48d8c
--- /dev/null
+++ b/src/WixToolset.Core/Link/FlattenAndProcessBundleTablesCommand.cs
@@ -0,0 +1,91 @@
+// 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.Link
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Linq;
+    using WixToolset.Data;
+    using WixToolset.Data.Burn;
+    using WixToolset.Data.Symbols;
+    using WixToolset.Extensibility.Services;
+
+    internal class FlattenAndProcessBundleTablesCommand
+    {
+        public FlattenAndProcessBundleTablesCommand(IntermediateSection entrySection, IMessaging messaging)
+        {
+            this.EntrySection = entrySection;
+            this.Messaging = messaging;
+        }
+
+        private IntermediateSection EntrySection { get; }
+
+        private IMessaging Messaging { get; }
+
+        public void Execute()
+        {
+            this.FlattenBundleTables();
+
+            if (this.Messaging.EncounteredError)
+            {
+                return;
+            }
+
+            this.ProcessBundleComplexReferences();
+        }
+
+        /// <summary>
+        /// Flattens the tables used in a Bundle.
+        /// </summary>
+        private void FlattenBundleTables()
+        {
+            // We need to flatten the nested PayloadGroups and PackageGroups under
+            // UX, Chain, and any Containers. When we're done, the WixGroups table
+            // will hold Payloads under UX, ChainPackages (references?) under Chain,
+            // and ChainPackages/Payloads under the attached and any detatched
+            // Containers.
+            var groups = new WixGroupingOrdering(this.EntrySection, this.Messaging);
+
+            // Create UX payloads and Package payloads
+            groups.UseTypes(new[] { ComplexReferenceParentType.Container, ComplexReferenceParentType.Layout, ComplexReferenceParentType.PackageGroup, ComplexReferenceParentType.PayloadGroup, ComplexReferenceParentType.Package },
+                            new[] { ComplexReferenceChildType.PackageGroup, ComplexReferenceChildType.Package, ComplexReferenceChildType.PackagePayload, ComplexReferenceChildType.PayloadGroup, ComplexReferenceChildType.Payload });
+            groups.FlattenAndRewriteGroups(ComplexReferenceParentType.Package, false);
+            groups.FlattenAndRewriteGroups(ComplexReferenceParentType.Container, false);
+            groups.FlattenAndRewriteGroups(ComplexReferenceParentType.Layout, false);
+
+            // Create Chain packages...
+            groups.UseTypes(new[] { ComplexReferenceParentType.PackageGroup }, new[] { ComplexReferenceChildType.Package, ComplexReferenceChildType.PackageGroup });
+            groups.FlattenAndRewriteRows(ComplexReferenceChildType.PackageGroup, "WixChain", false);
+
+            groups.RemoveUsedGroupRows();
+        }
+
+        private void ProcessBundleComplexReferences()
+        {
+            var groups = this.EntrySection.Symbols.OfType<WixGroupSymbol>().ToList();
+            var payloadsById = this.EntrySection.Symbols.OfType<WixBundlePayloadSymbol>().ToDictionary(c => c.Id.Id);
+
+            // Assign authored payloads to authored containers.
+            // Compressed Payloads not assigned to a container here will get assigned to the default attached container during binding.
+            foreach (var groupSymbol in groups)
+            {
+                if (ComplexReferenceChildType.Payload == groupSymbol.ChildType)
+                {
+                    var payloadSymbol = payloadsById[groupSymbol.ChildId];
+
+                    if (ComplexReferenceParentType.Container == groupSymbol.ParentType)
+                    {
+                        // TODO: v3 didn't warn if we overwrote the payload's container.
+                        // Should we warn now?
+                        payloadSymbol.ContainerRef = groupSymbol.ParentId;
+                    }
+                    else if (ComplexReferenceParentType.Layout == groupSymbol.ParentType)
+                    {
+                        payloadSymbol.LayoutOnly = true;
+                    }
+                }
+            }
+
+        }
+    }
+}
diff --git a/src/WixToolset.Core/Linker.cs b/src/WixToolset.Core/Linker.cs
index bf7130db..47671f26 100644
--- a/src/WixToolset.Core/Linker.cs
+++ b/src/WixToolset.Core/Linker.cs
@@ -312,7 +312,11 @@ namespace WixToolset.Core
                 }
 
                 // Bundles have groups of data that must be flattened in a way different from other types.
-                this.FlattenBundleTables(resolvedSection);
+                if (resolvedSection.Type == SectionType.Bundle)
+                {
+                    var command = new FlattenAndProcessBundleTablesCommand(resolvedSection, this.Messaging);
+                    command.Execute();
+                }
 
                 if (this.Messaging.EncounteredError)
                 {
@@ -854,37 +858,6 @@ namespace WixToolset.Core
         */
 
         /// <summary>
-        /// Flattens the tables used in a Bundle.
-        /// </summary>
-        /// <param name="entrySection">Output containing the tables to process.</param>
-        private void FlattenBundleTables(IntermediateSection entrySection)
-        {
-            if (SectionType.Bundle != entrySection.Type)
-            {
-                return;
-            }
-
-            // We need to flatten the nested PayloadGroups and PackageGroups under
-            // UX, Chain, and any Containers. When we're done, the WixGroups table
-            // will hold Payloads under UX, ChainPackages (references?) under Chain,
-            // and ChainPackages/Payloads under the attached and any detatched
-            // Containers.
-            var groups = new WixGroupingOrdering(entrySection, this.Messaging);
-
-            // Create UX payloads and Package payloads
-            groups.UseTypes(new[] { ComplexReferenceParentType.Container, ComplexReferenceParentType.Layout, ComplexReferenceParentType.PackageGroup, ComplexReferenceParentType.PayloadGroup, ComplexReferenceParentType.Package }, new[] { ComplexReferenceChildType.PackageGroup, ComplexReferenceChildType.Package, ComplexReferenceChildType.PackagePayload, ComplexReferenceChildType.PayloadGroup, ComplexReferenceChildType.Payload });
-            groups.FlattenAndRewriteGroups(ComplexReferenceParentType.Package, false);
-            groups.FlattenAndRewriteGroups(ComplexReferenceParentType.Container, false);
-            groups.FlattenAndRewriteGroups(ComplexReferenceParentType.Layout, false);
-
-            // Create Chain packages...
-            groups.UseTypes(new[] { ComplexReferenceParentType.PackageGroup }, new[] { ComplexReferenceChildType.Package, ComplexReferenceChildType.PackageGroup });
-            groups.FlattenAndRewriteRows(ComplexReferenceChildType.PackageGroup, "WixChain", false);
-
-            groups.RemoveUsedGroupRows();
-        }
-
-        /// <summary>
         /// Resolves the features connected to other features in the active output.
         /// </summary>
         /// <param name="featuresToFeatures">Feature to feature complex references.</param>
diff --git a/src/test/WixToolsetTest.CoreIntegration/ContainerFixture.cs b/src/test/WixToolsetTest.CoreIntegration/ContainerFixture.cs
index 43fa3f55..f24429f7 100644
--- a/src/test/WixToolsetTest.CoreIntegration/ContainerFixture.cs
+++ b/src/test/WixToolsetTest.CoreIntegration/ContainerFixture.cs
@@ -142,17 +142,34 @@ namespace WixToolsetTest.CoreIntegration
                 var extractResult = BundleExtractor.ExtractBAContainer(null, bundlePath, baFolderPath, extractFolderPath);
                 extractResult.AssertSuccess();
 
-                var payloads = extractResult.SelectManifestNodes("/burn:BurnManifest/burn:Chain/burn:MsiPackage/burn:PayloadRef")
+                var ignoreAttributes = new Dictionary<string, List<string>>
+                {
+                    { "MsiPackage", new List<string> { "CacheId", "InstallSize", "Size", "ProductCode" } },
+                    { "Provides", new List<string> { "Key" } },
+                };
+                var msiPackages = extractResult.SelectManifestNodes("/burn:BurnManifest/burn:Chain/burn:MsiPackage")
                                             .Cast<XmlElement>()
-                                            .Select(e => e.GetTestXml())
+                                            .Select(e => e.GetTestXml(ignoreAttributes))
                                             .ToArray();
                 WixAssert.CompareLineByLine(new string[]
                 {
-                    "<PayloadRef Id='FirstX86.msi' />",
-                    "<PayloadRef Id='fk1m38Cf9RZ2Bx_ipinRY6BftelU' />",
-                    "<PayloadRef Id='FirstX64.msi' />",
-                    "<PayloadRef Id='fC0n41rZK8oW3JK8LzHu6AT3CjdQ' />",
-                }, payloads);
+                    "<MsiPackage Id='FirstX86.msi' Cache='keep' CacheId='*' InstallSize='*' Size='*' PerMachine='yes' Permanent='no' Vital='yes' RollbackBoundaryForward='WixDefaultBoundary' LogPathVariable='WixBundleLog_FirstX86.msi' RollbackLogPathVariable='WixBundleRollbackLog_FirstX86.msi' ProductCode='*' Language='1033' Version='1.0.0.0' UpgradeCode='{12E4699F-E774-4D05-8A01-5BDD41BBA127}'>" +
+                      "<MsiProperty Id='ARPSYSTEMCOMPONENT' Value='1' />" +
+                      "<Provides Key='*' Version='1.0.0.0' DisplayName='MsiPackage' />" +
+                      "<RelatedPackage Id='{12E4699F-E774-4D05-8A01-5BDD41BBA127}' MaxVersion='1.0.0.0' MaxInclusive='no' OnlyDetect='no' LangInclusive='no'><Language Id='1033' /></RelatedPackage>" +
+                      "<RelatedPackage Id='{12E4699F-E774-4D05-8A01-5BDD41BBA127}' MinVersion='1.0.0.0' MinInclusive='no' OnlyDetect='yes' LangInclusive='no'><Language Id='1033' /></RelatedPackage>" +
+                      "<PayloadRef Id='FirstX86.msi' />" +
+                      "<PayloadRef Id='fk1m38Cf9RZ2Bx_ipinRY6BftelU' />" +
+                    "</MsiPackage>",
+                    "<MsiPackage Id='FirstX64.msi' Cache='keep' CacheId='*' InstallSize='*' Size='*' PerMachine='yes' Permanent='no' Vital='yes' RollbackBoundaryBackward='WixDefaultBoundary' LogPathVariable='WixBundleLog_FirstX64.msi' RollbackLogPathVariable='WixBundleRollbackLog_FirstX64.msi' ProductCode='*' Language='1033' Version='1.0.0.0' UpgradeCode='{12E4699F-E774-4D05-8A01-5BDD41BBA127}'>" +
+                      "<MsiProperty Id='ARPSYSTEMCOMPONENT' Value='1' />" +
+                      "<Provides Key='*' Version='1.0.0.0' DisplayName='MsiPackage' />" +
+                      "<RelatedPackage Id='{12E4699F-E774-4D05-8A01-5BDD41BBA127}' MaxVersion='1.0.0.0' MaxInclusive='no' OnlyDetect='no' LangInclusive='no'><Language Id='1033' /></RelatedPackage>" +
+                      "<RelatedPackage Id='{12E4699F-E774-4D05-8A01-5BDD41BBA127}' MinVersion='1.0.0.0' MinInclusive='no' OnlyDetect='yes' LangInclusive='no'><Language Id='1033' /></RelatedPackage>" +
+                      "<PayloadRef Id='FirstX64.msi' />" +
+                      "<PayloadRef Id='fC0n41rZK8oW3JK8LzHu6AT3CjdQ' />" +
+                    "</MsiPackage>",
+                }, msiPackages);
             }
         }