Fix building and perf for bundles with 1000s of package and BA payloads

Fixes other half of #5300

9 files changed, 54 insertions, 53 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index a99a3987..c126d63c 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -1175,7 +1175,7 @@ static HRESULT ExtractContainer(
     {
         BOOL fExtracted = FALSE;
 
-        hr = PayloadFindEmbeddedBySourcePath(pContext->pPayloads, sczStreamName, &pExtract);
+        hr = PayloadFindEmbeddedBySourcePath(pContainer->sdhPayloads, sczStreamName, &pExtract);
         if (E_NOTFOUND != hr)
         {
             ExitOnFailure(hr, "Failed to find embedded payload by source path: %ls container: %ls", sczStreamName, pContainer->sczId);
diff --git a/src/burn/engine/burnextension.cpp b/src/burn/engine/burnextension.cpp
index 475df1c5..ee4b1542 100644
--- a/src/burn/engine/burnextension.cpp
+++ b/src/burn/engine/burnextension.cpp
@@ -26,6 +26,7 @@ EXTERN_C HRESULT BurnExtensionParseFromXml(
     IXMLDOMNodeList* pixnNodes = NULL;
     IXMLDOMNode* pixnNode = NULL;
     DWORD cNodes = 0;
+    LPWSTR scz = NULL;
 
     // Select BundleExtension nodes.
     hr = XmlSelectNodes(pixnBundle, L"BundleExtension", &pixnNodes);
@@ -59,11 +60,11 @@ EXTERN_C HRESULT BurnExtensionParseFromXml(
         ExitOnFailure(hr, "Failed to get @Id.");
 
         // @EntryPayloadId
-        hr = XmlGetAttributeEx(pixnNode, L"EntryPayloadId", &pExtension->sczEntryPayloadId);
-        ExitOnFailure(hr, "Failed to get @EntryPayloadId.");
+        hr = XmlGetAttributeEx(pixnNode, L"EntryPayloadSourcePath", &scz);
+        ExitOnFailure(hr, "Failed to get @EntryPayloadSourcePath.");
 
-        hr = PayloadFindById(pBaPayloads, pExtension->sczEntryPayloadId, &pExtension->pEntryPayload);
-        ExitOnFailure(hr, "Failed to find BundleExtension EntryPayload '%ls'.", pExtension->sczEntryPayloadId);
+        hr = PayloadFindEmbeddedBySourcePath(pBaPayloads->sdhPayloads, scz, &pExtension->pEntryPayload);
+        ExitOnFailure(hr, "Failed to find BundleExtension EntryPayload '%ls'.", pExtension->sczId);
 
         // prepare next iteration
         ReleaseNullObject(pixnNode);
@@ -72,6 +73,7 @@ EXTERN_C HRESULT BurnExtensionParseFromXml(
     hr = S_OK;
 
 LExit:
+    ReleaseStr(scz);
     ReleaseObject(pixnNode);
     ReleaseObject(pixnNodes);
 
@@ -92,7 +94,6 @@ EXTERN_C void BurnExtensionUninitialize(
         {
             BURN_EXTENSION* pExtension = &pBurnExtensions->rgExtensions[i];
 
-            ReleaseStr(pExtension->sczEntryPayloadId);
             ReleaseStr(pExtension->sczId);
         }
         MemFree(pBurnExtensions->rgExtensions);
diff --git a/src/burn/engine/burnextension.h b/src/burn/engine/burnextension.h
index 370ddd2d..3529ef38 100644
--- a/src/burn/engine/burnextension.h
+++ b/src/burn/engine/burnextension.h
@@ -13,7 +13,6 @@ typedef struct _BURN_EXTENSION_ENGINE_CONTEXT BURN_EXTENSION_ENGINE_CONTEXT;
 
 typedef struct _BURN_EXTENSION
 {
-    LPWSTR sczEntryPayloadId;
     LPWSTR sczId;
 
     BURN_PAYLOAD* pEntryPayload;
diff --git a/src/burn/engine/container.cpp b/src/burn/engine/container.cpp
index c6f2ada8..aa123927 100644
--- a/src/burn/engine/container.cpp
+++ b/src/burn/engine/container.cpp
@@ -35,7 +35,7 @@ extern "C" HRESULT ContainersParseFromXml(
 
     pContainers->cContainers = cNodes;
 
-    // parse search elements
+    // parse container elements
     for (DWORD i = 0; i < cNodes; ++i)
     {
         BURN_CONTAINER* pContainer = &pContainers->rgContainers[i];
@@ -181,6 +181,7 @@ extern "C" void ContainersUninitialize(
             ReleaseStr(pContainer->downloadSource.sczUser);
             ReleaseStr(pContainer->downloadSource.sczPassword);
             ReleaseStr(pContainer->sczUnverifiedPath);
+            ReleaseDict(pContainer->sdhPayloads);
         }
         MemFree(pContainers->rgContainers);
     }
diff --git a/src/burn/engine/container.h b/src/burn/engine/container.h
index d454a248..f35f1da5 100644
--- a/src/burn/engine/container.h
+++ b/src/burn/engine/container.h
@@ -72,6 +72,9 @@ typedef struct _BURN_CONTAINER
     LPWSTR sczFilePath;         // relative path to container.
     DOWNLOAD_SOURCE downloadSource;
 
+    DWORD cParsedPayloads;
+    STRINGDICT_HANDLE sdhPayloads; // value is BURN_PAYLOAD*
+
     BYTE* pbHash;
     DWORD cbHash;
     BURN_CONTAINER_VERIFICATION verification;
diff --git a/src/burn/engine/manifest.h b/src/burn/engine/manifest.h
index 8c527279..40748c52 100644
--- a/src/burn/engine/manifest.h
+++ b/src/burn/engine/manifest.h
@@ -2,8 +2,6 @@
 // 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.
 
 
-interface IBurnPayload; // forward declare.
-
 #if defined(__cplusplus)
 extern "C" {
 #endif
diff --git a/src/burn/engine/payload.cpp b/src/burn/engine/payload.cpp
index 84c32eec..e0362a76 100644
--- a/src/burn/engine/payload.cpp
+++ b/src/burn/engine/payload.cpp
@@ -5,13 +5,6 @@
 
 // internal function declarations
 
-static HRESULT FindEmbeddedBySourcePath(
-    __in BURN_PAYLOADS* pPayloads,
-    __in_opt BURN_CONTAINER* pContainer,
-    __in_z LPCWSTR wzStreamName,
-    __out BURN_PAYLOAD** ppPayload
-    );
-
 
 // function definitions
 
@@ -29,6 +22,7 @@ extern "C" HRESULT PayloadsParseFromXml(
     LPWSTR scz = NULL;
     BOOL fChainPayload = pContainers && pLayoutPayloads; // These are required when parsing chain payloads.
     BOOL fValidFileSize = FALSE;
+    size_t cByteOffset = fChainPayload ? offsetof(BURN_PAYLOAD, sczKey) : offsetof(BURN_PAYLOAD, sczSourcePath);
 
     // select payload nodes
     hr = XmlSelectNodes(pixnBundle, L"Payload", &pixnNodes);
@@ -49,7 +43,11 @@ extern "C" HRESULT PayloadsParseFromXml(
 
     pPayloads->cPayloads = cNodes;
 
-    // parse search elements
+    // create dictionary for payloads
+    hr = DictCreateWithEmbeddedKey(&pPayloads->sdhPayloads, pPayloads->cPayloads, reinterpret_cast<void**>(&pPayloads->rgPayloads), cByteOffset, DICT_FLAG_NONE);
+    ExitOnFailure(hr, "Failed to create dictionary for payloads.");
+
+    // parse payload elements
     for (DWORD i = 0; i < cNodes; ++i)
     {
         BURN_PAYLOAD* pPayload = &pPayloads->rgPayloads[i];
@@ -104,6 +102,8 @@ extern "C" HRESULT PayloadsParseFromXml(
                 // find container
                 hr = ContainerFindById(pContainers, scz, &pPayload->pContainer);
                 ExitOnFailure(hr, "Failed to to find container: %ls", scz);
+
+                pPayload->pContainer->cParsedPayloads += 1;
             }
 
             // @LayoutOnly
@@ -190,12 +190,37 @@ extern "C" HRESULT PayloadsParseFromXml(
             }
         }
 
+        hr = DictAddValue(pPayloads->sdhPayloads, pPayload);
+        ExitOnFailure(hr, "Failed to add payload to payloads dictionary.");
+
         // prepare next iteration
         ReleaseNullObject(pixnNode);
     }
 
     hr = S_OK;
 
+    if (pContainers && pContainers->cContainers)
+    {
+        for (DWORD i = 0; i < pPayloads->cPayloads; ++i)
+        {
+            BURN_PAYLOAD* pPayload = &pPayloads->rgPayloads[i];
+            BURN_CONTAINER* pContainer = pPayload->pContainer;
+
+            if (!pContainer)
+            {
+                continue;
+            }
+            else if (!pContainer->sdhPayloads)
+            {
+                hr = DictCreateWithEmbeddedKey(&pContainer->sdhPayloads, pContainer->cParsedPayloads, NULL, offsetof(BURN_PAYLOAD, sczSourcePath), DICT_FLAG_NONE);
+                ExitOnFailure(hr, "Failed to create dictionary for container payloads.");
+            }
+
+            hr = DictAddValue(pContainer->sdhPayloads, pPayload);
+            ExitOnFailure(hr, "Failed to add payload to container dictionary.");
+        }
+    }
+
 LExit:
     ReleaseObject(pixnNodes);
     ReleaseObject(pixnNode);
@@ -237,6 +262,8 @@ extern "C" void PayloadsUninitialize(
         MemFree(pPayloads->rgPayloads);
     }
 
+    ReleaseDict(pPayloads->sdhPayloads);
+
     // clear struct
     memset(pPayloads, 0, sizeof(BURN_PAYLOADS));
 }
@@ -265,7 +292,7 @@ extern "C" HRESULT PayloadExtractUXContainer(
         ExitOnFailure(hr, "Failed to get next stream.");
 
         // find payload by stream name
-        hr = PayloadFindEmbeddedBySourcePath(pPayloads, sczStreamName, &pPayload);
+        hr = PayloadFindEmbeddedBySourcePath(pPayloads->sdhPayloads, sczStreamName, &pPayload);
         ExitOnFailure(hr, "Failed to find embedded payload: %ls", sczStreamName);
 
         // make file path
@@ -313,51 +340,22 @@ extern "C" HRESULT PayloadFindById(
     )
 {
     HRESULT hr = S_OK;
-    BURN_PAYLOAD* pPayload = NULL;
 
-    for (DWORD i = 0; i < pPayloads->cPayloads; ++i)
-    {
-        pPayload = &pPayloads->rgPayloads[i];
-
-        if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, pPayload->sczKey, -1, wzId, -1))
-        {
-            *ppPayload = pPayload;
-            ExitFunction1(hr = S_OK);
-        }
-    }
+    hr = DictGetValue(pPayloads->sdhPayloads, wzId, reinterpret_cast<void**>(ppPayload));
 
-    hr = E_NOTFOUND;
-
-LExit:
     return hr;
 }
 
 extern "C" HRESULT PayloadFindEmbeddedBySourcePath(
-    __in BURN_PAYLOADS* pPayloads,
+    __in STRINGDICT_HANDLE sdhPayloads,
     __in_z LPCWSTR wzStreamName,
     __out BURN_PAYLOAD** ppPayload
     )
 {
     HRESULT hr = S_OK;
-    BURN_PAYLOAD* pPayload = NULL;
 
-    for (DWORD i = 0; i < pPayloads->cPayloads; ++i)
-    {
-        pPayload = &pPayloads->rgPayloads[i];
-
-        if (BURN_PAYLOAD_PACKAGING_EMBEDDED == pPayload->packaging)
-        {
-            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, pPayload->sczSourcePath, -1, wzStreamName, -1))
-            {
-                *ppPayload = pPayload;
-                ExitFunction1(hr = S_OK);
-            }
-        }
-    }
+    hr = DictGetValue(sdhPayloads, wzStreamName, reinterpret_cast<void**>(ppPayload));
 
-    hr = E_NOTFOUND;
-
-LExit:
     return hr;
 }
 
diff --git a/src/burn/engine/payload.h b/src/burn/engine/payload.h
index c12fbe66..14738506 100644
--- a/src/burn/engine/payload.h
+++ b/src/burn/engine/payload.h
@@ -66,6 +66,7 @@ typedef struct _BURN_PAYLOADS
 {
     BURN_PAYLOAD* rgPayloads;
     DWORD cPayloads;
+    STRINGDICT_HANDLE sdhPayloads; // value is BURN_PAYLOAD*
 } BURN_PAYLOADS;
 
 typedef struct _BURN_PAYLOAD_GROUP_ITEM
@@ -109,7 +110,7 @@ HRESULT PayloadFindById(
     __out BURN_PAYLOAD** ppPayload
     );
 HRESULT PayloadFindEmbeddedBySourcePath(
-    __in BURN_PAYLOADS* pPayloads,
+    __in STRINGDICT_HANDLE sdhPayloads,
     __in_z LPCWSTR wzStreamName,
     __out BURN_PAYLOAD** ppPayload
     );
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index 6a5ae697..584bef14 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -15,7 +15,7 @@ const DWORD MB_RETRYTRYAGAIN = 0xF;
 
 // structs
 
-typedef struct _BOOTSTRAPPER_ENGINE_CONTEXT BOOTSTRAPPER_ENGINE_CONTEXT;
+typedef struct _BOOTSTRAPPER_ENGINE_CONTEXT BOOTSTRAPPER_ENGINE_CONTEXT; // forward declare
 
 typedef struct _BURN_USER_EXPERIENCE
 {