Mods: retain bundle-owned custom programs alongside verified stock aliases

Engine/Resource/Test/shader_lookup_tests.cpp

@@ -48,9 +48,43 @@ Lookup lookup(std::uint32_t records, const std::map<std::uint64_t, std::vector<s
     append(out.meta, std::uint32_t{0});
     return out;
 }
+// A complete preserved stock heap with an appended program array. This is
+// intentionally different from merely sharing record counts or numeric indices.
+Lookup preserved_lookup(bool extended, const std::map<std::uint64_t, std::vector<std::uint32_t>>& rows) {
+    Lookup out;
+    const std::uint32_t stockSize = 208;
+    const std::uint64_t start = extended ? 208 : 16;
+    const std::uint32_t count = extended ? 2 : 1;
+    const std::uint32_t size = extended ? 544 : stockSize;
+    out.resource.resize(size);
+    const std::uint64_t heap = 192;
+    std::memcpy(out.resource.data(), &count, 4);
+    std::memcpy(out.resource.data() + 4, &start, 8);
+    // The first array's count/pointer, within the stock record at 16.
+    const std::uint32_t one = 1;
+    std::memcpy(out.resource.data() + 28, &one, 4);
+    std::memcpy(out.resource.data() + 32, &heap, 8);
+    out.resource[192] = std::byte{7};
+    if (extended) std::memcpy(out.resource.data() + 208, out.resource.data() + 16, 168);
+    append(out.resource, std::uint32_t{4});
+    append(out.resource, std::uint32_t{32});
+    if (extended) append(out.resource, std::uint32_t{224});
+    append(out.resource, static_cast<std::uint32_t>(rows.size()));
+    for (const auto& [key, values] : rows) {
+        append(out.resource, key);
+        append(out.resource, static_cast<std::uint32_t>(values.size()));
+        for (const auto v : values) append(out.resource, v);
+    }
+    const std::uint32_t fix = extended ? 12 : 8;
+    append(out.meta, size); append(out.meta, fix);
+    append(out.meta, static_cast<std::uint32_t>(out.resource.size() - size - fix));
+    append(out.meta, std::uint32_t{});
+    return out;
+}
+
 // The merged rows, read back.
 std::map<std::uint64_t, std::vector<std::uint32_t>> rows_of(const fb::shader::MergedTable& merged) {
-    std::size_t at = read<std::uint32_t>(merged.resourceMeta, 0);
+    std::size_t at = static_cast<std::size_t>(read<std::uint32_t>(merged.resourceMeta, 0)) + read<std::uint32_t>(merged.resourceMeta, 4);
     const auto count = read<std::uint32_t>(merged.resource, at);
     at += 4;
     std::map<std::uint64_t, std::vector<std::uint32_t>> rows;
@@ -97,6 +131,44 @@ int main() {
     check(std::equal(selected.resource.begin(), selected.resource.begin() + 8,
                      retained.resource.begin()), "Opaque program bytes remain unchanged");
     check(retained.resource.size() == alias.resource.size(), "Aliases never duplicate program heaps");
+    const auto stock = preserved_lookup(false, {{1, {0}}});
+    const auto stockAlias = preserved_lookup(false, {{1, {0}}, {2, {0}}});
+    const auto extension = preserved_lookup(true, {{1, {0}}, {9, {1}}});
+    const std::vector<fb::shader::Table> stockEdits{stockAlias.table()};
+    check(fb::shader::only_stock_aliases(stockAlias.table(), stock.table()) &&
+          !fb::shader::only_stock_aliases(extension.table(), stock.table()),
+          "Selection distinguishes stock aliases from a private graph");
+    const auto replacement = preserved_lookup(false, {{1, {0, 0}}, {2, {0}}});
+    check(!fb::shader::only_stock_aliases(replacement.table(), stock.table()),
+          "Changing an original material row is not an alias-only edit");
+    const auto compatible = fb::shader::merge_program_lookup(extension.table(), stockEdits, stock.table());
+    check(compatible.added == 1 && compatible.outside == 0 && rows_of(compatible).contains(9),
+          "Verified stock aliases coexist with an appended private program array");
+    check(std::equal(extension.resource.begin(), extension.resource.begin() + 556, compatible.resource.begin()),
+          "Stock alias merge leaves the selected graph and all relocations untouched");
+    auto changedHeap = extension;
+    changedHeap.resource[192] = std::byte{8};
+    check(fb::shader::merge_program_lookup(changedHeap.table(), stockEdits, stock.table()).outside == 1,
+          "A changed stock heap cannot borrow stock aliases");
+    auto missingFixup = extension;
+    const std::uint32_t other = 0;
+    std::memcpy(missingFixup.resource.data() + 552, &other, 4);
+    check(fb::shader::merge_program_lookup(missingFixup.table(), stockEdits, stock.table()).outside == 1,
+          "A copied record without its relocation cannot borrow stock aliases");
+    auto extraFixup = extension;
+    const std::uint32_t unexpected = 64, extraSize = 16;
+    const auto insertion = extraFixup.resource.begin() + 556;
+    const auto* encoded = reinterpret_cast<const std::byte*>(&unexpected);
+    extraFixup.resource.insert(insertion, encoded, encoded + 4);
+    std::memcpy(extraFixup.meta.data() + 4, &extraSize, 4);
+    check(fb::shader::merge_program_lookup(extraFixup.table(), stockEdits, stock.table()).outside == 1,
+          "Extra fixups cannot reinterpret the unchanged stock heap");
+    const auto invalidAlias = preserved_lookup(false, {{1, {0}}, {2, {1}}});
+    const std::vector<fb::shader::Table> invalidEdits{invalidAlias.table()};
+    check(fb::shader::merge_program_lookup(extension.table(), invalidEdits, stock.table()).outside == 1,
+          "Stock aliases cannot address the selected graph's private program slots");
+    check(fb::shader::merge_program_lookup(extension.table(), stockEdits).outside == 1,
+          "The stock exception requires the caller's actual game table");
     if (failures) return 1;
     std::cout << "shader lookup tests passed\n";
     return 0;

Engine/Resource/shader_lookup.cpp

@@ -42,6 +42,90 @@ std::size_t lookup_offset(const Table& table) {
     return static_cast<std::size_t>(offset);
 }
 
+// Equality includes the opaque heap, relocation list and graph metadata;
+// lookup rows themselves can differ.
+bool same_graph(const Table& left, const Table& right) {
+    if (left.resourceMeta.size() != 16 || right.resourceMeta.size() != 16) return false;
+    const auto end = lookup_offset(left);
+    return end == lookup_offset(right) &&
+        std::equal(left.resourceMeta.begin(), left.resourceMeta.begin() + 8, right.resourceMeta.begin()) &&
+        std::equal(left.resourceMeta.begin() + 12, left.resourceMeta.end(), right.resourceMeta.begin() + 12) &&
+        std::equal(left.resource.begin(), left.resource.begin() + static_cast<std::ptrdiff_t>(end), right.resource.begin());
+}
+
+// Some authoring tools retain the complete stock heap and append a new program
+// array whose initial records are exact copies of the stock records. Stock-only
+// aliases can still address those programs, without importing any bytes. This
+// optional path recognises the pinned game's 168-byte records; an unfamiliar
+// layout simply cannot use it, rather than failing the ordinary alias merge.
+bool preserves_stock(const Table& retained, const Table& stock) {
+    if (stock.resourceMeta.size() != 16) return false;
+    try {
+        (void)lookup_offset(stock);
+        const auto stockSize = read<std::uint32_t>(stock.resourceMeta, 0);
+        const auto keptSize = read<std::uint32_t>(retained.resourceMeta, 0);
+        const auto count = read<std::uint32_t>(stock.resource, 0);
+        const auto keptCount = read<std::uint32_t>(retained.resource, 0);
+        const auto start = read<std::uint64_t>(stock.resource, 4);
+        const auto keptStart = read<std::uint64_t>(retained.resource, 4);
+        const auto recordBytes = static_cast<std::uint64_t>(count) * 168;
+        if (!count || count > keptCount || stockSize < 12 || keptSize < stockSize ||
+            start < 12 || start > stockSize || recordBytes > stockSize - start ||
+            keptStart < 12 || keptStart > keptSize || recordBytes > keptSize - keptStart ||
+            !std::equal(stock.resourceMeta.begin() + 12, stock.resourceMeta.end(), retained.resourceMeta.begin() + 12) ||
+            !std::equal(stock.resource.begin() + 12, stock.resource.begin() + stockSize, retained.resource.begin() + 12) ||
+            !std::equal(stock.resource.begin() + static_cast<std::ptrdiff_t>(start),
+                        stock.resource.begin() + static_cast<std::ptrdiff_t>(start + recordBytes),
+                        retained.resource.begin() + static_cast<std::ptrdiff_t>(keptStart))) return false;
+        const auto relocations = [](const Table& table, std::uint32_t size) {
+            const auto bytes = read<std::uint32_t>(table.resourceMeta, 4);
+            if (bytes % 4) throw std::runtime_error("relocation layout");
+            std::vector<std::uint32_t> fields;
+            for (std::size_t at = 0; at < bytes; at += 4) {
+                const auto field = read<std::uint32_t>(table.resource, static_cast<std::size_t>(size) + at);
+                if (field > size || size - field < 8 || read<std::uint64_t>(table.resource, field) > size)
+                    throw std::runtime_error("relocation bounds");
+                fields.push_back(field);
+            }
+            std::ranges::sort(fields);
+            if (std::adjacent_find(fields.begin(), fields.end()) != fields.end())
+                throw std::runtime_error("duplicate relocation");
+            return fields;
+        };
+        const auto fields = relocations(stock, stockSize);
+        const auto keptFields = relocations(retained, keptSize);
+        if (!std::binary_search(fields.begin(), fields.end(), 4U) ||
+            !std::binary_search(keptFields.begin(), keptFields.end(), 4U)) return false;
+        for (const auto field : fields) {
+            if (!std::binary_search(keptFields.begin(), keptFields.end(), field)) return false;
+            if (field >= start && field < start + recordBytes) {
+                const auto copied = static_cast<std::uint32_t>(keptStart + field - start);
+                if (!std::binary_search(keptFields.begin(), keptFields.end(), copied)) return false;
+            }
+        }
+        // Additional fixups in either preserved region would reinterpret its
+        // unchanged bytes as pointers, so membership must agree both ways.
+        for (const auto field : keptFields) {
+            if (field < stockSize && !std::binary_search(fields.begin(), fields.end(), field)) return false;
+            if (field >= keptStart && field < keptStart + recordBytes) {
+                const auto original = static_cast<std::uint32_t>(start + field - keptStart);
+                if (!std::binary_search(fields.begin(), fields.end(), original)) return false;
+            }
+        }
+        struct Array { std::size_t count, pointer, stride; };
+        constexpr Array arrays[]{{0x0c,0x10,16},{0x18,0x1c,8},{0x30,0x34,16},
+                                {0x3c,0x40,6},{0x48,0x4c,8},{0x54,0x58,2},{0x60,0x64,8}};
+        for (std::uint32_t i = 0; i < count; ++i) for (const auto& array : arrays) {
+            const auto record = static_cast<std::size_t>(start + i * 168ULL);
+            const auto items = read<std::uint32_t>(stock.resource, record + array.count);
+            const auto pointer = read<std::uint64_t>(stock.resource, record + array.pointer);
+            if (pointer > stockSize || static_cast<std::uint64_t>(items) * array.stride > stockSize - pointer ||
+                (items && !std::binary_search(fields.begin(), fields.end(), static_cast<std::uint32_t>(record + array.pointer)))) return false;
+        }
+        return true;
+    } catch (const std::exception&) { return false; }
+}
+
 ProgramRows read_program_rows(const Table& table) {
     auto offset = lookup_offset(table);
     const auto count = read<std::uint32_t>(table.resource, offset);
@@ -78,28 +162,47 @@ TextureRows read_texture_rows(const Table& table) {
 
 } // namespace
 
-MergedTable merge_program_lookup(const Table& base, const std::span<const Table> edits) {
+bool preserves_stock_programs(const Table& candidate, const Table& stock) {
+    return preserves_stock(candidate, stock);
+}
+
+bool only_stock_aliases(const Table& candidate, const Table& stock) {
+    try {
+        if (!same_graph(candidate, stock)) return false;
+        const auto rows = read_program_rows(candidate);
+        const auto original = read_program_rows(stock);
+        const auto count = read<std::uint32_t>(stock.resource, 0);
+        for (const auto& [key, values] : original) {
+            const auto at = rows.find(key);
+            if (at == rows.end() || at->second != values) return false;
+        }
+        for (const auto& [key, values] : rows) {
+            (void)key;
+            if (std::any_of(values.begin(), values.end(), [&](auto index) { return index >= count; })) return false;
+        }
+        return true;
+    } catch (const std::exception&) { return false; }
+}
+
+MergedTable merge_program_lookup(const Table& base, const std::span<const Table> edits, const Table& stock) {
     const auto start = lookup_offset(base);
     const auto records = read<std::uint32_t>(base.resource, 0); // the base blob's record count
     auto rows = read_program_rows(base);
     const auto baseRows = rows;
 
+    const bool stockPreserved = preserves_stock(base, stock);
+    const auto stockRecords = stockPreserved ? read<std::uint32_t>(stock.resource, 0) : 0;
     MergedTable merged;
     for (const auto& edit : edits) {
         // Indices have meaning only within their owning graph. Keep the selected
         // bundle's heap and relocations verbatim: importing another graph can
         // make the game instantiate programs before their compiled files load.
-        const auto editStart = lookup_offset(edit);
-        const bool sameGraph = editStart == start &&
-            std::equal(base.resourceMeta.begin(), base.resourceMeta.begin() + 8,
-                       edit.resourceMeta.begin()) &&
-            std::equal(base.resourceMeta.begin() + 12, base.resourceMeta.end(),
-                       edit.resourceMeta.begin() + 12) &&
-            std::equal(base.resource.begin(), base.resource.begin() +
-                       static_cast<std::ptrdiff_t>(start), edit.resource.begin());
+        const bool sameGraph = same_graph(base, edit);
+        const bool stockGraph = stockPreserved && same_graph(stock, edit);
+        const auto available = sameGraph ? records : stockRecords;
         for (const auto& [key, values] : read_program_rows(edit)) {
             if (baseRows.contains(key)) continue;
-            if (!sameGraph || std::any_of(values.begin(), values.end(), [&](std::uint32_t index) { return index >= records; })) {
+            if ((!sameGraph && !stockGraph) || std::any_of(values.begin(), values.end(), [&](std::uint32_t index) { return index >= available; })) {
                 ++merged.outside;
                 continue;
             }

Engine/Resource/shader_lookup.h

@@ -14,6 +14,7 @@ namespace dingosdk::frostbite::shader {
 
 inline constexpr std::uint32_t programLookupType = 0xE2E6955B;
 inline constexpr std::uint32_t textureLookupType = 0x2D254A89;
+inline constexpr std::uint32_t bytecodeType = 0x0E2B02F7;
 
 struct Table {
     std::span<const std::byte> resource;
@@ -30,9 +31,14 @@ struct MergedTable {
 
 // The selected bundle contributor owns the retained program graph. Merge only
 // aliases whose program heap and relocation bytes match it exactly. No program
-// records or relocations are interpreted, copied or rebased. `outside` includes
+// records or relocations are copied or rebased. An optional stock table allows
+// stock aliases when its heap, records and relocations survive unchanged. `outside` includes
 // rows from incompatible graphs as well as out-of-range indices.
-[[nodiscard]] MergedTable merge_program_lookup(const Table& base, std::span<const Table> edits);
+// Selection helpers for contributors to the SAME bundle. A caller selecting a
+// different contributor must also retain that contributor's compiled resources.
+[[nodiscard]] bool preserves_stock_programs(const Table& candidate, const Table& stock);
+[[nodiscard]] bool only_stock_aliases(const Table& candidate, const Table& stock);
+[[nodiscard]] MergedTable merge_program_lookup(const Table& base, std::span<const Table> edits, const Table& stock = {});
 
 [[nodiscard]] MergedTable merge_texture_lookup(const Table& base, std::span<const Table> edits);
 

Engine/Vfs/mod_merge_combine.cpp

@@ -57,6 +57,8 @@ struct Contribution {
     fb::Sha1 sha1;
     std::vector<std::byte> resourceMeta;
     bool base{};
+    std::uint64_t resourceId{};
+    std::uint32_t resourceType{};
 };
 
 // Region files run parallel to the manifest's ebx, then resources, then chunks.
@@ -986,7 +988,7 @@ fb::TocDocument combine(const fs::path& baseToc, const fs::path& baseRoot,
             if (const auto original = repairedFrom.find(firstFile + position); original != repairedFrom.end()) {
                 // Repaired above: what this mod changed is its copy as it ships it.
                 state.history[id].push_back({&root, original->second.first, original->second.second,
-                                             entry.asset.resourceMeta, isBase});
+                                             entry.asset.resourceMeta, isBase, entry.asset.resourceId, entry.asset.resourceType});
             } else if (combinable && !isBase && state.matches_game_copy(id, entry.asset.sha1)) {
                 ++report.unchangedCopies;
             } else if (combinable) {
@@ -997,7 +999,7 @@ fb::TocDocument combine(const fs::path& baseToc, const fs::path& baseRoot,
                     if (const auto* directory = store.find_directory(unshifted.location.installChunk))
                         unshift(*placement, *directory, unshifted.location.archive, unshifted.offset);
                 state.history[id].push_back(
-                    {&root, unshifted, entry.asset.sha1, entry.asset.resourceMeta, isBase});
+                    {&root, unshifted, entry.asset.sha1, entry.asset.resourceMeta, isBase, entry.asset.resourceId, entry.asset.resourceType});
             }
             if (isBase) state.baseSha.insert_or_assign(id, entry.asset.sha1);
             const auto at = state.seen.find(id);
@@ -1148,19 +1150,56 @@ fb::TocDocument combine(const fs::path& baseToc, const fs::path& baseRoot,
                         const Contribution* retained = baseCopy;
                         for (const auto& contribution : found->second)
                             if (contribution.sha1 == entry.asset.sha1) retained = &contribution;
-                        const auto retainedBytes = decode(*retained);
+                        const auto* selected = retained;
+                        auto retainedBytes = decode(*retained);
+                        const fb::shader::Table stockTable{baseBytes, baseCopy->resourceMeta};
+                        if (entry.asset.resourceType == fb::shader::programLookupType &&
+                            fb::shader::only_stock_aliases({retainedBytes, retained->resourceMeta}, stockTable)) {
+                            // A stock-alias edit owns no new programs. It must not
+                            // erase a graph another contributor already delivers
+                            // to this same bundle. Highest custom contributor wins.
+                            for (std::size_t i = 0; i < edits.size(); ++i) {
+                                const fb::shader::Table candidate{editBytes[i], edits[i]->resourceMeta};
+                                if (fb::shader::only_stock_aliases(candidate, stockTable) ||
+                                    !fb::shader::preserves_stock_programs(candidate, stockTable)) continue;
+                                bool dependenciesRetained = true;
+                                bool bytecodeDelivered = false;
+                                for (const auto& resource : state.assets) {
+                                    if (resource.asset.kind != fb::AssetKind::resource ||
+                                        resource.asset.resourceType == fb::shader::programLookupType ||
+                                        resource.asset.resourceType == fb::shader::textureLookupType ||
+                                        fb::bundle_ref::is_table(resource.asset.name)) continue;
+                                    const auto key = asset_key(resource.asset);
+                                    const auto history = state.history.find(key);
+                                    if (history == state.history.end()) continue;
+                                    const auto shipped = state.baseSha.find(key);
+                                    for (const auto& copy : history->second) {
+                                        if (copy.root != edits[i]->root || copy.base ||
+                                            (shipped != state.baseSha.end() && copy.sha1 == shipped->second)) continue;
+                                        if (copy.resourceType == fb::shader::bytecodeType) bytecodeDelivered = true;
+                                        if (copy.sha1 != resource.asset.sha1 ||
+                                            copy.resourceMeta != resource.asset.resourceMeta ||
+                                            copy.resourceId != resource.asset.resourceId ||
+                                            copy.resourceType != resource.asset.resourceType) dependenciesRetained = false;
+                                    }
+                                }
+                                if (!dependenciesRetained || !bytecodeDelivered) continue;
+                                retained = edits[i];
+                                retainedBytes = editBytes[i];
+                            }
+                        }
                         const fb::shader::Table baseTable =
                             entry.asset.resourceType == fb::shader::programLookupType
                                 ? fb::shader::Table{retainedBytes, retained->resourceMeta}
                                 : fb::shader::Table{baseBytes, baseCopy->resourceMeta};
                         const auto table = entry.asset.resourceType == fb::shader::programLookupType
-                            ? fb::shader::merge_program_lookup(baseTable, tables)
+                            ? fb::shader::merge_program_lookup(baseTable, tables, {baseBytes, baseCopy->resourceMeta})
                             : fb::shader::merge_texture_lookup(baseTable, tables);
                         if (table.outside)
                             report.notes.push_back(entry.asset.name + ": " + std::to_string(table.outside) +
                                 " material row(s) left out: their graph or indices do not match the retained "
                                 "bundle shader database");
-                        if (!table.added) continue;
+                        if (!table.added && retained == selected) continue;
                         rebuilt = table.resource;
                         entry.asset.resourceMeta = table.resourceMeta;
                         what = std::to_string(table.added) + " material row(s)";