From f4329d17538c14569ab7058c1c378aa5b2297b2c Mon Sep 17 00:00:00 2001
From: Sean Hall <r.sean.hall@gmail.com>
Date: Sun, 15 Jan 2023 20:29:17 -0600
Subject: Implement SourceDir substitution for Payloads.

7160
---
 src/tools/heat/UtilFinalizeHarvesterMutator.cs     | 35 ++++++++++
 .../DirectoryToPayloadGroupFixture.cs              | 80 ++++++++++++++++++++++
 2 files changed, 115 insertions(+)
 create mode 100644 src/tools/test/WixToolsetTest.Heat/DirectoryToPayloadGroupFixture.cs

diff --git a/src/tools/heat/UtilFinalizeHarvesterMutator.cs b/src/tools/heat/UtilFinalizeHarvesterMutator.cs
index 61bb8731..d41d6e1e 100644
--- a/src/tools/heat/UtilFinalizeHarvesterMutator.cs
+++ b/src/tools/heat/UtilFinalizeHarvesterMutator.cs
@@ -25,6 +25,7 @@ namespace WixToolset.Harvesters
         private Hashtable filePaths;
         private ArrayList files;
         private ArrayList registryValues;
+        private List<Wix.Payload> payloads;
         private bool suppressCOMElements;
         private bool suppressVB6COMElements;
         private string preprocessorVariable;
@@ -40,6 +41,7 @@ namespace WixToolset.Harvesters
             this.filePaths = new Hashtable();
             this.files = new ArrayList();
             this.registryValues = new ArrayList();
+            this.payloads = new List<Wix.Payload>();
         }
 
         /// <summary>
@@ -93,6 +95,7 @@ namespace WixToolset.Harvesters
             this.filePaths.Clear();
             this.files.Clear();
             this.registryValues.Clear();
+            this.payloads.Clear();
 
             // index elements in this wix document
             this.IndexElement(wix);
@@ -100,6 +103,7 @@ namespace WixToolset.Harvesters
             this.MutateDirectories();
             this.MutateFiles();
             this.MutateRegistryValues();
+            this.MutatePayloads();
 
             // must occur after all the registry values have been formatted
             this.MutateComponents();
@@ -131,6 +135,10 @@ namespace WixToolset.Harvesters
             {
                 this.registryValues.Add(element);
             }
+            else if (element is Wix.Payload payloadElement)
+            {
+                this.payloads.Add(payloadElement);
+            }
 
             // index the child elements
             if (element is Wix.IParentElement)
@@ -1005,6 +1013,33 @@ namespace WixToolset.Harvesters
             }
         }
 
+        /// <summary>
+        /// Mutate the payloads.
+        /// </summary>
+        private void MutatePayloads()
+        {
+            string sourceDirSubstitution = this.preprocessorVariable;
+            if (sourceDirSubstitution == null)
+            {
+                return;
+            }
+
+            string prefix = "$(";
+            if (sourceDirSubstitution.StartsWith("wix.", StringComparison.Ordinal))
+            {
+                prefix = "!(";
+            }
+            sourceDirSubstitution = String.Concat(prefix, sourceDirSubstitution, ")");
+
+            foreach (var payload in this.payloads)
+            {
+                if (payload.SourceFile != null && payload.SourceFile.StartsWith("SourceDir\\", StringComparison.Ordinal))
+                {
+                    payload.SourceFile = payload.SourceFile.Substring(9).Insert(0, sourceDirSubstitution);
+                }
+            }
+        }
+
         /// <summary>
         /// Mutate an individual registry string, according to a collection of replacement items.
         /// </summary>
diff --git a/src/tools/test/WixToolsetTest.Heat/DirectoryToPayloadGroupFixture.cs b/src/tools/test/WixToolsetTest.Heat/DirectoryToPayloadGroupFixture.cs
new file mode 100644
index 00000000..c9b8ee30
--- /dev/null
+++ b/src/tools/test/WixToolsetTest.Heat/DirectoryToPayloadGroupFixture.cs
@@ -0,0 +1,80 @@
+// 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 WixToolsetTest.Heat
+{
+    using System;
+    using System.IO;
+    using System.Linq;
+    using WixInternal.TestSupport;
+    using Xunit;
+
+    public class DirectoryToPayloadGroupFixture
+    {
+        [Fact]
+        public void CanHarvestSimpleDirectory()
+        {
+            var folder = TestData.Get("TestData", "SingleFile");
+
+            using (var fs = new DisposableFileSystem())
+            {
+                var outputPath = Path.Combine(fs.GetFolder(), "out.wxs");
+
+                var args = new[]
+                {
+                    "dir", folder,
+                    "-generate", "payloadgroup",
+                    "-o", outputPath
+                };
+
+                var result = HeatRunner.Execute(args);
+                result.AssertSuccess();
+
+                var wxs = File.ReadAllLines(outputPath).Select(s => s.Replace("\"", "'")).ToArray();
+                WixAssert.CompareLineByLine(new[]
+                {
+                    "<Wix xmlns='http://wixtoolset.org/schemas/v4/wxs'>",
+                    "    <Fragment>",
+                    "        <PayloadGroup Id='TARGETDIR'>",
+                    "            <Payload SourceFile='SourceDir\\a.txt' />",
+                    "        </PayloadGroup>",
+                    "    </Fragment>",
+                    "</Wix>",
+                }, wxs);
+            }
+        }
+
+        [Fact]
+        public void CanHarvestSimpleDirectoryWithSourceDirSubstitution()
+        {
+            var folder = TestData.Get("TestData", "SingleFile");
+
+            using (var fs = new DisposableFileSystem())
+            {
+                var outputPath = Path.Combine(fs.GetFolder(), "out.wxs");
+
+                var args = new[]
+                {
+                    "dir", folder,
+                    "-generate", "payloadgroup",
+                    "-var", "var.RootDir",
+                    "-o", outputPath
+                };
+
+                var result = HeatRunner.Execute(args);
+                result.AssertSuccess();
+
+                var wxs = File.ReadAllLines(outputPath).Select(s => s.Replace("\"", "'")).ToArray();
+                WixAssert.CompareLineByLine(new[]
+                {
+                    "<Wix xmlns='http://wixtoolset.org/schemas/v4/wxs'>",
+                    "    <Fragment>",
+                    "        <PayloadGroup Id='TARGETDIR'>",
+                    "            <Payload SourceFile='$(var.RootDir)\\a.txt' />",
+                    "        </PayloadGroup>",
+                    "    </Fragment>",
+                    "</Wix>",
+                }, wxs);
+            }
+        }
+    }
+}
-- 
cgit v1.2.3-55-g6feb