Thunks/gen: Implement assisted struct repacking

This can be used to allow automatically handling structures that require
special behavior for one member but are automatically repackable otherwise.
The feature is enabled using the new custom_repack annotation and requires
additional repacking functions to be defined in the host file for each
customized member.

ThunkLibs/Generator/analysis.cpp

@@ -300,6 +300,24 @@ void AnalysisAction::ParseInterface(clang::ASTContext& context) {
 
                 if (auto emitted_function = llvm::dyn_cast<clang::FunctionDecl>(template_args[0].getAsDecl())) {
                     // Process later
+                } else if (auto annotated_member = llvm::dyn_cast<clang::FieldDecl>(template_args[0].getAsDecl())) {
+                    {
+                        if (decl->getNumBases() != 1 || decl->bases_begin()->getType().getAsString() != "fexgen::custom_repack") {
+                            throw report_error(template_arg_loc, "Unsupported member annotation(s)");
+                        }
+
+                        if (!annotated_member->getType()->isPointerType() && !annotated_member->getType()->isArrayType()) {
+                            throw report_error(template_arg_loc, "custom_repack annotation requires pointer member");
+                        }
+                    }
+
+                    // Get or add parent type to list of structure types
+                    auto repack_info_it = types.emplace(context.getCanonicalType(annotated_member->getParent()->getTypeForDecl()), RepackedType {}).first;
+                    if (repack_info_it->second.assumed_compatible) {
+                        throw report_error(template_arg_loc, "May not annotate members of opaque types");
+                    }
+                    // Add member to its list of members
+                    repack_info_it->second.custom_repacked_members.insert(annotated_member->getNameAsString());
                 } else {
                     throw report_error(template_arg_loc, "Cannot annotate this kind of symbol");
                 }
@@ -356,14 +374,17 @@ void AnalysisAction::ParseInterface(clang::ASTContext& context) {
                             }
 
                             for (auto* member : type->getAsStructureType()->getDecl()->fields()) {
-                                  /*if (!member->getType()->isPointerType())*/ {
-                                      // TODO: Perform more elaborate validation for non-pointers to ensure ABI compatibility
-                                      continue;
-                                  }
-
-                                  throw report_error(member->getBeginLoc(), "Unannotated pointer member")
-                                        .addNote(report_error(param_loc, "in struct type", clang::DiagnosticsEngine::Note))
-                                        .addNote(report_error(template_arg_loc, "used in annotation here", clang::DiagnosticsEngine::Note));
+                                auto annotated_type = types.find(type->getCanonicalTypeUnqualified().getTypePtr());
+                                if (annotated_type == types.end() || !annotated_type->second.UsesCustomRepackFor(member)) {
+                                    /*if (!member->getType()->isPointerType())*/ {
+                                        // TODO: Perform more elaborate validation for non-pointers to ensure ABI compatibility
+                                        continue;
+                                    }
+
+                                    throw report_error(member->getBeginLoc(), "Unannotated pointer member")
+                                          .addNote(report_error(param_loc, "in struct type", clang::DiagnosticsEngine::Note))
+                                          .addNote(report_error(template_arg_loc, "used in annotation here", clang::DiagnosticsEngine::Note));
+                                }
                             }
                         };
 
@@ -518,7 +539,7 @@ void AnalysisAction::CoverReferencedTypes(clang::ASTContext& context) {
                     }
                     continue;
                 }
-                if (member_type->isUnionType() && !types.contains(member_type)) {
+                if (member_type->isUnionType() && !types.contains(member_type) && !type_repack_info.UsesCustomRepackFor(member)) {
                     throw std::runtime_error(fmt::format("\"{}\" has unannotated member \"{}\" of union type \"{}\"",
                                                          clang::QualType { type, 0 }.getAsString(),
                                                          member->getNameAsString(),

ThunkLibs/Generator/analysis.h

@@ -122,6 +122,16 @@ class AnalysisAction : public clang::ASTFrontendAction {
 
         // If true, emit guest_layout/host_layout definitions even if the type is non-repackable
         bool emit_layout_wrappers = false;
+
+        // Set of members (identified by their field name) with custom repacking
+        std::unordered_set<std::string> custom_repacked_members;
+
+        bool UsesCustomRepackFor(const clang::FieldDecl* member) const {
+            return custom_repacked_members.contains(member->getNameAsString());
+        }
+        bool UsesCustomRepackFor(const std::string& member_name) const {
+            return custom_repacked_members.contains(member_name);
+        }
     };
 
 protected:

ThunkLibs/Generator/data_layout.cpp

@@ -272,7 +272,9 @@ TypeCompatibility DataLayoutCompareAction::GetTypeCompatibility(
                 // * Pointer member is annotated
                 // TODO: Don't restrict this to structure types. it applies to pointers to builtin types too!
                 auto host_member_pointee_type = context.getCanonicalType(host_member_type->getPointeeType().getTypePtr());
-                if (types.contains(host_member_pointee_type) && types.at(host_member_pointee_type).assumed_compatible) {
+                if (types.at(type).UsesCustomRepackFor(host_member_field)) {
+                    member_compat.push_back(TypeCompatibility::Repackable);
+                } else if (types.contains(host_member_pointee_type) && types.at(host_member_pointee_type).assumed_compatible) {
                     // Pointee doesn't need repacking, but pointer needs extending on 32-bit
                     member_compat.push_back(is_32bit ? TypeCompatibility::Repackable : TypeCompatibility::Full);
                 } else if (host_member_pointee_type->isPointerType()) {

ThunkLibs/Generator/gen.cpp

@@ -234,11 +234,19 @@ void GenerateThunkLibsAction::EmitLayoutWrappers(
             };
             // Prefer initialization via the constructor's initializer list if possible (to detect unintended narrowing), otherwise initialize in the body
             for (auto* member : type->getAsStructureType()->getDecl()->fields()) {
-                map_field(member, true);
+                if (!type_repack_info.UsesCustomRepackFor(member)) {
+                    map_field(member, true);
+                } else {
+                  // Leave field uninitialized
+                }
             }
             fmt::print(file, "    }} {{\n");
             for (auto* member : type->getAsStructureType()->getDecl()->fields()) {
-                map_field(member, false);
+                if (!type_repack_info.UsesCustomRepackFor(member)) {
+                    map_field(member, false);
+                } else {
+                  // Leave field uninitialized
+                }
             }
         }
         fmt::print(file, "  }}\n");
@@ -274,16 +282,46 @@ void GenerateThunkLibsAction::EmitLayoutWrappers(
 
             // Prefer initialization via the constructor's initializer list if possible (to detect unintended narrowing), otherwise initialize in the body
             for (auto& member : guest_abi.at(struct_name).get_if_struct()->members) {
-                map_field2(member, true);
+                if (!type_repack_info.UsesCustomRepackFor(member.member_name)) {
+                    map_field2(member, true);
+                } else {
+                  // Leave field uninitialized
+                }
             }
             fmt::print(file, "  }} }};\n");
             for (auto& member : guest_abi.at(struct_name).get_if_struct()->members) {
-                map_field2(member, false);
+                if (!type_repack_info.UsesCustomRepackFor(member.member_name)) {
+                    map_field2(member, false);
+                } else {
+                  // Leave field uninitialized
+                }
             }
         }
         fmt::print(file, "  return ret;\n");
         fmt::print(file, "}}\n\n");
 
+        // Forward-declare user-provided repacking functions
+        if (type_repack_info.custom_repacked_members.empty()) {
+            fmt::print(file, "void fex_apply_custom_repacking_entry(host_layout<{}>& source, const guest_layout<{}>& from) {{\n", struct_name, struct_name);
+            fmt::print(file, "}}\n");
+            fmt::print(file, "bool fex_apply_custom_repacking_exit(guest_layout<{}>& into, host_layout<{}>& from) {{\n", struct_name, struct_name);
+            fmt::print(file, "  return false;\n");
+            fmt::print(file, "}}\n");
+        } else {
+            fmt::print(file, "void fex_custom_repack_entry(host_layout<{}>& into, const guest_layout<{}>& from);\n",
+                       struct_name, struct_name);
+            fmt::print(file, "bool fex_custom_repack_exit(guest_layout<{}>& into, const host_layout<{}>& from);\n\n",
+                       struct_name, struct_name);
+
+            fmt::print(file, "void fex_apply_custom_repacking_entry(host_layout<{}>& source, const guest_layout<{}>& from) {{\n", struct_name, struct_name);
+            fmt::print(file, "  fex_custom_repack_entry(source, from);\n");
+            fmt::print(file, "}}\n");
+
+            fmt::print(file, "bool fex_apply_custom_repacking_exit(guest_layout<{}>& into, host_layout<{}>& from) {{\n", struct_name, struct_name);
+            fmt::print(file, "  return fex_custom_repack_exit(into, from);\n");
+            fmt::print(file, "}}\n");
+        }
+
         fmt::print(file, "template<> inline constexpr bool has_compatible_data_layout<{}> = {};\n",
                    struct_name, (type_compat.at(type) == TypeCompatibility::Full));
     }
@@ -603,6 +641,7 @@ void GenerateThunkLibsAction::OnAnalysisComplete(clang::ASTContext& context) {
                     continue;
                 }
 
+                // Layout repacking happens here
                 if (!param_type->isPointerType() || (is_assumed_compatible || pointee_compat == TypeCompatibility::Full) ||
                     param_type->getPointeeType()->isBuiltinType() /* TODO: handle size_t. Actually, properly check for data layout compatibility */) {
                     // Fully compatible

ThunkLibs/include/common/GeneratorInterface.h

@@ -12,6 +12,22 @@ struct callback_annotation_base {
 };
 struct callback_stub : callback_annotation_base {};
 
+// Member annotation to mark members handled by custom repacking. This enables
+// automatic struct repacking of structs with non-trivial members (pointers,
+// unions, ...). Repacking logic is auto-generated as usual, with the
+// difference that an external function is called to manually repack the
+// annotated members.
+//
+// Two functions must be implemented for the parent struct type:
+// * fex_custom_repack_entry, called after automatic repacking of the other members
+// * fex_custom_repack_exit, called on exit but before automatic exit-repacking
+//     of the other members. Non-trivial implementations must perform host->guest
+//     repacking manually and return the boolean value true.
+//
+// If multiple members of the same struct are annotated as custom_repack,
+// they must be handled in the same fex_custom_repack_entry/exit functions.
+struct custom_repack {};
+
 // Type annotation to indicate that guest_layout/host_layout definitions should
 // be emitted even if the type is non-repackable. Pointer members will be
 // copied (or zero-extended) without regard for the referred data.

ThunkLibs/include/common/Host.h

@@ -262,17 +262,31 @@ struct repack_wrapper {
   repack_wrapper(guest_layout<GuestT>& orig_arg_) : orig_arg(orig_arg_) {
     if (orig_arg.get_pointer()) {
       data = { *orig_arg_.get_pointer() };
+
+      if constexpr (!std::is_enum_v<T>) {
+        constexpr bool is_compatible = has_compatible_data_layout<T> && std::is_same_v<T, GuestT>;
+        if constexpr (!is_compatible && std::is_class_v<std::remove_pointer_t<T>>) {
+          fex_apply_custom_repacking_entry(*data, *orig_arg_.get_pointer());
+        }
+      }
     }
   }
 
   ~repack_wrapper() {
     // TODO: Properly detect opaque types
     if constexpr (requires(guest_layout<T> t, decltype(data) h) { t.get_pointer(); (bool)h; *data; }) {
-      if constexpr (!std::is_const_v<std::remove_pointer_t<T>>) { // Skip exit-repacking for const pointees
-        if (data) {
-          constexpr bool is_compatible = has_compatible_data_layout<T> && std::is_same_v<T, GuestT>;
-          if constexpr (!is_compatible && std::is_class_v<std::remove_pointer_t<T>>) {
-            *orig_arg.get_pointer() = to_guest(*data); // TODO: Only if annotated as out-parameter
+      // NOTE: It's assumed that the native host library didn't modify any
+      //       const-pointees, so we skip automatic exit repacking for them.
+      //       However, *custom* repacking must still be applied since it might
+      //       have unrelated side effects (such as deallocation of memory
+      //       reserved on entry)
+      if (!fex_apply_custom_repacking_exit(*orig_arg.get_pointer(), *data)) {
+        if constexpr (!std::is_const_v<std::remove_pointer_t<T>>) { // Skip exit-repacking for const pointees
+          if (data) {
+            constexpr bool is_compatible = has_compatible_data_layout<T> && std::is_same_v<T, GuestT>;
+            if constexpr (!is_compatible && std::is_class_v<std::remove_pointer_t<T>>) {
+              *orig_arg.get_pointer() = to_guest(*data); // TODO: Only if annotated as out-parameter
+            }
           }
         }
       }

unittests/ThunkLibs/abi.cpp

@@ -596,6 +596,31 @@ TEST_CASE_METHOD(Fixture, "DataLayoutPointers") {
             CHECK(action->GetTypeCompatibility("struct B") == TypeCompatibility::None);
             CHECK(action->GetTypeCompatibility("struct A") == TypeCompatibility::None);
         }
+
+        SECTION("Innermost struct is incompatible but the pointer member is annotated") {
+            auto action = compute_data_layout(
+                "#include <thunks_common.h>\n"
+                "#include <cstdint>\n",
+                "#ifdef HOST\n"
+                "struct C { int32_t a; int32_t b; };\n"
+                "#else\n"
+                "struct C { int32_t b; int32_t a; };\n"
+                "#endif\n"
+                "struct B { C* a; int16_t b; };\n"
+                "struct A { int32_t a; B b; };\n"
+                "template<> struct fex_gen_config<&B::a> : fexgen::custom_repack {};\n"
+                "template<> struct fex_gen_type<A> {};\n", guest_abi);
+
+            INFO(FormatDataLayout(action->host_layout));
+
+            REQUIRE(action->guest_layout->contains("A"));
+            REQUIRE(action->guest_layout->contains("B"));
+            REQUIRE(action->guest_layout->contains("C"));
+
+            CHECK(action->GetTypeCompatibility("struct C") == TypeCompatibility::Repackable);
+            CHECK(action->GetTypeCompatibility("struct B") == TypeCompatibility::Repackable);
+            CHECK(action->GetTypeCompatibility("struct A") == TypeCompatibility::Repackable);
+        }
     }
 
     SECTION("Unannotated pointer to union type") {
@@ -623,6 +648,21 @@ TEST_CASE_METHOD(Fixture, "DataLayoutPointers") {
         CHECK(action->GetTypeCompatibility("struct A") == compat_full64_repackable32);
     }
 
+    SECTION("Pointer to union type with custom_repack annotation") {
+        auto action = compute_data_layout(
+            "#include <thunks_common.h>\n"
+            "#include <cstdint>\n",
+            "union B { int32_t a; uint32_t b; };\n"
+            "struct A { B* a; };\n"
+            "template<> struct fex_gen_config<&A::a> : fexgen::custom_repack {};\n"
+            "template<> struct fex_gen_type<A> {};\n", guest_abi);
+
+        INFO(FormatDataLayout(action->host_layout));
+
+        REQUIRE(action->guest_layout->contains("A"));
+        CHECK(action->GetTypeCompatibility("struct A") == TypeCompatibility::Repackable);
+    }
+
     SECTION("Pointer to opaque type") {
         auto action = compute_data_layout(
             "#include <thunks_common.h>\n"
@@ -638,6 +678,52 @@ TEST_CASE_METHOD(Fixture, "DataLayoutPointers") {
         CHECK(action->GetTypeCompatibility("struct A") == compat_full64_repackable32);
     }
 
+    SECTION("Pointer member with custom repacking code") {
+        // Data layout analysis only needs to know about the custom_repack
+        // annotation. The actual custom repacking code isn't needed for the
+        // test.
+
+        auto action = compute_data_layout(
+            "#include <thunks_common.h>\n"
+            "#include <cstdint>\n",
+            "#ifdef HOST\n"
+            "struct B { int32_t a; };\n"
+            "#else\n"
+            "struct B { int32_t b; };\n"
+            "#endif\n"
+            "struct A { B* a; };\n"
+            "template<> struct fex_gen_config<&A::a> : fexgen::custom_repack {};\n"
+            "template<> struct fex_gen_type<A> {};\n", guest_abi);
+
+        INFO(FormatDataLayout(action->host_layout));
+
+        REQUIRE(action->guest_layout->contains("A"));
+        REQUIRE(action->guest_layout->contains("B"));
+        CHECK(action->GetTypeCompatibility("struct A") == TypeCompatibility::Repackable);
+        CHECK(action->GetTypeCompatibility("struct B") == TypeCompatibility::None);
+    }
+
+    SECTION("Custom repacking induces repacking requirement") {
+        // Data layout analysis only needs to know about the custom_repack
+        // annotation. The actual custom repacking code isn't needed for the
+        // test.
+
+        auto action = compute_data_layout(
+            "#include <thunks_common.h>\n"
+            "#include <cstdint>\n",
+            "struct B {};\n"
+            "struct A { B* a; };\n"
+            "template<> struct fex_gen_config<&A::a> : fexgen::custom_repack {};\n"
+            "template<> struct fex_gen_type<A> {};\n", guest_abi);
+
+        INFO(FormatDataLayout(action->host_layout));
+
+        REQUIRE(action->guest_layout->contains("A"));
+        REQUIRE(action->guest_layout->contains("B"));
+        CHECK(action->GetTypeCompatibility("struct B") == TypeCompatibility::Full);
+        CHECK(action->GetTypeCompatibility("struct A") == TypeCompatibility::Repackable);
+    }
+
     SECTION("Self-referencing struct (like VkBaseOutStructure)") {
         // Without annotation
         auto action = compute_data_layout(

unittests/ThunkLibs/common.h

@@ -90,6 +90,7 @@ struct custom_host_impl {};
 struct callback_annotation_base { bool prevent_multiple; };
 struct callback_stub : callback_annotation_base {};
 
+struct custom_repack {};
 struct emit_layout_wrappers {};
 
 struct opaque_type {};

unittests/ThunkLibs/generator.cpp

@@ -711,6 +711,11 @@ TEST_CASE_METHOD(Fixture, "StructRepacking") {
         SECTION("Parameter annotated as ptr_passthrough") {
             CHECK_NOTHROW(run_thunkgen_host(prelude, code + "template<> struct fex_gen_param<func, 0, A*> : fexgen::ptr_passthrough {};\n", guest_abi));
         }
+
+        SECTION("Struct member annotated as custom_repack") {
+            CHECK_NOTHROW(run_thunkgen_host("struct A { void* a; };\n",
+                  code + "template<> struct fex_gen_config<&A::a> : fexgen::custom_repack {};\n", guest_abi));
+        }
     }
 
     SECTION("Pointer to struct with pointer member of consistent data layout") {
@@ -781,4 +786,15 @@ TEST_CASE_METHOD(Fixture, "VoidPointerParameter") {
             CHECK_NOTHROW(run_thunkgen_host(prelude, code, guest_abi));
         }
     }
+
+    SECTION("Custom repack in struct") {
+        const char* prelude =
+            "struct A { void* a; };\n";
+        const char* code =
+            "#include <thunks_common.h>\n"
+            "void func(A*);\n"
+            "template<> struct fex_gen_config<&A::a> : fexgen::custom_repack {};\n"
+            "template<> struct fex_gen_config<func> {};\n";
+        CHECK_NOTHROW(run_thunkgen_host(prelude, code, guest_abi));
+    }
 }